llm-in-text/backend/session_store.py

import hashlib
import secrets
import threading
import time
from dataclasses import dataclass
from typing import Any

try:
    import psycopg
    from psycopg.rows import dict_row
except Exception:  # pragma: no cover
    psycopg = None
    dict_row = None


def _now_ms() -> int:
    return int(time.time() * 1000)


def hash_value(value: str) -> str:
    return hashlib.sha256(value.encode("utf-8")).hexdigest()


def new_session_id() -> str:
    return secrets.token_urlsafe(32)


@dataclass
class SessionRecord:
    session_id: str
    session_hash: str
    created_at_ms: int
    last_seen_at_ms: int
    first_ip_hash: str
    last_ip_hash: str
    user_agent_hash: str
    risk_score: int = 0
    blocked_until_ms: int = 0
    is_new: bool = False


class BaseSessionStore:
    def get_or_create(self, session_id: str | None, *, client_ip_hash: str, user_agent_hash: str) -> SessionRecord:
        raise NotImplementedError


class InMemorySessionStore(BaseSessionStore):
    def __init__(self) -> None:
        self.sessions: dict[str, SessionRecord] = {}
        self.lock = threading.Lock()

    def get_or_create(self, session_id: str | None, *, client_ip_hash: str, user_agent_hash: str) -> SessionRecord:
        now = _now_ms()
        with self.lock:
            if session_id:
                session_hash = hash_value(session_id)
                record = self.sessions.get(session_hash)
                if record is not None:
                    record.last_seen_at_ms = now
                    record.last_ip_hash = client_ip_hash
                    record.user_agent_hash = user_agent_hash
                    record.is_new = False
                    return record
            next_session_id = new_session_id()
            next_hash = hash_value(next_session_id)
            record = SessionRecord(
                session_id=next_session_id,
                session_hash=next_hash,
                created_at_ms=now,
                last_seen_at_ms=now,
                first_ip_hash=client_ip_hash,
                last_ip_hash=client_ip_hash,
                user_agent_hash=user_agent_hash,
                is_new=True,
            )
            self.sessions[next_hash] = record
            return record


class PostgresSessionStore(BaseSessionStore):
    def __init__(self, database_url: str) -> None:
        if psycopg is None or dict_row is None:
            raise RuntimeError("psycopg 未安装，无法使用 PostgreSQL session 存储")
        self.database_url = database_url
        self._init_lock = threading.Lock()
        self._initialized = False

    def _connect(self):
        return psycopg.connect(self.database_url, autocommit=True, row_factory=dict_row)

    def _ensure_initialized(self) -> None:
        if self._initialized:
            return
        with self._init_lock:
            if self._initialized:
                return
            with self._connect() as conn:
                with conn.cursor() as cur:
                    cur.execute(
                        """
                        CREATE TABLE IF NOT EXISTS anonymous_sessions (
                            session_hash TEXT PRIMARY KEY,
                            created_at_ms BIGINT NOT NULL,
                            last_seen_at_ms BIGINT NOT NULL,
                            first_ip_hash TEXT NOT NULL,
                            last_ip_hash TEXT NOT NULL,
                            user_agent_hash TEXT NOT NULL,
                            risk_score INTEGER NOT NULL DEFAULT 0,
                            blocked_until_ms BIGINT NOT NULL DEFAULT 0,
                            metadata_json JSONB NOT NULL DEFAULT '{}'::jsonb
                        )
                        """
                    )
                    cur.execute(
                        "CREATE INDEX IF NOT EXISTS anonymous_sessions_last_seen_idx ON anonymous_sessions(last_seen_at_ms)"
                    )
            self._initialized = True

    def get_or_create(self, session_id: str | None, *, client_ip_hash: str, user_agent_hash: str) -> SessionRecord:
        self._ensure_initialized()
        now = _now_ms()
        if session_id:
            session_hash = hash_value(session_id)
            with self._connect() as conn:
                with conn.cursor() as cur:
                    cur.execute(
                        """
                        UPDATE anonymous_sessions
                        SET last_seen_at_ms = %s,
                            last_ip_hash = %s,
                            user_agent_hash = %s
                        WHERE session_hash = %s
                        RETURNING session_hash, created_at_ms, last_seen_at_ms, first_ip_hash, last_ip_hash, user_agent_hash, risk_score, blocked_until_ms
                        """,
                        (now, client_ip_hash, user_agent_hash, session_hash),
                    )
                    row = cur.fetchone()
                    if row is not None:
                        return SessionRecord(
                            session_id=session_id,
                            session_hash=row["session_hash"],
                            created_at_ms=int(row["created_at_ms"]),
                            last_seen_at_ms=int(row["last_seen_at_ms"]),
                            first_ip_hash=row["first_ip_hash"],
                            last_ip_hash=row["last_ip_hash"],
                            user_agent_hash=row["user_agent_hash"],
                            risk_score=int(row["risk_score"] or 0),
                            blocked_until_ms=int(row["blocked_until_ms"] or 0),
                            is_new=False,
                        )
        next_session_id = new_session_id()
        next_hash = hash_value(next_session_id)
        with self._connect() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    INSERT INTO anonymous_sessions (
                        session_hash, created_at_ms, last_seen_at_ms, first_ip_hash, last_ip_hash, user_agent_hash
                    )
                    VALUES (%s, %s, %s, %s, %s, %s)
                    """,
                    (next_hash, now, now, client_ip_hash, client_ip_hash, user_agent_hash),
                )
        return SessionRecord(
            session_id=next_session_id,
            session_hash=next_hash,
            created_at_ms=now,
            last_seen_at_ms=now,
            first_ip_hash=client_ip_hash,
            last_ip_hash=client_ip_hash,
            user_agent_hash=user_agent_hash,
            is_new=True,
        )


_session_store: BaseSessionStore | None = None


def get_session_store(database_url: str | None = None) -> BaseSessionStore:
    global _session_store
    if _session_store is not None:
        return _session_store
    if database_url:
        _session_store = PostgresSessionStore(database_url)
    else:
        _session_store = InMemorySessionStore()
    return _session_store


def reset_session_store() -> None:
    global _session_store
    _session_store = None