Merge pull request #4082 from nesquena/stage-mm

Release MM (v0.51.374): custom-provider context-length probes carry the API key (#4059)

CHANGELOG.md

@@ -3,6 +3,12 @@
 
 ## [Unreleased]
 
+## [v0.51.374] — 2026-06-12 — Release MM (custom-provider context-length probes carry the API key)
+
+### Fixed
+
+- **Custom provider context-length probes now include the provider API key during session hydration and streaming fallbacks (#4059).** Static session loads and fallback usage/session-save paths now reuse the matched custom provider credentials when querying `/v1/models`, so authenticated custom endpoints do not fall back to the default 256K window and clobber a larger persisted context length / compression threshold. (#4059)
+
 ## [v0.51.373] — 2026-06-12 — Release ML (sidebar lineage + source-count fixes)
 
 ### Fixed

api/routes.py

@@ -2297,7 +2297,7 @@ def _model_matches_configured_default(
 
 
 class _ContextLengthLookupInputs:
-    __slots__ = ("config_context_length", "custom_providers", "base_url", "provider")
+    __slots__ = ("config_context_length", "custom_providers", "base_url", "provider", "api_key")
 
     def __init__(
         self,
@@ -2306,11 +2306,13 @@ class _ContextLengthLookupInputs:
         custom_providers: list | None = None,
         base_url: str = "",
         provider: str = "",
+        api_key: str = "",
     ) -> None:
         self.config_context_length = config_context_length
         self.custom_providers = custom_providers
         self.base_url = base_url
         self.provider = provider
+        self.api_key = api_key
 
 
 def _positive_context_length(value) -> int | None:
@@ -2400,6 +2402,44 @@ def _providers_match_for_context(config_key: object, requested_provider: str) ->
     )
 
 
+def _custom_provider_api_key_for_context(entry: dict, provider: str) -> str:
+    """Resolve the API key for a matched ``custom_providers`` entry.
+
+    Static session hydration/update routes already have a per-profile config
+    snapshot. Resolve from the matched entry instead of re-reading global config,
+    while preserving the same literal, ``${ENV_VAR}``, ``key_env``, and
+    sanitized-env shapes used by streaming/provider resolution.
+    """
+    raw_api_key = entry.get("api_key")
+    if raw_api_key is not None:
+        api_key_text = str(raw_api_key).strip()
+        if api_key_text.startswith("${") and api_key_text.endswith("}") and len(api_key_text) > 3:
+            env_name = api_key_text[2:-1]
+            resolved = os.getenv(env_name, "").strip()
+            if resolved:
+                return resolved
+            logger.debug(
+                "Custom provider %s api_key references %s, but the environment variable is unset or empty",
+                provider,
+                api_key_text,
+            )
+        elif api_key_text:
+            return api_key_text
+
+    key_env = str(entry.get("key_env") or "").strip()
+    if key_env:
+        resolved = os.getenv(key_env, "").strip()
+        if resolved:
+            return resolved
+
+    try:
+        from api.config import _lookup_custom_api_key_env
+
+        return _lookup_custom_api_key_env(provider) or ""
+    except Exception:
+        return ""
+
+
 def _context_length_lookup_inputs_for_model(
     model: str | None,
     provider: str | None = None,
@@ -2459,6 +2499,7 @@ def _context_length_lookup_inputs_for_model(
             break
 
     custom_context_length = None
+    effective_api_key = ""
     if custom_providers:
         target_base = effective_base_url.rstrip("/")
         model_candidates = set(_model_lookup_candidates(bare_model or model_for_lookup))
@@ -2488,6 +2529,7 @@ def _context_length_lookup_inputs_for_model(
                 effective_provider = entry_slug
             if not effective_base_url and entry_base:
                 effective_base_url = entry_base
+            effective_api_key = _custom_provider_api_key_for_context(entry, effective_provider or entry_slug)
             custom_context_length = _models_config_context_length(models_cfg, bare_model or model_for_lookup)
             break
 
@@ -2510,6 +2552,7 @@ def _context_length_lookup_inputs_for_model(
         custom_providers=custom_providers,
         base_url=effective_base_url,
         provider=effective_provider,
+        api_key=effective_api_key,
     )
 
 
@@ -3023,6 +3066,7 @@ def _resolve_context_length_for_session_model(
             return _get_cl(
                 model_for_lookup,
                 _ctx_lookup.base_url,
+                api_key=_ctx_lookup.api_key,
                 config_context_length=_ctx_lookup.config_context_length,
                 provider=_ctx_lookup.provider or provider or "",
                 custom_providers=_ctx_lookup.custom_providers,

api/streaming.py

@@ -7294,11 +7294,13 @@ def _run_agent_streaming(
                         )
                         _cfg_ctx_len = _ctx_lookup.config_context_length
                         _cfg_custom_providers = _ctx_lookup.custom_providers
+                        _cfg_api_key = _ctx_lookup.api_key or getattr(agent, 'api_key', '') or resolved_api_key or ''
                         _cfg_base_url = _ctx_lookup.base_url or _cfg_base_url
                         _cfg_provider = _ctx_lookup.provider or resolved_provider or ''
                         _resolved_cl = get_model_context_length(
                             getattr(agent, 'model', resolved_model or '') or '',
                             _cfg_base_url,
+                            api_key=_cfg_api_key,
                             config_context_length=_cfg_ctx_len,
                             provider=_cfg_provider,
                             custom_providers=_cfg_custom_providers,
@@ -7552,12 +7554,14 @@ def _run_agent_streaming(
                     )
                     _cfg_ctx_len = _ctx_lookup.config_context_length
                     _cfg_custom_providers = _ctx_lookup.custom_providers
+                    _cfg_api_key = _ctx_lookup.api_key or getattr(agent, 'api_key', '') or resolved_api_key or ''
                     _cfg_base_url = _ctx_lookup.base_url
                     _cfg_provider = _ctx_lookup.provider or resolved_provider or ''
                     try:
                         _fb_cl = _get_cl(
                             getattr(agent, 'model', resolved_model or '') or '',
                             _cfg_base_url,
+                            api_key=_cfg_api_key,
                             config_context_length=_cfg_ctx_len,
                             provider=_cfg_provider,
                             custom_providers=_cfg_custom_providers,

tests/test_issue1896_context_length_fallback_args.py

@@ -21,6 +21,8 @@ config-override args again.
 """
 
 from pathlib import Path
+import sys
+import types
 
 
 REPO = Path(__file__).resolve().parent.parent
@@ -233,3 +235,178 @@ def test_routes_session_load_fallback_passes_config_overrides():
         "session-load fallback must catch TypeError to support older "
         "hermes-agent builds without the new kwargs."
     )
+
+
+def test_context_lookup_returns_custom_provider_api_key_from_entry():
+    """#4059: static session hydration must carry custom-provider API keys.
+
+    A named ``custom_providers`` entry can require auth for its ``/v1/models``
+    endpoint. The route-side lookup helper already identifies the matching
+    provider/base/model; it must also return that entry's API key so
+    ``get_model_context_length`` does not probe anonymously and fall back to
+    256K.
+    """
+    from api.routes import _context_length_lookup_inputs_for_model
+
+    lookup = _context_length_lookup_inputs_for_model(
+        "custom-model-id",
+        "custom:llm-proxy",
+        cfg={
+            "custom_providers": [
+                {
+                    "name": "llm-proxy",
+                    "base_url": "https://llm.example.test/v1",
+                    "api_key": "sk-test-entry",
+                    "model": "custom-model-id",
+                }
+            ]
+        },
+    )
+
+    assert lookup.provider == "custom:llm-proxy"
+    assert lookup.base_url == "https://llm.example.test/v1"
+    assert lookup.api_key == "sk-test-entry"
+
+
+def test_context_lookup_resolves_custom_provider_api_key_env_template(monkeypatch):
+    """#4059: ``${ENV_VAR}`` custom-provider keys resolve before metadata probes."""
+    from api.routes import _context_length_lookup_inputs_for_model
+
+    monkeypatch.setenv("ISSUE_4059_CONTEXT_KEY", "env-template-key")
+
+    lookup = _context_length_lookup_inputs_for_model(
+        "custom-model-id",
+        "custom:llm-proxy",
+        cfg={
+            "custom_providers": [
+                {
+                    "name": "llm-proxy",
+                    "base_url": "https://llm.example.test/v1",
+                    "api_key": "${ISSUE_4059_CONTEXT_KEY}",
+                    "model": "custom-model-id",
+                }
+            ]
+        },
+    )
+
+    assert lookup.api_key == "env-template-key"
+
+
+def test_context_lookup_logs_unresolved_custom_provider_api_key_env_template(monkeypatch, caplog):
+    """#4059: unresolved ``${ENV_VAR}`` keys get a DEBUG diagnostic.
+
+    Behavior stays permissive: after logging the unresolved template, the helper
+    still falls through to ``key_env`` and sanitized provider env lookup.
+    """
+    import logging
+
+    from api.routes import _context_length_lookup_inputs_for_model
+
+    monkeypatch.delenv("ISSUE_4059_MISSING_CONTEXT_KEY", raising=False)
+    monkeypatch.setenv("ISSUE_4059_KEY_ENV_AFTER_TEMPLATE", "key-env-fallback")
+    caplog.set_level(logging.DEBUG, logger="api.routes")
+
+    lookup = _context_length_lookup_inputs_for_model(
+        "custom-model-id",
+        "custom:llm-proxy",
+        cfg={
+            "custom_providers": [
+                {
+                    "name": "llm-proxy",
+                    "base_url": "https://llm.example.test/v1",
+                    "api_key": "${ISSUE_4059_MISSING_CONTEXT_KEY}",
+                    "key_env": "ISSUE_4059_KEY_ENV_AFTER_TEMPLATE",
+                    "model": "custom-model-id",
+                }
+            ]
+        },
+    )
+
+    assert lookup.api_key == "key-env-fallback"
+    assert "${ISSUE_4059_MISSING_CONTEXT_KEY}" in caplog.text
+    assert "unset or empty" in caplog.text
+
+
+def test_context_lookup_resolves_custom_provider_key_env(monkeypatch):
+    """#4059: ``key_env`` custom-provider keys resolve before metadata probes."""
+    from api.routes import _context_length_lookup_inputs_for_model
+
+    monkeypatch.setenv("ISSUE_4059_KEY_ENV", "key-env-value")
+
+    lookup = _context_length_lookup_inputs_for_model(
+        "custom-model-id",
+        "custom:llm-proxy",
+        cfg={
+            "custom_providers": [
+                {
+                    "name": "llm-proxy",
+                    "base_url": "https://llm.example.test/v1",
+                    "key_env": "ISSUE_4059_KEY_ENV",
+                    "model": "custom-model-id",
+                }
+            ]
+        },
+    )
+
+    assert lookup.api_key == "key-env-value"
+
+
+def test_routes_session_model_resolver_passes_custom_provider_api_key(monkeypatch):
+    """#4059: ``_resolve_context_length_for_session_model`` passes api_key.
+
+    This is the session-load/static-update path that was clobbering persisted
+    500K context metadata back to the unauthenticated 256K fallback.
+    """
+    from api import config as cfg_mod
+    from api import routes
+
+    # Some unrelated route tests install ``sys.modules["agent"]`` as a plain
+    # module stub at collection time. Make this regression test own a package-
+    # shaped temporary ``agent.model_metadata`` import target so the production
+    # helper's local import is order-independent.
+    fake_agent = types.ModuleType("agent")
+    fake_agent.__path__ = []
+    metadata = types.ModuleType("agent.model_metadata")
+    fake_agent.model_metadata = metadata  # type: ignore[attr-defined]
+    monkeypatch.setitem(sys.modules, "agent", fake_agent)
+    monkeypatch.setitem(sys.modules, "agent.model_metadata", metadata)
+
+    seen = {}
+
+    monkeypatch.setattr(
+        cfg_mod,
+        "get_config",
+        lambda: {
+            "custom_providers": [
+                {
+                    "name": "llm-proxy",
+                    "base_url": "https://llm.example.test/v1",
+                    "api_key": "sk-test-route",
+                    "model": "custom-model-id",
+                }
+            ]
+        },
+    )
+
+    def fake_get_model_context_length(model, base_url, **kwargs):
+        seen.update(model=model, base_url=base_url, kwargs=kwargs)
+        return 500_000 if kwargs.get("api_key") == "sk-test-route" else 256_000
+
+    monkeypatch.setattr(metadata, "get_model_context_length", fake_get_model_context_length, raising=False)
+
+    assert routes._resolve_context_length_for_session_model(
+        "custom-model-id",
+        "custom:llm-proxy",
+    ) == 500_000
+    assert seen["kwargs"]["api_key"] == "sk-test-route"
+
+
+def test_streaming_context_length_fallbacks_pass_api_key():
+    """#4059: both streaming fallback probes must pass the resolved api_key."""
+    blocks = _both_callsites()
+    for i, block in enumerate(blocks):
+        assert "api_key=" in block, (
+            f"Callsite #{i+1} is missing api_key=. Authenticated custom "
+            f"provider /v1/models probes then fail and fall back to 256K. "
+            f"See #4059.\n\nBlock:\n{block}"
+        )

tests/test_issue3256_context_length_default_only_guard.py

@@ -20,8 +20,9 @@ def _install_fake_get_model_context_length(monkeypatch, recorder):
     """Install a fake get_model_context_length into a stand-in agent.model_metadata."""
     mod = types.ModuleType("agent.model_metadata")
 
-    def _fake(model, base_url="", config_context_length=None, provider="", custom_providers=None):
+    def _fake(model, base_url="", api_key="", config_context_length=None, provider="", custom_providers=None):
         recorder["model"] = model
+        recorder["api_key"] = api_key
         recorder["config_context_length"] = config_context_length
         # Pretend the real per-model metadata window is 1,000,000 unless the
         # caller forced a config cap, in which case honor the cap (mirrors the