fix(discovery): bound tool-call probe response reads (#1041)

contextual_orchestrator/model_discovery.py

@@ -195,6 +195,12 @@ def probe_discovered_model_tool_call_capability(
     This is real runtime evidence, not a model-name heuristic. It is deliberately
     separate from :func:`discover_all_models` so callers decide when the extra
     latency and token cost are justified.
+
+    Both the success body and the 400 error body are untrusted network input,
+    so each read is capped at :data:`MAX_DISCOVERY_RESPONSE_BYTES` (plus one
+    byte to detect an overage) exactly like the sibling discovery fetches; an
+    oversized body is treated as ambiguous evidence and returns ``None``
+    rather than being buffered whole.
     """
     api_key = get_credential(discovered.credential_name)
     if not api_key:
@@ -260,18 +266,24 @@ def probe_discovered_model_tool_call_capability(
     request = urllib.request.Request(url, data=data, headers=headers, method="POST")
     try:
         with client._open_provider(request, destination, timeout=timeout) as response:
-            body = response.read().decode("utf-8", errors="replace")
+            body = response.read(MAX_DISCOVERY_RESPONSE_BYTES + 1)
     except urllib.error.HTTPError as exc:
         if exc.code != 400:
             return None
-        body = exc.read().decode("utf-8", errors="replace")
+        error_body = exc.read(MAX_DISCOVERY_RESPONSE_BYTES + 1)
+        if len(error_body) > MAX_DISCOVERY_RESPONSE_BYTES:
+            return None
+        body = error_body.decode("utf-8", errors="replace")
         try:
             error_payload = json.loads(body)
         except json.JSONDecodeError:
             error_payload = {"message": body}
         return _tool_call_parallelism_from_error(error_payload)
     except (urllib.error.URLError, OSError, TimeoutError, ValueError):
         return None
+    if len(body) > MAX_DISCOVERY_RESPONSE_BYTES:
+        return None
+    body = body.decode("utf-8", errors="replace")
     try:
         response_payload = json.loads(body)
     except json.JSONDecodeError:

tests/test_model_discovery_boundaries.py

@@ -8,6 +8,8 @@
 
 from __future__ import annotations
 
+import io
+import json
 import ssl
 import urllib.error
 from dataclasses import replace
@@ -33,6 +35,7 @@
     _valid_price_component,
     agent_from_discovered,
     discover_provider_models,
+    probe_discovered_model_tool_call_capability,
     refresh_price_book,
     select_bootstrap_discovered_agents,
     select_top_n_cheapest_discovered_agents,
@@ -212,6 +215,157 @@ def read(self, amt: int | None = None) -> bytes:
     assert reads == [MAX_DISCOVERY_RESPONSE_BYTES + 1]
 
 
+def _tool_call_probe_model() -> DiscoveredModel:
+    return DiscoveredModel(
+        provider_name="openrouter",
+        model_id="probe-model",
+        credential_name="OPENROUTER_API_KEY",
+        chat_base_url="https://openrouter.example/v1",
+        auth_scheme="Bearer",
+    )
+
+
+def test_tool_call_probe_caps_success_body_read() -> None:
+    """The capability probe reads the success body through the shared size bound.
+
+    Regression for the unbounded ``response.read()`` in
+    :func:`probe_discovered_model_tool_call_capability`: a compromised or
+    misbehaving provider could stream an arbitrarily large body into memory
+    before JSON parsing ever runs. The probe must request at most
+    ``MAX_DISCOVERY_RESPONSE_BYTES + 1`` bytes, exactly like its sibling
+    fetches, and still return its evidence on a normal body.
+    """
+    register_credential("OPENROUTER_API_KEY", "sk-router")
+    reads: list[int | None] = []
+    response = _Response(
+        {
+            "choices": [
+                {
+                    "message": {
+                        "tool_calls": [
+                            {"type": "function", "function": {"name": "probe_a"}},
+                            {"type": "function", "function": {"name": "probe_b"}},
+                        ]
+                    }
+                }
+            ]
+        }
+    )
+    original_read = response.read
+
+    def bounded_read(amt: int | None = None) -> bytes:
+        reads.append(amt)
+        return original_read(amt)
+
+    response.read = bounded_read  # type: ignore[method-assign]
+    with (
+        patch(
+            "contextual_orchestrator.model_discovery.ModelClient._validate_provider",
+            return_value=object(),
+        ),
+        patch(
+            "contextual_orchestrator.model_discovery.ModelClient._open_provider",
+            return_value=response,
+        ),
+    ):
+        assert probe_discovered_model_tool_call_capability(_tool_call_probe_model()) is True
+    assert reads == [MAX_DISCOVERY_RESPONSE_BYTES + 1]
+
+
+def test_tool_call_probe_rejects_oversized_success_body_without_buffering_it() -> None:
+    """An oversized probe body is fail-closed to ``None``, never fully buffered."""
+    register_credential("OPENROUTER_API_KEY", "sk-router")
+    oversized = b"0" * (MAX_DISCOVERY_RESPONSE_BYTES + 1024)
+    reads: list[int | None] = []
+
+    class OversizedResponse:
+        def __enter__(self):
+            return self
+
+        def __exit__(self, *_args):
+            return False
+
+        def read(self, amt: int | None = None) -> bytes:
+            reads.append(amt)
+            return oversized if amt is None else oversized[:amt]
+
+    with (
+        patch(
+            "contextual_orchestrator.model_discovery.ModelClient._validate_provider",
+            return_value=object(),
+        ),
+        patch(
+            "contextual_orchestrator.model_discovery.ModelClient._open_provider",
+            return_value=OversizedResponse(),
+        ),
+    ):
+        assert probe_discovered_model_tool_call_capability(_tool_call_probe_model()) is None
+    assert reads == [MAX_DISCOVERY_RESPONSE_BYTES + 1]
+
+
+def test_tool_call_probe_caps_single_call_rejection_body_read() -> None:
+    """A 400 body is read through the same bound before the negative verdict.
+
+    The explicit single-tool-call rejection is the one 400 the probe trusts;
+    its body is still untrusted network input and must not be buffered whole.
+    """
+    register_credential("OPENROUTER_API_KEY", "sk-router")
+    oversized = b"0" * (MAX_DISCOVERY_RESPONSE_BYTES + 1024)
+    reads: list[int | None] = []
+
+    class OversizedHTTPError(urllib.error.HTTPError):
+        def read(self, amt: int | None = None) -> bytes:
+            reads.append(amt)
+            return oversized if amt is None else oversized[:amt]
+
+    error = OversizedHTTPError(
+        "https://openrouter.example/v1/chat/completions",
+        400,
+        "bad request",
+        None,
+        None,
+    )
+    with (
+        patch(
+            "contextual_orchestrator.model_discovery.ModelClient._validate_provider",
+            return_value=object(),
+        ),
+        patch(
+            "contextual_orchestrator.model_discovery.ModelClient._open_provider",
+            side_effect=error,
+        ),
+    ):
+        assert probe_discovered_model_tool_call_capability(_tool_call_probe_model()) is None
+    assert reads == [MAX_DISCOVERY_RESPONSE_BYTES + 1]
+
+
+def test_tool_call_probe_still_maps_in_budget_single_call_rejection_to_false() -> None:
+    """The new read bound must not swallow a normal explicit rejection body."""
+    register_credential("OPENROUTER_API_KEY", "sk-router")
+    error = urllib.error.HTTPError(
+        "https://openrouter.example/v1/chat/completions",
+        400,
+        "bad request",
+        None,
+        io.BytesIO(
+            json.dumps(
+                {"error": {"message": "this model only supports a single tool call"}}
+            ).encode()
+        ),
+    )
+    with (
+        patch(
+            "contextual_orchestrator.model_discovery.ModelClient._validate_provider",
+            return_value=object(),
+        ),
+        patch(
+            "contextual_orchestrator.model_discovery.ModelClient._open_provider",
+            side_effect=error,
+        ),
+    ):
+        assert probe_discovered_model_tool_call_capability(_tool_call_probe_model()) is False
+
+
 def test_malformed_json_maps_to_invalid_response_code() -> None:
     """A non-JSON provider body is invalid_response, not a crash."""
     register_credential("OPENAI_API_KEY", "sk-openai")