diff --git a/src/harness_sdk/instrumentation/grpc/__init__.py b/src/harness_sdk/instrumentation/grpc/__init__.py index 728f8a3..bdb4d83 100644 --- a/src/harness_sdk/instrumentation/grpc/__init__.py +++ b/src/harness_sdk/instrumentation/grpc/__init__.py @@ -56,7 +56,8 @@ def server_wrapper(*args, **kwargs) -> None: logger.debug('Entering wrapper interceptors set') logger.debug( 'Setting server_interceptor_wrapper() as interceptor.') - kwargs["interceptors"] = [server_interceptor_wrapper(self)] + existing = list(kwargs.get("interceptors") or []) + kwargs["interceptors"] = [server_interceptor_wrapper(self)] + existing return self._original_wrapper_func(*args, **kwargs) grpc.server = server_wrapper diff --git a/src/harness_sdk/instrumentation/requests/__init__.py b/src/harness_sdk/instrumentation/requests/__init__.py index 7e265d8..7accd6f 100644 --- a/src/harness_sdk/instrumentation/requests/__init__.py +++ b/src/harness_sdk/instrumentation/requests/__init__.py @@ -2,6 +2,7 @@ import logging from opentelemetry.instrumentation.requests import RequestsInstrumentor +from harness_sdk.gen_ai.exceptions import ControlRequestBlocked from harness_sdk.plugins.control import get_control_registry from harness_sdk.instrumentation import BaseInstrumentorWrapper @@ -28,7 +29,16 @@ def request_hook(self, span, request_obj): '''capture request data''' url = request_obj.url self.generic_request_handler(request_obj.headers, request_obj.body, span) - get_control_registry().evaluate(span, url, request_obj.headers, request_obj.body, False) + control_result = get_control_registry().evaluate( + span, url, request_obj.headers, request_obj.body, False + ) + if control_result.block: + logger.debug( + "requests request blocked by control plugin: url=%s status=%s", + url, + control_result.response_status_code, + ) + raise ControlRequestBlocked(control_result) def response_hook(self, span, _, response): diff --git a/test/instrumentation/grpc/test_grpc_server_interceptors.py b/test/instrumentation/grpc/test_grpc_server_interceptors.py new file mode 100644 index 0000000..38d0ba1 --- /dev/null +++ b/test/instrumentation/grpc/test_grpc_server_interceptors.py @@ -0,0 +1,30 @@ +"""grpc.server wrapper must preserve user-supplied interceptors.""" +from unittest.mock import MagicMock, patch + +import harness_sdk.instrumentation.grpc as grpc_module +from harness_sdk.instrumentation.grpc import GrpcInstrumentorServerWrapper + + +def test_server_wrapper_prepends_harness_interceptor(): + wrapper = GrpcInstrumentorServerWrapper() + existing_interceptor = object() + original_server = MagicMock(return_value="server-instance") + + with patch( + "harness_sdk.instrumentation.grpc.GrpcInstrumentorServer._instrument", + return_value=None, + ): + with patch( + "harness_sdk.instrumentation.grpc.server_interceptor_wrapper", + return_value="harness-interceptor", + ): + with patch("harness_sdk.instrumentation.grpc.grpc.server", original_server): + wrapper._instrument() # pylint: disable=protected-access + + grpc_module.grpc.server( + futures=MagicMock(), + interceptors=[existing_interceptor], + ) + + _, kwargs = original_server.call_args + assert kwargs["interceptors"] == ["harness-interceptor", existing_interceptor] diff --git a/test/instrumentation/requests/test_requests_blocking.py b/test/instrumentation/requests/test_requests_blocking.py new file mode 100644 index 0000000..cae45aa --- /dev/null +++ b/test/instrumentation/requests/test_requests_blocking.py @@ -0,0 +1,60 @@ +"""Tests for requests client blocking via control plugins.""" +import pytest +import requests + +from harness_sdk.gen_ai.exceptions import ControlRequestBlocked +from harness_sdk.instrumentation.requests import RequestsInstrumentorWrapper +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 requests_wrapper(): + wrapper = RequestsInstrumentorWrapper() + yield wrapper + if wrapper.is_instrumented_by_opentelemetry: + wrapper.uninstrument() + + +def test_requests_request_hook_raises_when_control_blocks(requests_wrapper): + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + trace.set_tracer_provider(provider) + + get_control_registry().register(_BlockingPlugin()) + requests_wrapper.instrument(tracer_provider=provider) + + request_obj = requests.Request("GET", "https://example.com/get", headers={"x-test": "1"}) + prepared = request_obj.prepare() + tracer = provider.get_tracer("test") + with tracer.start_as_current_span("requests") as span: + with pytest.raises(ControlRequestBlocked) as exc_info: + requests_wrapper.request_hook(span, prepared) + + assert exc_info.value.result.response_status_code == 403 + get_control_registry().clear()