126 lines
4.1 KiB
Python
126 lines
4.1 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import time
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
from yuxi.utils.logging_config import logger
|
||
|
|
|
||
|
|
|
||
|
|
class ThreadBinding:
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
thread_key: str,
|
||
|
|
subagent_id: str | None = None,
|
||
|
|
acp_binding: str | None = None,
|
||
|
|
idle_timeout_ms: int = 86_400_000,
|
||
|
|
):
|
||
|
|
self.thread_key = thread_key
|
||
|
|
self.subagent_id = subagent_id
|
||
|
|
self.acp_binding = acp_binding
|
||
|
|
self.idle_timeout_ms = idle_timeout_ms
|
||
|
|
self.last_activity: float = time.monotonic()
|
||
|
|
self.created_at: float = time.monotonic()
|
||
|
|
|
||
|
|
def touch(self) -> None:
|
||
|
|
self.last_activity = time.monotonic()
|
||
|
|
|
||
|
|
def is_expired(self) -> bool:
|
||
|
|
elapsed = (time.monotonic() - self.last_activity) * 1000
|
||
|
|
return elapsed > self.idle_timeout_ms
|
||
|
|
|
||
|
|
|
||
|
|
class ThreadBindingManager:
|
||
|
|
def __init__(self, config: dict[str, Any] | None = None):
|
||
|
|
cfg = config or {}
|
||
|
|
bindings_cfg = cfg.get("thread_bindings", {})
|
||
|
|
self._enabled = bindings_cfg.get("enabled", False)
|
||
|
|
self._default_idle_timeout_ms = bindings_cfg.get("idle_timeout_ms", 86_400_000)
|
||
|
|
self._cleanup_interval = bindings_cfg.get("cleanup_interval", 300)
|
||
|
|
self._bindings: dict[str, ThreadBinding] = {}
|
||
|
|
self._cleanup_task: asyncio.Task | None = None
|
||
|
|
|
||
|
|
@property
|
||
|
|
def enabled(self) -> bool:
|
||
|
|
return self._enabled
|
||
|
|
|
||
|
|
def bind(
|
||
|
|
self,
|
||
|
|
thread_key: str,
|
||
|
|
subagent_id: str | None = None,
|
||
|
|
acp_binding: str | None = None,
|
||
|
|
idle_timeout_ms: int | None = None,
|
||
|
|
) -> ThreadBinding:
|
||
|
|
timeout = idle_timeout_ms or self._default_idle_timeout_ms
|
||
|
|
binding = ThreadBinding(
|
||
|
|
thread_key=thread_key,
|
||
|
|
subagent_id=subagent_id,
|
||
|
|
acp_binding=acp_binding,
|
||
|
|
idle_timeout_ms=timeout,
|
||
|
|
)
|
||
|
|
self._bindings[thread_key] = binding
|
||
|
|
logger.debug(f"[Telegram] Thread bound: {thread_key} -> subagent={subagent_id}")
|
||
|
|
return binding
|
||
|
|
|
||
|
|
def unbind(self, thread_key: str) -> bool:
|
||
|
|
if thread_key in self._bindings:
|
||
|
|
del self._bindings[thread_key]
|
||
|
|
logger.debug(f"[Telegram] Thread unbound: {thread_key}")
|
||
|
|
return True
|
||
|
|
return False
|
||
|
|
|
||
|
|
def get_binding(self, thread_key: str) -> ThreadBinding | None:
|
||
|
|
binding = self._bindings.get(thread_key)
|
||
|
|
if binding and binding.is_expired():
|
||
|
|
self.unbind(thread_key)
|
||
|
|
return None
|
||
|
|
return binding
|
||
|
|
|
||
|
|
def get_subagent_id(self, thread_key: str) -> str | None:
|
||
|
|
binding = self.get_binding(thread_key)
|
||
|
|
return binding.subagent_id if binding else None
|
||
|
|
|
||
|
|
def touch(self, thread_key: str) -> None:
|
||
|
|
binding = self._bindings.get(thread_key)
|
||
|
|
if binding:
|
||
|
|
binding.touch()
|
||
|
|
|
||
|
|
def cleanup_expired(self) -> int:
|
||
|
|
expired = [k for k, v in self._bindings.items() if v.is_expired()]
|
||
|
|
for k in expired:
|
||
|
|
self.unbind(k)
|
||
|
|
if expired:
|
||
|
|
logger.info(f"[Telegram] Cleaned up {len(expired)} expired thread bindings")
|
||
|
|
return len(expired)
|
||
|
|
|
||
|
|
async def start_cleanup_loop(self) -> None:
|
||
|
|
if not self._enabled:
|
||
|
|
return
|
||
|
|
|
||
|
|
async def _loop():
|
||
|
|
while True:
|
||
|
|
await asyncio.sleep(self._cleanup_interval)
|
||
|
|
self.cleanup_expired()
|
||
|
|
|
||
|
|
self._cleanup_task = asyncio.create_task(_loop())
|
||
|
|
|
||
|
|
async def stop_cleanup_loop(self) -> None:
|
||
|
|
if self._cleanup_task and not self._cleanup_task.done():
|
||
|
|
self._cleanup_task.cancel()
|
||
|
|
try:
|
||
|
|
await self._cleanup_task
|
||
|
|
except asyncio.CancelledError:
|
||
|
|
pass
|
||
|
|
self._cleanup_task = None
|
||
|
|
|
||
|
|
def get_all_bindings(self) -> dict[str, dict[str, Any]]:
|
||
|
|
return {
|
||
|
|
k: {
|
||
|
|
"thread_key": v.thread_key,
|
||
|
|
"subagent_id": v.subagent_id,
|
||
|
|
"acp_binding": v.acp_binding,
|
||
|
|
"idle_remaining_ms": max(0, int(v.idle_timeout_ms - (time.monotonic() - v.last_activity) * 1000)),
|
||
|
|
}
|
||
|
|
for k, v in self._bindings.items()
|
||
|
|
}
|