ForcePilot/backend/package/yuxi/channels/adapters/feishu/thread_bindings.py

from __future__ import annotations

import asyncio
import time


class ThreadBindingManager:
    DEFAULT_TTL_S = 3600
    CLEANUP_INTERVAL_S = 300

    def __init__(self, ttl_s: int = DEFAULT_TTL_S):
        self._bindings: dict[str, str] = {}
        self._reverse: dict[str, str] = {}
        self._timestamps: dict[str, float] = {}
        self._ttl_s = ttl_s
        self._lock = asyncio.Lock()
        self._cleanup_task: asyncio.Task | None = None

    async def bind(self, thread_id: str, session_id: str) -> None:
        async with self._lock:
            self._bindings[thread_id] = session_id
            self._reverse[session_id] = thread_id
            self._timestamps[thread_id] = time.monotonic()

    async def unbind_thread(self, thread_id: str) -> None:
        async with self._lock:
            session_id = self._bindings.pop(thread_id, None)
            if session_id:
                self._reverse.pop(session_id, None)
            self._timestamps.pop(thread_id, None)

    async def unbind_session(self, session_id: str) -> None:
        async with self._lock:
            thread_id = self._reverse.pop(session_id, None)
            if thread_id:
                self._bindings.pop(thread_id, None)
                self._timestamps.pop(thread_id, None)

    async def get_session(self, thread_id: str) -> str | None:
        async with self._lock:
            session_id = self._bindings.get(thread_id)
            if session_id:
                ts = self._timestamps.get(thread_id, 0)
                if time.monotonic() - ts > self._ttl_s:
                    await self._unbind_thread_unlocked(thread_id)
                    return None
            return session_id

    async def get_thread(self, session_id: str) -> str | None:
        async with self._lock:
            return self._reverse.get(session_id)

    async def clear(self) -> None:
        async with self._lock:
            self._bindings.clear()
            self._reverse.clear()
            self._timestamps.clear()

    def start_cleanup(self) -> None:
        if self._cleanup_task is None or self._cleanup_task.done():
            self._cleanup_task = asyncio.create_task(self._cleanup_loop())

    def stop_cleanup(self) -> None:
        if self._cleanup_task and not self._cleanup_task.done():
            self._cleanup_task.cancel()
            self._cleanup_task = None

    async def _cleanup_loop(self) -> None:
        while True:
            try:
                await asyncio.sleep(self.CLEANUP_INTERVAL_S)
                await self._purge_expired()
            except asyncio.CancelledError:
                return

    async def _purge_expired(self) -> None:
        async with self._lock:
            now = time.monotonic()
            expired = [tid for tid, ts in self._timestamps.items() if now - ts > self._ttl_s]
            for thread_id in expired:
                await self._unbind_thread_unlocked(thread_id)

    async def _unbind_thread_unlocked(self, thread_id: str) -> None:
        session_id = self._bindings.pop(thread_id, None)
        if session_id:
            self._reverse.pop(session_id, None)
        self._timestamps.pop(thread_id, None)