diff --git a/src/harness_sdk/gen_ai/__init__.py b/src/harness_sdk/gen_ai/__init__.py index f44eec7..15109da 100644 --- a/src/harness_sdk/gen_ai/__init__.py +++ b/src/harness_sdk/gen_ai/__init__.py @@ -2,6 +2,6 @@ from __future__ import annotations -from harness_sdk.gen_ai.exceptions import ControlEvaluationBlocked +from harness_sdk.gen_ai.exceptions import ControlEvaluationBlocked, ControlRequestBlocked __all__ = ["ControlEvaluationBlocked"] diff --git a/src/harness_sdk/gen_ai/exceptions.py b/src/harness_sdk/gen_ai/exceptions.py index 095387d..7d316c6 100644 --- a/src/harness_sdk/gen_ai/exceptions.py +++ b/src/harness_sdk/gen_ai/exceptions.py @@ -15,3 +15,12 @@ def __init__(self, result: "ControlResult") -> None: self.result = result msg = getattr(result, "response_message", None) or "Forbidden" super().__init__(msg) + + +class ControlRequestBlocked(Exception): + """Raised when control registry ``evaluate`` returns ``block=True`` for an outbound request.""" + + def __init__(self, result: "ControlResult") -> None: + self.result = result + msg = getattr(result, "response_message", None) or "Forbidden" + super().__init__(msg) diff --git a/src/harness_sdk/instrumentation/httpx/__init__.py b/src/harness_sdk/instrumentation/httpx/__init__.py index e40ef04..5619756 100644 --- a/src/harness_sdk/instrumentation/httpx/__init__.py +++ b/src/harness_sdk/instrumentation/httpx/__init__.py @@ -4,6 +4,7 @@ from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor from harness_sdk.plugins.control import get_control_registry +from harness_sdk.gen_ai.exceptions import ControlRequestBlocked from harness_sdk.instrumentation import BaseInstrumentorWrapper from harness_sdk.instrumentation.httpx.utils import ( decode_response_body_for_capture, @@ -42,7 +43,14 @@ def _process_request(self, span, request_info): headers = headers_from_httpx(request_info.headers) body = read_request_body(request_info.stream) self.generic_request_handler(headers, body, span) - get_control_registry().evaluate(span, url, headers, body, False) + control_result = get_control_registry().evaluate(span, url, headers, body, False) + if control_result.block: + logger.debug( + "httpx request blocked by control plugin: url=%s status=%s", + url, + control_result.response_status_code, + ) + raise ControlRequestBlocked(control_result) def _process_response(self, span, response_info): headers = headers_from_httpx(response_info.headers) diff --git a/test/instrumentation/httpx_client/test_httpx_blocking.py b/test/instrumentation/httpx_client/test_httpx_blocking.py new file mode 100644 index 0000000..b51d7e8 --- /dev/null +++ b/test/instrumentation/httpx_client/test_httpx_blocking.py @@ -0,0 +1,59 @@ +"""Tests for httpx client blocking via control plugins.""" +import httpx +import pytest + +from harness_sdk.gen_ai.exceptions import ControlRequestBlocked +from harness_sdk.instrumentation.httpx import HTTPXClientInstrumentorWrapper +from harness_sdk.plugins.control import ControlResult, get_control_registry +from opentelemetry import trace +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + + +class _BlockingPlugin: + name = "test_blocking" + provides_blocking = True + + def on_init(self, config): # pylint: disable=unused-argument + pass + + def evaluate(self, span, url, headers, body, is_grpc): # pylint: disable=unused-argument + return ControlResult( + block=True, + response_status_code=403, + response_message="Blocked by policy", + ) + + def evaluate_agent_span(self, span, body=""): # pylint: disable=unused-argument + return ControlResult() + + def shutdown(self): + pass + + +@pytest.fixture +def httpx_wrapper(): + wrapper = HTTPXClientInstrumentorWrapper() + yield wrapper + if wrapper.is_instrumented_by_opentelemetry: + wrapper.uninstrument() + + +def test_httpx_request_hook_raises_when_control_blocks(httpx_wrapper): + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + trace.set_tracer_provider(provider) + + get_control_registry().register(_BlockingPlugin()) + httpx_wrapper.instrument(tracer_provider=provider) + + request_info = httpx.Request("GET", "https://example.com/get", headers={"x-test": "1"}) + tracer = provider.get_tracer("test") + with tracer.start_as_current_span("httpx") as span: + with pytest.raises(ControlRequestBlocked) as exc_info: + httpx_wrapper._process_request(span, request_info) # pylint: disable=protected-access + + assert exc_info.value.result.response_status_code == 403 + get_control_registry().clear()