Spaces:
Sleeping
Sleeping
File size: 3,583 Bytes
116524e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | import logging
import time
import asyncio
from dataclasses import dataclass, field
from typing import Dict, Any
from ace.runners import ACELiteLLM
from ace.integrations.mcp.config import MCPServerConfig
from ace.integrations.mcp.errors import SessionNotFoundError
logger = logging.getLogger(__name__)
@dataclass
class Session:
session_id: str
runner: ACELiteLLM
created_at: float = field(default_factory=time.time)
last_accessed: float = field(default_factory=time.time)
lock: asyncio.Lock = field(default_factory=asyncio.Lock)
class SessionRegistry:
def __init__(self, config: MCPServerConfig):
self.config = config
self._sessions: Dict[str, Session] = {}
self._registry_lock = asyncio.Lock()
async def get_or_create(
self, session_id: str, model: str | None = None, **runner_kwargs: Any
) -> Session:
"""Get an existing session or create a new one, sweeping expired sessions first."""
async with self._registry_lock:
expired = self._collect_expired()
if session_id in self._sessions:
session = self._sessions[session_id]
session.last_accessed = time.time()
result = session
else:
# Create new runner
target_model = model or self.config.default_model
runner = ACELiteLLM.from_model(target_model, **runner_kwargs)
session = Session(session_id=session_id, runner=runner)
self._sessions[session_id] = session
result = session
# Drain outside the lock so other callers aren't blocked
self._drain_sessions(expired)
return result
async def get(self, session_id: str) -> Session:
"""Get an existing session. Raises SessionNotFoundError if not found."""
async with self._registry_lock:
expired = self._collect_expired()
if session_id not in self._sessions:
# Drain before raising so we don't leak
self._drain_sessions(expired)
raise SessionNotFoundError(session_id)
session = self._sessions[session_id]
session.last_accessed = time.time()
self._drain_sessions(expired)
return session
async def delete(self, session_id: str) -> None:
"""Delete a session if it exists."""
async with self._registry_lock:
session = self._sessions.pop(session_id, None)
if session is not None:
self._drain_sessions([session])
def _collect_expired(self) -> list[Session]:
"""Remove and return sessions that have exceeded the TTL.
Must be called while holding ``_registry_lock``.
"""
now = time.time()
ttl = self.config.session_ttl_seconds
expired_ids = [
sid
for sid, session in self._sessions.items()
if now - session.last_accessed > ttl
]
return [self._sessions.pop(sid) for sid in expired_ids]
@staticmethod
def _drain_sessions(sessions: list[Session]) -> None:
"""Best-effort wait for any in-progress background learning."""
for session in sessions:
try:
session.runner.wait_for_background(timeout=2.0)
except Exception:
logger.debug(
"Failed to drain session %s", session.session_id, exc_info=True
)
|