feat(gateway): session model pool — concurrency-aware auto-assignment with auxiliary slot tracking

agent/auxiliary_client.py

@@ -539,6 +539,38 @@ def _peek_pool_entry(provider: str) -> Optional[Any]:
     return None
 
 
+# ---- SessionModelPool helper (auxiliary slot tracking) ----
+
+_UNSET = object()  # sentinel: distinguishes "never tried" from "pool is None/disabled"
+_aux_pool_cache = _UNSET
+
+
+def _get_session_model_pool():
+    """Return the SessionModelPool singleton (or None if disabled/unavailable).
+
+    Uses a sentinel so that a ``None`` result (pool disabled/not configured)
+    is not cached permanently — if the pool is created later in the process
+    lifetime, subsequent calls will find it.
+
+    Import failures are also not cached, so transient import errors (e.g.
+    during startup) don't permanently disable pool tracking.
+    """
+    global _aux_pool_cache
+    if _aux_pool_cache is not _UNSET:
+        return _aux_pool_cache
+    try:
+        from gateway.session_model_pool import get_session_model_pool
+        _pool = get_session_model_pool({})
+        if _pool is not None:
+            # Cache only non-None (enabled) results.
+            _aux_pool_cache = _pool
+        # None (disabled/unconfigured) is NOT cached — next call retries.
+    except Exception:
+        # Import failure is NOT cached — next call retries.
+        pass
+    return _aux_pool_cache
+
+
 def _pool_runtime_api_key(entry: Any) -> str:
     if entry is None:
         return ""
@@ -4875,7 +4907,10 @@ def call_llm(
         extra_body: Additional request body fields.
 
     Returns:
-        Response object with .choices[0].message.content
+        Response object with .choices[0].message.content.
+        Returns ``None`` when the SessionModelPool auxiliary slot is
+        blocked (pool enabled + model saturated). Callers should check
+        for ``None`` before accessing response attributes.
 
     Raises:
         RuntimeError: If no provider is configured.
@@ -4968,6 +5003,27 @@ def call_llm(
 
     # Handle unsupported temperature, max_tokens vs max_completion_tokens retry,
     # then payment fallback.
+    #
+    # Session Model Pool: acquire an auxiliary slot before making the call
+    # so the pool can throttle concurrent auxiliary requests to the same model.
+    _pool_aux_acquired = False
+    _pool = _get_session_model_pool()
+    try:
+        if _pool and _pool.enabled:
+            _pool_aux_acquired = _pool.acquire_auxiliary_slot(final_model or "", resolved_provider or "")
+    except Exception:
+        pass
+
+    if not _pool_aux_acquired and _pool and _pool.enabled:
+        logger.warning(
+            "Auxiliary %s: blocked by SessionModelPool for %s:%s — throttled",
+            task or "call", resolved_provider, final_model,
+        )
+        raise RuntimeError(
+            f"Auxiliary call '{task or 'call'}' throttled by SessionModelPool: "
+            f"no auxiliary slots available for {resolved_provider}:{final_model}"
+        )
+
     try:
         return _validate_llm_response(
             client.chat.completions.create(**kwargs), task)
@@ -5261,6 +5317,18 @@ def call_llm(
                 logger.debug("Auxiliary: cache eviction after connection error failed",
                              exc_info=True)
         raise
+    finally:
+        # Session Model Pool: release auxiliary slot after the call
+        # completes (success, error, or fallback).
+        if _pool_aux_acquired:
+            try:
+                if _pool and _pool.enabled:
+                    _pool.release_auxiliary_slot(final_model or "", resolved_provider or "")
+            except Exception as _exc:
+                logger.error(
+                    "SessionModelPool: FAILED to release auxiliary slot for %s:%s — "
+                    "slot may be leaked: %s", resolved_provider, final_model, _exc,
+                )
 
 
 def extract_content_or_reasoning(response) -> str:
@@ -5337,6 +5405,9 @@ async def async_call_llm(
     """Centralized asynchronous LLM call.
 
     Same as call_llm() but async. See call_llm() for full documentation.
+
+    Includes SessionModelPool auxiliary slot tracking (acquire before call,
+    release in finally block), mirroring the synchronous implementation.
     """
     resolved_provider, resolved_model, resolved_base_url, resolved_api_key, resolved_api_mode = _resolve_task_provider_model(
         task, provider, model, base_url, api_key)
@@ -5409,6 +5480,28 @@ async def async_call_llm(
     if _is_anthropic_compat_endpoint(resolved_provider, _client_base):
         kwargs["messages"] = _convert_openai_images_to_anthropic(kwargs["messages"])
 
+    # Session Model Pool: acquire an auxiliary slot before making the async call
+    # so the pool can throttle concurrent auxiliary requests to the same model.
+    _async_pool_aux_acquired = False
+    _async_pool = _get_session_model_pool()
+    try:
+        if _async_pool and _async_pool.enabled:
+            _async_pool_aux_acquired = _async_pool.acquire_auxiliary_slot(
+                final_model or "", resolved_provider or ""
+            )
+    except Exception:
+        pass
+
+    if not _async_pool_aux_acquired and _async_pool and _async_pool.enabled:
+        logger.warning(
+            "Auxiliary %s (async): blocked by SessionModelPool for %s:%s — throttled",
+            task or "call", resolved_provider, final_model,
+        )
+        raise RuntimeError(
+            f"Async auxiliary call '{task or 'call'}' throttled by SessionModelPool: "
+            f"no auxiliary slots available for {resolved_provider}:{final_model}"
+        )
+
     try:
         return _validate_llm_response(
             await client.chat.completions.create(**kwargs), task)
@@ -5660,3 +5753,17 @@ async def async_call_llm(
                 logger.debug("Auxiliary (async): cache eviction after connection error failed",
                              exc_info=True)
         raise
+    finally:
+        # Session Model Pool: release auxiliary slot after the async call
+        # completes (success, error, or fallback).
+        if _async_pool_aux_acquired:
+            try:
+                if _async_pool and _async_pool.enabled:
+                    _async_pool.release_auxiliary_slot(
+                        final_model or "", resolved_provider or ""
+                    )
+            except Exception as _exc:
+                logger.error(
+                    "SessionModelPool: FAILED to release async auxiliary slot for %s:%s — "
+                    "slot may be leaked: %s", resolved_provider, final_model, _exc,
+                )

gateway/run.py

@@ -24,6 +24,7 @@
     # means UTF-8 stdio setup is skipped on Windows; POSIX is unaffected.
     pass
 
+
 import asyncio
 import dataclasses
 import inspect
@@ -1876,6 +1877,12 @@ def __init__(self, config: Optional[GatewayConfig] = None):
         # Per-session reasoning effort overrides from /reasoning.
         # Key: session_key, Value: parsed reasoning config dict.
         self._session_reasoning_overrides: Dict[str, Dict[str, Any]] = {}
+        # Pool-assigned models (populated by SessionModelPool when enabled).
+        # Key: session_key, Value: dict with model/provider/context_length.
+        # These are weaker than manual /model overrides and are released
+        # when the session ends or when a manual override takes effect.
+        self._pool_assigned_models: Dict[str, Dict[str, Any]] = {}
+        self._pool_assigned_models_lock = threading.Lock()
         self._kanban_notifier_profile = self._active_profile_name()
         # Teams meeting pipeline runtime (bound later when msgraph_webhook adapter exists).
         self._teams_pipeline_runtime = None
@@ -2513,6 +2520,41 @@ def _recover_telegram_topic_thread_id(
                 return None
         return None
 
+    def _mark_pool_override(self, session_key: str) -> None:
+        """Mark a session as manually overridden in the pool.
+
+        Called after ``_release_pool_slot`` so the pool knows not to
+        reassign a model on the next turn.
+        """
+        try:
+            from gateway.session_model_pool import get_session_model_pool as _get_pool
+            _p = _get_pool({})
+            if _p:
+                _p.mark_manual_override(session_key)
+        except Exception as _exc:
+            logger.debug("SessionModelPool: failed to mark override for %s: %s", session_key, _exc)
+
+    def _release_pool_slot(self, session_key: str) -> None:
+        """Release a pool-assigned slot for a session (if one exists).
+
+        Centralizes the release pattern used in 3 places: session reset,
+        /model override, in-place model switch, and any other override path.
+        Thread-safe: acquires ``_pool_assigned_models_lock`` internally.
+        """
+        try:
+            with self._pool_assigned_models_lock:
+                _old_pool = self._pool_assigned_models.pop(session_key, None)
+            if _old_pool:
+                from gateway.session_model_pool import get_session_model_pool as _get_pool
+                # The singleton ignores config after first init; pass {}
+                # to avoid unnecessary disk I/O via _load_gateway_config().
+                _p = _get_pool({})
+                if _p:
+                    _p.release_session_slot(session_key)
+        except Exception as _exc:
+            logger.debug("SessionModelPool: failed to release slot for %s: %s", session_key, _exc)
+
+
     def _resolve_session_agent_runtime(
         self,
         *,
@@ -2535,6 +2577,42 @@ def _resolve_session_agent_runtime(
 
         model = _resolve_gateway_model(user_config)
         override = self._session_model_overrides.get(resolved_session_key) if resolved_session_key else None
+        # Will be set by pool integration below if a pool-assigned provider
+        # is available. Applied after runtime_kwargs is created.
+        _pool_provider_override = None
+
+        # --- Session Model Pool integration ---
+        # If no manual override exists for this session, check whether the
+        # pool wants to assign a different model. Pool assignments are
+        # weaker than manual /model overrides and are released when the
+        # session ends or when a manual override takes effect.
+        if not override and resolved_session_key:
+            try:
+                from gateway.session_model_pool import get_session_model_pool as _get_pool
+                _cfg = user_config if user_config else _load_gateway_config()
+                _pool = _get_pool(_cfg)
+                if _pool and _pool.enabled:
+                    # Always call acquire_session_slot — it is thread-safe
+                    # internally and refreshes the session timestamp on every
+                    # call. This prevents premature eviction of active sessions
+                    # and avoids a TOCTOU race between the local cache check
+                    # and the pool's own state.
+                    _pool_assign = _pool.acquire_session_slot(resolved_session_key)
+                    if _pool_assign:
+                        with self._pool_assigned_models_lock:
+                            self._pool_assigned_models[resolved_session_key] = _pool_assign
+                        model = _pool_assign.get("model", model)
+                        # Stash pool provider so it can be injected into
+                        # runtime_kwargs after _resolve_runtime_agent_kwargs().
+                        _pool_provider_override = _pool_assign.get("provider")
+                        logger.debug(
+                            "SessionModelPool: session=%s using pool-assigned model=%s provider=%s",
+                            resolved_session_key, model, _pool_assign.get("provider"),
+                        )
+            except Exception as _pool_exc:
+                logger.debug("SessionModelPool lookup failed: %s", _pool_exc)
+        # --- End Session Model Pool integration ---
+
         if override:
             override_model = override.get("model", model)
             override_runtime = {
@@ -2577,6 +2655,10 @@ def _resolve_session_agent_runtime(
                 resolved_session_key, model, runtime_kwargs
             )
 
+        # Apply pool-assigned provider (set during pool integration above).
+        if not override and _pool_provider_override:
+            runtime_kwargs["provider"] = _pool_provider_override
+
         # When the config has no model.default but a provider was resolved
         # (e.g. user ran `hermes auth add openai-codex` without `hermes model`),
         # fall back to the provider's first catalog model so the API call
@@ -8743,6 +8825,16 @@ async def _handle_message_with_agent(self, event, source, _quick_key: str, run_g
             self._set_session_reasoning_override(session_key, None)
             if hasattr(self, "_pending_model_notes"):
                 self._pending_model_notes.pop(session_key, None)
+            # Release pool-assigned slot for the reset session.
+            self._release_pool_slot(session_key)
+            # Clear manual override so the pool can reassign on next turn.
+            try:
+                from gateway.session_model_pool import get_session_model_pool as _get_pool_rst
+                _p_rst = _get_pool_rst({})
+                if _p_rst:
+                    _p_rst.clear_manual_override(session_key)
+            except Exception as _exc:
+                logger.debug("SessionModelPool: failed to clear override for %s: %s", session_key, _exc)
 
         # Emit session:start for new or auto-reset sessions
         _is_new_session = (
@@ -10953,6 +11045,11 @@ async def _on_model_selected(
                             "api_mode": result.api_mode,
                         }
 
+                        # Release pool-assigned slot for this session if one
+                        # exists — the manual override takes precedence.
+                        self._release_pool_slot(_session_key)
+                        self._mark_pool_override(_session_key)
+
                         # Evict cached agent so the next turn creates a fresh
                         # agent from the override rather than relying on the
                         # stale cache signature to trigger a rebuild.
@@ -11107,6 +11204,11 @@ async def _on_model_selected(
             "api_mode": result.api_mode,
         }
 
+        # Release pool-assigned slot for this session if one exists —
+        # the manual override takes precedence.
+        self._release_pool_slot(session_key)
+        self._mark_pool_override(session_key)
+
         # Evict cached agent so the next turn creates a fresh agent from the
         # override rather than relying on cache signature mismatch detection.
         self._evict_cached_agent(session_key)