diff --git a/contextual_orchestrator/model_discovery.py b/contextual_orchestrator/model_discovery.py index 761220ee4..a5f2d72aa 100644 --- a/contextual_orchestrator/model_discovery.py +++ b/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,11 +266,14 @@ 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: @@ -272,6 +281,9 @@ def probe_discovered_model_tool_call_capability( 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: diff --git a/tests/test_model_discovery_boundaries.py b/tests/test_model_discovery_boundaries.py index 5e193da25..593001c95 100644 --- a/tests/test_model_discovery_boundaries.py +++ b/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")