From 4281151ae859241351ba14d8c7682dc67ff4c126 Mon Sep 17 00:00:00 2001
From: kshitijk4poor <82637225+kshitijk4poor@users.noreply.github.com>
Date: Sat, 11 Jul 2026 23:33:48 +0530
Subject: [PATCH] test(gateway): cover effective context budget

---
 gateway/run.py                                |  1 +
 .../test_context_ref_expansion_runtime.py     | 46 +++++++++++++++++++
 2 files changed, 47 insertions(+)

diff --git a/gateway/run.py b/gateway/run.py
index b36e69bda0a..6e5d8eb5d66 100644
--- a/gateway/run.py
+++ b/gateway/run.py
@@ -10613,6 +10613,7 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
                 _msg_cwd = os.environ.get("TERMINAL_CWD", os.path.expanduser("~"))
                 _msg_config_ctx = None
                 _msg_cfg = None
+                _msg_model_cfg = {}
                 _msg_custom_providers = []
                 try:
                     _msg_cfg = _load_gateway_config()
diff --git a/tests/gateway/test_context_ref_expansion_runtime.py b/tests/gateway/test_context_ref_expansion_runtime.py
index fd268fac26f..ee98d57da10 100644
--- a/tests/gateway/test_context_ref_expansion_runtime.py
+++ b/tests/gateway/test_context_ref_expansion_runtime.py
@@ -279,6 +279,52 @@ async def test_at_reference_passes_compatible_custom_provider_context(monkeypatc
     assert captured["custom_providers"] == custom_providers
 
 
+@pytest.mark.asyncio
+async def test_at_reference_applies_custom_runtime_budget_to_preprocessor(monkeypatch):
+    """The custom runtime's real budget must reach reference expansion."""
+    runner = _make_runner()
+    source = _source()
+    captured = {}
+    custom_providers = [{
+        "name": "private",
+        "base_url": "https://private.example/v1",
+        "models": {"session/model": {"context_length": 32768}},
+    }]
+    monkeypatch.setattr(
+        gateway_run,
+        "_load_gateway_config",
+        lambda: {"model": {"default": "global/model", "context_length": 128000}, "custom_providers": custom_providers},
+    )
+    monkeypatch.setattr(runner, "_resolve_session_agent_runtime", lambda **_kwargs: (
+        "session/model",
+        {"provider": "custom:private", "api_key": "test", "base_url": "https://private.example/v1"},
+    ))
+
+    import hermes_cli.config as config_mod
+    import agent.model_metadata as model_meta_mod
+    import agent.context_references as ctx_mod
+
+    monkeypatch.setattr(config_mod, "get_compatible_custom_providers", lambda _cfg: custom_providers)
+    monkeypatch.setattr(config_mod, "get_custom_provider_context_length", lambda **_kwargs: 32768)
+
+    async def _fake_get_context(_model, **kwargs):
+        captured["config_context_length"] = kwargs["config_context_length"]
+        return kwargs["config_context_length"]
+
+    async def _preprocess(message, *, context_length, **_kwargs):
+        captured["preprocessor_budget"] = context_length
+        return ContextReferenceResult(message="expanded", original_message=message, expanded=True)
+
+    monkeypatch.setattr(model_meta_mod, "get_model_context_length_async", _fake_get_context)
+    monkeypatch.setattr(ctx_mod, "preprocess_context_references_async", _preprocess)
+
+    result = await runner._prepare_inbound_message_text(
+        event=MessageEvent(text="@file:note", source=source), source=source, history=[]
+    )
+    assert result == "expanded"
+    assert captured == {"config_context_length": 32768, "preprocessor_budget": 32768}
+
+
 @pytest.mark.asyncio
 async def test_at_reference_ignores_global_context_for_session_model_override(monkeypatch):
     """A session model override must not inherit another model's global limit."""