299 lines
10 KiB
Python
299 lines
10 KiB
Python
"""
|
|
Session Manager for ComfyUI-Doctor.
|
|
|
|
Manages a reusable aiohttp ClientSession for LLM API calls to reduce
|
|
resource overhead and prevent "Unclosed client session" warnings.
|
|
|
|
R7 Enhancement: Adds rate limiting and concurrency control for LLM requests.
|
|
"""
|
|
|
|
import asyncio
|
|
import aiohttp
|
|
import atexit
|
|
import logging
|
|
from typing import Optional
|
|
|
|
try:
|
|
from .rate_limiter import RateLimiter, ConcurrencyLimiter
|
|
except ImportError as import_error:
|
|
from import_compat import ensure_absolute_import_fallback_allowed
|
|
ensure_absolute_import_fallback_allowed(import_error)
|
|
from rate_limiter import RateLimiter, ConcurrencyLimiter
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class SessionManager:
|
|
"""
|
|
Manages a shared aiohttp ClientSession for LLM API calls.
|
|
|
|
Benefits:
|
|
- Connection reuse (HTTP keep-alive)
|
|
- Reduced resource overhead
|
|
- Proper cleanup on shutdown
|
|
|
|
Usage:
|
|
session = await SessionManager.get_session()
|
|
async with session.get(url) as response:
|
|
...
|
|
"""
|
|
_session: Optional[aiohttp.ClientSession] = None
|
|
_lock: Optional[asyncio.Lock] = None
|
|
_closed: bool = False
|
|
|
|
# Default timeout for LLM API calls
|
|
DEFAULT_TIMEOUT = aiohttp.ClientTimeout(total=60)
|
|
|
|
# R7: Rate limiters (shared across all requests)
|
|
_core_limiter: Optional[RateLimiter] = None # analyze, chat
|
|
_light_limiter: Optional[RateLimiter] = None # verify_key, list_models
|
|
_concurrency: Optional[ConcurrencyLimiter] = None # max concurrent LLM calls
|
|
|
|
# R7: Default limits (configurable via CONFIG)
|
|
DEFAULT_CORE_RATE = 30 # req/min for heavy endpoints
|
|
DEFAULT_LIGHT_RATE = 10 # req/min for light endpoints
|
|
DEFAULT_CONCURRENCY = 3 # simultaneous LLM requests
|
|
DEFAULT_RATE_WINDOW_SECONDS = 60
|
|
|
|
# S13: Outbound proxy trust boundary policy (secure by default).
|
|
PROXY_POLICY_STRICT_OFF = "strict_off"
|
|
PROXY_POLICY_INHERIT_ENV = "inherit_env"
|
|
ENV_PROXY_POLICY_KEY = "DOCTOR_OUTBOUND_PROXY_POLICY"
|
|
_proxy_policy: str = PROXY_POLICY_STRICT_OFF
|
|
_proxy_policy_source: str = "default"
|
|
|
|
@classmethod
|
|
def configure_limits(
|
|
cls,
|
|
*,
|
|
core_rate_limit: Optional[int] = None,
|
|
light_rate_limit: Optional[int] = None,
|
|
max_concurrent: Optional[int] = None,
|
|
rate_window_seconds: Optional[int] = None,
|
|
) -> None:
|
|
"""
|
|
Configure rate/concurrency limits (typically from config).
|
|
|
|
This resets cached limiter instances so new values take effect.
|
|
"""
|
|
def _coerce_positive_int(value: Optional[int], default: int) -> int:
|
|
if value is None:
|
|
return default
|
|
try:
|
|
coerced = int(value)
|
|
except (TypeError, ValueError):
|
|
return default
|
|
return coerced if coerced > 0 else default
|
|
|
|
cls.DEFAULT_CORE_RATE = _coerce_positive_int(core_rate_limit, cls.DEFAULT_CORE_RATE)
|
|
cls.DEFAULT_LIGHT_RATE = _coerce_positive_int(light_rate_limit, cls.DEFAULT_LIGHT_RATE)
|
|
cls.DEFAULT_CONCURRENCY = _coerce_positive_int(max_concurrent, cls.DEFAULT_CONCURRENCY)
|
|
cls.DEFAULT_RATE_WINDOW_SECONDS = _coerce_positive_int(
|
|
rate_window_seconds,
|
|
cls.DEFAULT_RATE_WINDOW_SECONDS,
|
|
)
|
|
|
|
cls._core_limiter = None
|
|
cls._light_limiter = None
|
|
cls._concurrency = None
|
|
|
|
@classmethod
|
|
def _get_lock(cls) -> asyncio.Lock:
|
|
"""Get or create the async lock (must be created in event loop)."""
|
|
if cls._lock is None:
|
|
cls._lock = asyncio.Lock()
|
|
return cls._lock
|
|
|
|
@classmethod
|
|
def _normalize_proxy_policy(cls, policy: Optional[str]) -> str:
|
|
normalized = (policy or "").strip().lower()
|
|
if normalized in {cls.PROXY_POLICY_STRICT_OFF, cls.PROXY_POLICY_INHERIT_ENV}:
|
|
return normalized
|
|
return cls.PROXY_POLICY_STRICT_OFF
|
|
|
|
@classmethod
|
|
def configure_proxy_policy(
|
|
cls,
|
|
*,
|
|
config_policy: Optional[str] = None,
|
|
env_policy: Optional[str] = None,
|
|
) -> None:
|
|
"""
|
|
Configure proxy trust policy with env override precedence.
|
|
|
|
Precedence:
|
|
1) env_policy (DOCTOR_OUTBOUND_PROXY_POLICY)
|
|
2) config_policy
|
|
3) strict_off fallback
|
|
"""
|
|
chosen_source = "default"
|
|
chosen_value: Optional[str] = None
|
|
|
|
raw_env = (env_policy or "").strip()
|
|
if raw_env:
|
|
chosen_source = "env"
|
|
chosen_value = raw_env
|
|
else:
|
|
raw_config = (config_policy or "").strip()
|
|
if raw_config:
|
|
chosen_source = "config"
|
|
chosen_value = raw_config
|
|
|
|
normalized = cls._normalize_proxy_policy(chosen_value)
|
|
if chosen_value and normalized != chosen_value.lower():
|
|
logger.warning(
|
|
"Invalid outbound proxy policy '%s' from %s; falling back to '%s'",
|
|
chosen_value,
|
|
chosen_source,
|
|
cls.PROXY_POLICY_STRICT_OFF,
|
|
)
|
|
chosen_source = f"{chosen_source}_invalid_fallback"
|
|
|
|
cls._proxy_policy = normalized
|
|
cls._proxy_policy_source = chosen_source
|
|
|
|
@classmethod
|
|
def get_proxy_diagnostics(cls) -> dict:
|
|
"""Return effective outbound proxy trust settings for diagnostics."""
|
|
return {
|
|
"policy": cls._proxy_policy,
|
|
"trust_env": cls._proxy_policy == cls.PROXY_POLICY_INHERIT_ENV,
|
|
"source": cls._proxy_policy_source,
|
|
"env_key": cls.ENV_PROXY_POLICY_KEY,
|
|
}
|
|
|
|
@classmethod
|
|
def get_core_limiter(cls) -> RateLimiter:
|
|
"""
|
|
Get rate limiter for heavy LLM endpoints (analyze, chat).
|
|
|
|
Returns:
|
|
RateLimiter: Shared rate limiter instance (30 req/min default).
|
|
"""
|
|
if cls._core_limiter is None:
|
|
cls._core_limiter = RateLimiter(
|
|
max_per_minute=cls.DEFAULT_CORE_RATE,
|
|
window_seconds=cls.DEFAULT_RATE_WINDOW_SECONDS,
|
|
)
|
|
return cls._core_limiter
|
|
|
|
@classmethod
|
|
def get_light_limiter(cls) -> RateLimiter:
|
|
"""
|
|
Get rate limiter for light endpoints (verify_key, list_models).
|
|
|
|
Returns:
|
|
RateLimiter: Shared rate limiter instance (10 req/min default).
|
|
"""
|
|
if cls._light_limiter is None:
|
|
cls._light_limiter = RateLimiter(
|
|
max_per_minute=cls.DEFAULT_LIGHT_RATE,
|
|
window_seconds=cls.DEFAULT_RATE_WINDOW_SECONDS,
|
|
)
|
|
return cls._light_limiter
|
|
|
|
@classmethod
|
|
def get_concurrency_limiter(cls) -> ConcurrencyLimiter:
|
|
"""
|
|
Get concurrency limiter for LLM requests.
|
|
|
|
Returns:
|
|
ConcurrencyLimiter: Shared semaphore (3 concurrent default).
|
|
"""
|
|
if cls._concurrency is None:
|
|
cls._concurrency = ConcurrencyLimiter(max_concurrent=cls.DEFAULT_CONCURRENCY)
|
|
return cls._concurrency
|
|
|
|
@classmethod
|
|
async def get_session(cls) -> aiohttp.ClientSession:
|
|
"""
|
|
Get or create a shared aiohttp ClientSession.
|
|
|
|
Returns:
|
|
aiohttp.ClientSession: The shared session instance.
|
|
|
|
Note:
|
|
This method is thread-safe and will create the session
|
|
on first call. The session uses a default 60-second timeout.
|
|
"""
|
|
async with cls._get_lock():
|
|
if cls._session is None or cls._session.closed:
|
|
cls._closed = False
|
|
# CRITICAL: Keep strict-off default to avoid ambient proxy env inheritance.
|
|
cls._session = aiohttp.ClientSession(
|
|
timeout=cls.DEFAULT_TIMEOUT,
|
|
trust_env=(cls._proxy_policy == cls.PROXY_POLICY_INHERIT_ENV),
|
|
)
|
|
return cls._session
|
|
|
|
@classmethod
|
|
async def close(cls) -> None:
|
|
"""
|
|
Close the shared session.
|
|
|
|
This should be called during application shutdown to ensure
|
|
proper cleanup. Safe to call multiple times.
|
|
"""
|
|
async with cls._get_lock():
|
|
if cls._session is not None and not cls._session.closed:
|
|
await cls._session.close()
|
|
cls._session = None
|
|
cls._closed = True
|
|
|
|
@classmethod
|
|
def is_closed(cls) -> bool:
|
|
"""Check if the session has been explicitly closed."""
|
|
return cls._closed
|
|
|
|
@classmethod
|
|
def reset(cls) -> None:
|
|
"""
|
|
Reset the manager state (for testing purposes).
|
|
|
|
Warning: This does not close the session - call close() first.
|
|
"""
|
|
cls._session = None
|
|
cls._lock = None
|
|
cls._closed = False
|
|
cls._core_limiter = None
|
|
cls._light_limiter = None
|
|
cls._concurrency = None
|
|
cls._proxy_policy = cls.PROXY_POLICY_STRICT_OFF
|
|
cls._proxy_policy_source = "default"
|
|
|
|
|
|
def _sync_close_session():
|
|
"""
|
|
Synchronous wrapper for session cleanup at interpreter shutdown.
|
|
|
|
Called by atexit to ensure session is closed when Python exits.
|
|
"""
|
|
if SessionManager._session is not None and not SessionManager._session.closed:
|
|
try:
|
|
# Try to get or create an event loop
|
|
try:
|
|
loop = asyncio.get_running_loop()
|
|
except RuntimeError:
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
|
|
# Run the close coroutine
|
|
if loop.is_running():
|
|
# If loop is running, schedule the close
|
|
loop.create_task(SessionManager.close())
|
|
else:
|
|
loop.run_until_complete(SessionManager.close())
|
|
except Exception:
|
|
# At shutdown, errors are expected - just try to close directly
|
|
try:
|
|
if SessionManager._session is not None:
|
|
# Force close the connector
|
|
if SessionManager._session.connector is not None:
|
|
SessionManager._session.connector.close()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
# Register cleanup at interpreter exit
|
|
atexit.register(_sync_close_session)
|