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)