Mutual_Fund/service/client_agent/archive_schedule.py

"""Redis client_agent 会话自动归档调度数据结构。"""

from __future__ import annotations


class SessionArchiveSchedule:
    """管理空闲会话的归档截止时间，不负责实际归档。"""

    DUE_KEY = "session:archive:due"
    META_KEY = "session:{session_id}:meta"

    def __init__(self, redis):
        self.redis = redis

    @classmethod
    def meta_key(cls, session_id: str) -> str:
        """生成与短期记忆共用的会话元数据 Key。"""
        return cls.META_KEY.format(session_id=session_id)

    async def register(
        self,
        *,
        session_id: str,
        created_at: float,
        last_activity_at: float,
        archive_due_at: float,
        absolute_expire_at: float,
    ) -> None:
        """登记新会话的归档时间和生命周期信息。"""
        await self.redis.hset(
            self.meta_key(session_id),
            mapping={
                "created_at": str(created_at),
                "last_activity_at": str(last_activity_at),
                "archive_due_at": str(archive_due_at),
                "absolute_expire_at": str(absolute_expire_at),
                "archive_status": "pending",
                "archive_retry_count": "0",
            },
        )
        await self.redis.expire(
            self.meta_key(session_id),
            max(1, int(absolute_expire_at - created_at)),
        )
        await self.redis.zadd(self.DUE_KEY, {session_id: archive_due_at})

    async def touch(
        self,
        *,
        session_id: str,
        last_activity_at: float,
        archive_due_at: float,
    ) -> None:
        """更新会话最近活动时间和下一次归档截止时间。"""
        await self.redis.hset(
            self.meta_key(session_id),
            mapping={
                "last_activity_at": str(last_activity_at),
                "archive_due_at": str(archive_due_at),
                "archive_status": "pending",
            },
        )
        await self.redis.zadd(self.DUE_KEY, {session_id: archive_due_at})

    async def list_due(self, *, now: float, limit: int = 100) -> list[str]:
        """返回归档截止时间已到的会话。"""
        if limit <= 0:
            return []
        values = await self.redis.zrangebyscore(
            self.DUE_KEY, 0, now, start=0, num=limit
        )
        return [value.decode() if isinstance(value, bytes) else str(value) for value in values]

    async def get_meta(self, session_id: str) -> dict:
        """读取会话归档元数据。"""
        return await self.redis.hgetall(self.meta_key(session_id))

    async def reschedule(self, session_id: str, *, due_at: float) -> None:
        """更新下一次扫描时间，保留当前调度记录。"""
        await self.redis.zadd(self.DUE_KEY, {session_id: due_at})

    async def rebuild_due_index(self, *, limit: int = 10000) -> int:
        """从会话元数据重建丢失的到期索引。"""
        rebuilt = 0
        async for raw_key in self.redis.scan_iter(match="session:*:meta", count=100):
            key = raw_key.decode() if isinstance(raw_key, bytes) else raw_key
            prefix = "session:"
            suffix = ":meta"
            if not key.startswith(prefix) or not key.endswith(suffix):
                continue
            session_id = key[len(prefix) : -len(suffix)]
            meta = await self.redis.hgetall(key)
            due_at = meta.get("archive_due_at")
            if due_at is None or meta.get("archive_status") == "success":
                continue
            await self.redis.zadd(self.DUE_KEY, {session_id: float(due_at)})
            rebuilt += 1
            if rebuilt >= limit:
                break
        return rebuilt

    async def mark_processing(self, session_id: str) -> None:
        """标记会话进入归档处理状态。"""
        await self.redis.hset(
            self.meta_key(session_id), mapping={"archive_status": "processing"}
        )

    async def mark_retry(
        self,
        session_id: str,
        *,
        retry_count: int,
        error: str,
        next_retry_at: float,
    ) -> None:
        """记录归档失败信息和下一次重试时间。"""
        await self.redis.hset(
            self.meta_key(session_id),
            mapping={
                "archive_status": "retry",
                "archive_retry_count": str(retry_count),
                "last_archive_error": error[:500],
                "next_retry_at": str(next_retry_at),
            },
        )

    async def remove(self, session_id: str) -> None:
        """移除归档调度记录，不删除会话消息或短期记忆元数据。"""
        await self.redis.zrem(self.DUE_KEY, session_id)


__all__ = ["SessionArchiveSchedule"]