From 28baf2eeb280b7f4c8d3f35acd464fad81660adb Mon Sep 17 00:00:00 2001 From: Santosh Sahu Date: Mon, 21 Sep 2026 08:42:58 +0530 Subject: [PATCH] NO-TICKET: Fix gRPC interceptor chaining and requests client blocking Prepend the Harness gRPC server interceptor instead of replacing the interceptor list, and raise ControlRequestBlocked from the requests hook when policy blocks (parity with httpx). Co-authored-by: Cursor --- .../instrumentation/grpc/__init__.py | 3 +- .../instrumentation/requests/__init__.py | 12 +++- .../grpc/test_grpc_server_interceptors.py | 30 ++++++++++ .../requests/test_requests_blocking.py | 60 +++++++++++++++++++ 4 files changed, 103 insertions(+), 2 deletions(-) create mode 100644 test/instrumentation/grpc/test_grpc_server_interceptors.py create mode 100644 test/instrumentation/requests/test_requests_blocking.py 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()