Mutual_Fund/agent/customer_agent/session.py

"""Redis-backed anonymous customer-service session primitives."""
from __future__ import annotations

import time
import uuid
from inspect import isawaitable


class SessionOwnershipError(Exception):
    def __init__(self, code: int, message: str):
        self.code = code
        self.message = message
        super().__init__(message)


async def _config(config_getter, key: str, default):
    value = config_getter(key, str(default))
    if isawaitable(value):
        value = await value
    return type(default)(value)


class AnonymousSessionService:
    def __init__(self, redis, *, config_getter, clock=time.time):
        self.redis = redis
        self.config_getter = config_getter
        self.clock = clock

    async def create_session(self) -> str:
        session_id = uuid.uuid4().hex
        ttl = await _config(self.config_getter, "agent.customer.session.ttl", 1800)
        await self.redis.set(f"session:{session_id}", "anonymous", ex=ttl)
        # Redis lists cannot be created empty; the marker is ignored by readers.
        await self.redis.rpush(f"session:{session_id}:messages", "")
        await self.redis.expire(f"session:{session_id}:messages", ttl)
        return session_id

    async def verify_session_ownership(
        self,
        session_id: str,
        *,
        customer_id: str | None = None,
        session_customer_id: str | None = None,
    ) -> None:
        if not await self.redis.exists(f"session:{session_id}"):
            raise SessionOwnershipError(404, "会话不存在或已过期")
        if customer_id is not None and customer_id != session_customer_id:
            raise SessionOwnershipError(403, "无权访问该会话")

    async def consume_chat_quota(self, session_id: str) -> int | None:
        window = await _config(
            self.config_getter, "agent.customer.rate_limit.window_sec", 60
        )
        maximum = await _config(
            self.config_getter, "agent.customer.rate_limit.max_requests", 20
        )
        key = f"rate:limit:anon:{session_id}:chat"
        now = self.clock()
        await self.redis.zremrangebyscore(key, 0, (now - window) * 1000)
        count = await self.redis.zcard(key)
        if count >= maximum:
            entries = await self.redis.zrange(key, 0, 0, withscores=True)
            oldest = entries[0][1] if entries else now * 1000
            return max(1, int((oldest / 1000 + window) - now))
        await self.redis.zadd(key, {uuid.uuid4().hex: now * 1000})
        await self.redis.expire(key, window)
        return None