diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 7c3c726..3dcbddf 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -42,6 +42,8 @@ jobs: git diff --stat exit 1 } + + # Plugin tests - name: Install nemocheck plugin dependencies working-directory: ./plugins/examples/nemocheck run: | @@ -50,3 +52,13 @@ jobs: - name: Run nemocheck plugin tests working-directory: ./plugins/examples/nemocheck run: uv run pytest tests + + # Server tests + - name: Install server test dependencies + run: | + echo "Installing pytest and dependencies for server tests..." + pip install pytest pytest-asyncio + - name: Run server unit tests + run: | + echo "Running server unit tests..." + uv run pytest tests diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index a2cd02a..7e76baf 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -4,6 +4,7 @@ repos: hooks: # Run the linter. - id: ruff - args: [ --fix ] + args: [ --fix, --line-length=80 ] # Run the formatter. - id: ruff-format + args: [ --line-length=80 ] diff --git a/plugins/examples/nemo/nemo_wrapper_plugin.py b/plugins/examples/nemo/nemo_wrapper_plugin.py index 047c46f..f7c1a06 100644 --- a/plugins/examples/nemo/nemo_wrapper_plugin.py +++ b/plugins/examples/nemo/nemo_wrapper_plugin.py @@ -62,20 +62,24 @@ async def tool_pre_invoke( rails_response = await self._rails.generate_async( messages=[{"role": "user", "content": payload_args}] ) - except ( - asyncio.CancelledError - ): # asyncio.exceptions.CancelledError is thrown by nemo, need to catch - logging.exception("An error occurred in the nemo plugin except block:") + except asyncio.CancelledError: # asyncio.exceptions.CancelledError is thrown by nemo, need to catch + logging.exception( + "An error occurred in the nemo plugin except block:" + ) finally: logger.warning("[NemoWrapperPlugin] Async rails executed") logger.warning(rails_response) if rails_response and "PII detected" in rails_response["content"]: - logger.warning("[NemoWrapperPlugin] PII detected, stopping processing") + logger.warning( + "[NemoWrapperPlugin] PII detected, stopping processing" + ) return ToolPreInvokeResult( modified_payload=payload, continue_processing=False ) logger.warning("[NemoWrapperPlugin] No PII detected, continuing") - return ToolPreInvokeResult(modified_payload=payload, continue_processing=True) + return ToolPreInvokeResult( + modified_payload=payload, continue_processing=True + ) async def tool_post_invoke( self, payload: ToolPostInvokePayload, context: PluginContext diff --git a/plugins/examples/nemocheck/plugin.py b/plugins/examples/nemocheck/plugin.py index 1751590..8d1f704 100644 --- a/plugins/examples/nemocheck/plugin.py +++ b/plugins/examples/nemocheck/plugin.py @@ -62,7 +62,9 @@ def __init__(self, config: PluginConfig): ) else: self.check_endpoint = DEFAULT_CHECK_ENDPOINT - logger.warning("Plugin config is empty or invalid, using default endpoint") + logger.warning( + "Plugin config is empty or invalid, using default endpoint" + ) logger.info(f"Nemo Check endpoint: {self.check_endpoint}") async def prompt_pre_fetch( @@ -105,9 +107,11 @@ async def tool_pre_invoke( Returns: The result of the plugin's analysis, including whether the tool can proceed. """ - logger.info("tool_pre_invoke....") - logger.info(payload) - tool_name = payload.name # ("tool_name", None) + logger.info( + f"[NemoCheck] Starting tool pre invoke hook with payload {payload}" + ) + + tool_name = payload.name check_nemo_payload = { "model": MODEL_NAME, "messages": [ @@ -119,7 +123,9 @@ async def tool_pre_invoke( "type": "function", "function": { "name": tool_name, - "arguments": payload.args.get("tool_args", None), + "arguments": payload.args.get( + "tool_args", None + ), }, } ], @@ -135,7 +141,7 @@ async def tool_pre_invoke( if response.status_code == 200: data = response.json() status = data.get("status", "blocked") - logger.debug(f"rails reply: {data}") + logger.debug(f"[NemoCheck] Rails reply: {data}") if status == "success": metadata = data.get("rails_status") @@ -147,7 +153,7 @@ async def tool_pre_invoke( violation = PluginViolation( reason=f"Check tool rails:{status}.", description=json.dumps(data), - code=f"checkserver_http_status_code:{response.status_code}", + code="NEMO_RAILS_BLOCKED", details=metadata, ) return ToolPreInvokeResult( @@ -158,23 +164,25 @@ async def tool_pre_invoke( else: violation = PluginViolation( reason="Tool Check Unavailable", - description="Tool arguments check server returned error", - code=f"checkserver_http_status_code:{response.status_code}", - details={}, + description=f"Tool arguments check server returned error. Status code: {response.status_code}, Response: {response.text}", + code="NEMO_SERVER_ERROR", + details={"status_code": response.status_code}, ) return ToolPreInvokeResult( continue_processing=False, violation=violation ) except Exception as e: - logger.error(f"Error calling Nemo Check endpoint: {e}") + logger.error(f"[NemoCheck] Error checking tool arguments: {e}") violation = PluginViolation( reason="Tool Check Error", description=f"Failed to connect to check server: {str(e)}", - code="checkserver_connection_error", - details={}, + code="NEMO_CONNECTION_ERROR", + details={"error": str(e)}, + ) + return ToolPreInvokeResult( + continue_processing=False, violation=violation ) - return ToolPreInvokeResult(continue_processing=False, violation=violation) async def tool_post_invoke( self, payload: ToolPostInvokePayload, context: PluginContext @@ -188,4 +196,90 @@ async def tool_post_invoke( Returns: The result of the plugin's analysis, including whether the tool result should proceed. """ - return ToolPostInvokeResult(continue_processing=True) + logger.info( + f"[NemoCheck] Starting tool post invoke hook with payload {payload}" + ) + + # Extract content from payload.result + # payload.result format: {'content': [{'type': 'text', 'text': 'Hello, bob!'}]} + result_content = payload.result.get("content", []) + tool_name = payload.name + + if not result_content: + logger.warning( + "[NemoCheck] No content in tool result, skipping check" + ) + return ToolPostInvokeResult(continue_processing=True) + + # Extract text content from the content array + # TODO: what to do if there's actually multiple texts? + text_content = "" + for item in result_content: + if item.get("type") == "text": + text_content += item.get("text", "") + + # Build NeMo check payload for tool response + check_nemo_payload = { + "model": MODEL_NAME, # ideally optional + "messages": [ + {"role": "tool", "content": text_content, "name": tool_name} + ], + } + + logger.debug( + f"[NemoCheck] Payload for guardrail check: {check_nemo_payload}" + ) + + violation = None + try: + response = requests.post( + self.check_endpoint, headers=HEADERS, json=check_nemo_payload + ) + if response.status_code == 200: + data = response.json() + status = data.get("status", "blocked") + logger.debug(f"[NemoCheck] Rails reply: {data}") + + if status == "success": + metadata = data.get("rails_status") + result = ToolPostInvokeResult( + continue_processing=True, metadata=metadata + ) + else: # blocked + metadata = data.get("rails_status") + violation = PluginViolation( + reason=f"Check tool rails:{status}.", + description=json.dumps(data), + code="NEMO_RAILS_BLOCKED", + details=metadata, + ) + result = ToolPostInvokeResult( + continue_processing=False, + violation=violation, + metadata=metadata, + ) + else: + violation = PluginViolation( + reason="Tool Check Unavailable", + description=f"Tool response check server returned error. Status code: {response.status_code}, Response: {response.text}", + code="NEMO_SERVER_ERROR", + details={"status_code": response.status_code}, + ) + result = ToolPostInvokeResult( + continue_processing=False, violation=violation + ) + + logger.info(f"[NemoCheck] Tool post invoke result: {result}") + return result + + except Exception as e: + logger.error(f"[NemoCheck] Error checking tool response: {e}") + violation = PluginViolation( + reason="Tool Check Error", + description=f"Failed to connect to check server: {str(e)}", + code="NEMO_CONNECTION_ERROR", + details={"error": str(e)}, + ) + return ToolPostInvokeResult( + continue_processing=False, violation=violation + ) diff --git a/plugins/examples/nemocheck/tests/test_all.py b/plugins/examples/nemocheck/tests/test_all.py index 3212f77..f2263e5 100644 --- a/plugins/examples/nemocheck/tests/test_all.py +++ b/plugins/examples/nemocheck/tests/test_all.py @@ -48,9 +48,13 @@ async def test_prompt_pre_hook(plugin_manager: PluginManager): async def test_prompt_post_hook(plugin_manager: PluginManager): """Test prompt post hook across all registered plugins.""" # Customize payload for testing - message = Message(content=TextContent(type="text", text="prompt"), role=Role.USER) + message = Message( + content=TextContent(type="text", text="prompt"), role=Role.USER + ) prompt_result = PromptResult(messages=[message]) - payload = PromptPosthookPayload(prompt_id="test_prompt", result=prompt_result) + payload = PromptPosthookPayload( + prompt_id="test_prompt", result=prompt_result + ) global_context = GlobalContext(request_id="1") result, _ = await plugin_manager.invoke_hook( PromptHookType.PROMPT_POST_FETCH, payload, global_context diff --git a/plugins/examples/nemocheck/tests/test_nemocheck.py b/plugins/examples/nemocheck/tests/test_nemocheck.py index 8cd5561..7c1cbd7 100644 --- a/plugins/examples/nemocheck/tests/test_nemocheck.py +++ b/plugins/examples/nemocheck/tests/test_nemocheck.py @@ -1,35 +1,274 @@ """Tests for plugin.""" +# Standard +from unittest.mock import Mock, patch + # Third-Party import pytest # First-Party from mcpgateway.plugins.framework import ( PluginConfig, + PluginContext, GlobalContext, PromptPrehookPayload, + ToolPostInvokePayload, + ToolPreInvokePayload, ) # Local from plugin import NemoCheck -@pytest.mark.asyncio -async def test_nemocheck(): - """Test plugin prompt prefetch hook.""" +@pytest.fixture +def plugin(): + """Create a NemoCheck plugin instance.""" config = PluginConfig( name="test", kind="nemocheck.NemoCheck", - hooks=["prompt_pre_fetch"], - config={"setting_one": "test_value"}, + hooks=["prompt_pre_fetch", "tool_pre_invoke", "tool_post_invoke"], + config={}, ) + return NemoCheck(config) + + +@pytest.fixture +def context(): + """Create a PluginContext instance.""" + return PluginContext(global_context=GlobalContext(request_id="1")) + - plugin = NemoCheck(config) +def mock_http_response(status_code, response_data=None): + """Helper to create mock HTTP responses.""" + mock_response = Mock() + mock_response.status_code = status_code + if response_data: + mock_response.json.return_value = response_data + return mock_response - # Test your plugin logic + +@pytest.mark.asyncio +async def test_prompt_pre_fetch(plugin, context): + """Test plugin prompt prefetch hook.""" payload = PromptPrehookPayload( prompt_id="test_prompt", args={"arg0": "This is an argument"} ) - context = GlobalContext(request_id="1") result = await plugin.prompt_pre_fetch(payload, context) assert result.continue_processing + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "status_code,response_data,expected_continue,has_violation,expected_code", + [ + ( + 200, + { + "status": "success", + "rails_status": { + "detect senstitive data": {"status": "success"} + }, + }, + True, + False, + None, + ), + ( + 200, + { + "status": "blocked", + "rails_status": {"detect hap": {"status": "blocked"}}, + }, + False, + True, + "NEMO_RAILS_BLOCKED", + ), + (503, None, False, True, "NEMO_SERVER_ERROR"), + ], +) +async def test_tool_pre_invoke_scenarios( + plugin, + context, + status_code, + response_data, + expected_continue, + has_violation, + expected_code, +): + """Test tool_pre_invoke with various scenarios including error codes.""" + payload = ToolPreInvokePayload( + name="test_tool", + args={"tool_args": '{"param": "value"}'}, + ) + + with patch( + "plugin.requests.post", + return_value=mock_http_response(status_code, response_data), + ): + result = await plugin.tool_pre_invoke(payload, context) + + assert result.continue_processing == expected_continue + assert (result.violation is not None) == has_violation + if has_violation: + assert result.violation.code == expected_code + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "status_code,response_data,expected_continue,has_violation,expected_code", + [ + ( + 200, + { + "status": "success", + "rails_status": { + "detect senstitive data": {"status": "success"} + }, + }, + True, + False, + None, + ), + ( + 200, + { + "status": "blocked", + "rails_status": {"detect hap": {"status": "blocked"}}, + }, + False, + True, + "NEMO_RAILS_BLOCKED", + ), + (500, None, False, True, "NEMO_SERVER_ERROR"), + ], +) +async def test_tool_post_invoke_http_scenarios( + plugin, + context, + status_code, + response_data, + expected_continue, + has_violation, + expected_code, +): + """Test tool_post_invoke with various HTTP response scenarios including error codes.""" + payload = ToolPostInvokePayload( + name="test_tool", + result={"content": [{"type": "text", "text": "Test content"}]}, + ) + + with patch( + "plugin.requests.post", + return_value=mock_http_response(status_code, response_data), + ): + result = await plugin.tool_post_invoke(payload, context) + + assert result.continue_processing == expected_continue + assert (result.violation is not None) == has_violation + if has_violation: + assert result.violation.code == expected_code + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "result_data,should_continue", + [ + ({"content": []}, True), # Empty content + ({"output": "value"}, True), # No content key + ], +) +async def test_tool_post_invoke_passthrough_content_cases( + plugin, context, result_data, should_continue +): + """Test tool_post_invoke no/empty content cases that do not flag.""" + payload = ToolPostInvokePayload(name="test_tool", result=result_data) + result = await plugin.tool_post_invoke(payload, context) + assert result.continue_processing == should_continue + assert result.violation is None + + +@pytest.mark.asyncio +async def test_tool_post_invoke_concatenates_text(plugin, context): + """Test tool_post_invoke concatenates multiple text items.""" + payload = ToolPostInvokePayload( + name="test_tool", + result={ + "content": [ + {"type": "text", "text": "First. "}, + {"type": "text", "text": "Second."}, + ] + }, + ) + + with patch( + "plugin.requests.post", + return_value=mock_http_response( + 200, {"status": "success", "rails_status": {}} + ), + ) as mock_post: + result = await plugin.tool_post_invoke(payload, context) + + assert result.continue_processing + sent_content = mock_post.call_args[1]["json"]["messages"][0]["content"] + assert sent_content == "First. Second." + + +@pytest.mark.asyncio +async def test_tool_post_invoke_filters_non_text(plugin, context): + """Test tool_post_invoke filters non-text content.""" + payload = ToolPostInvokePayload( + name="test_tool", + result={ + "content": [ + {"type": "image", "url": "http://example.com/img.png"}, + {"type": "text", "text": "Text only"}, + ] + }, + ) + + with patch( + "plugin.requests.post", + return_value=mock_http_response( + 200, {"status": "success", "rails_status": {}} + ), + ) as mock_post: + result = await plugin.tool_post_invoke(payload, context) + + assert result.continue_processing + sent_content = mock_post.call_args[1]["json"]["messages"][0]["content"] + assert sent_content == "Text only" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "hook_name,payload_factory", + [ + ( + "tool_pre_invoke", + lambda: ToolPreInvokePayload( + name="test_tool", args={"tool_args": '{"param": "value"}'} + ), + ), + ( + "tool_post_invoke", + lambda: ToolPostInvokePayload( + name="test_tool", + result={"content": [{"type": "text", "text": "content"}]}, + ), + ), + ], +) +async def test_connection_error_handling( + plugin, context, hook_name, payload_factory +): + """Test both hooks fail closed on connection errors with NEMO_CONNECTION_ERROR code.""" + payload = payload_factory() + hook = getattr(plugin, hook_name) + + with patch("plugin.requests.post", side_effect=Exception("Network error")): + result = await hook(payload, context) + + assert not result.continue_processing + assert result.violation is not None + assert result.violation.code == "NEMO_CONNECTION_ERROR" + assert "Network error" in result.violation.description diff --git a/src/__init__.py b/src/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/server.py b/src/server.py index fca8319..e7785e5 100644 --- a/src/server.py +++ b/src/server.py @@ -18,7 +18,6 @@ PromptPrehookPayload, ToolPostInvokePayload, ToolPreInvokePayload, - PluginViolation, ) from mcpgateway.plugins.framework import PluginManager @@ -44,6 +43,56 @@ def set_result_in_body(body, result_args): body["params"]["arguments"] = result_args +def create_mcp_immediate_error_response(body, error_message, violation=None): + """ + Create an MCP error response using immediate_response. + + This helper creates a standardized error response that can be used + for both pre-invoke and post-invoke blocking scenarios. + + Args: + body: The original request/response body containing jsonrpc and id + error_message: Base error message + violation: Optional PluginViolation with reason and description + + Returns: + ProcessingResponse with immediate_response containing the error + """ + # Build error message with violation details if present + if violation is not None: + error_message = f"{violation.reason} -- {violation.description}" + + error_body = { + "jsonrpc": body["jsonrpc"], + "id": body["id"], + "error": {"code": -32000, "message": error_message}, + } + + return ep.ProcessingResponse( + immediate_response=ep.ImmediateResponse( + # Use 200 status with error in body for MCP protocol compatibility + status=http_status_pb2.HttpStatus(code=200), + headers=ep.HeaderMutation( + set_headers=[ + core.HeaderValueOption( + header=core.HeaderValue( + key="content-type", + raw_value="application/json".encode("utf-8"), + ) + ), + core.HeaderValueOption( + header=core.HeaderValue( + key="x-mcp-denied", + raw_value="True".encode("utf-8"), + ) + ), + ], + ), + body=(json.dumps(error_body)).encode("utf-8"), + ) + ) + + # ============================================================================ # MCP HOOK HANDLERS # ============================================================================ @@ -62,7 +111,9 @@ async def getToolPreInvokeResponse(body): "tool_args": body["params"]["arguments"], "client_session_id": "replaceme", } - payload = ToolPreInvokePayload(name=body["params"]["name"], args=payload_args) + payload = ToolPreInvokePayload( + name=body["params"]["name"], args=payload_args + ) # TODO: hard-coded ids global_context = GlobalContext(request_id="1", server_id="2") logger.debug(f"**** Invoking Tool Pre Invoke with payload: {payload} ****") @@ -71,36 +122,10 @@ async def getToolPreInvokeResponse(body): ) logger.debug(f"**** Tool Pre Invoke Result: {result} ****") if not result.continue_processing: - error_message = "No go - Tool args forbidden" - if result.violation is not None: - violation: PluginViolation = result.violation - error_message = f"{violation.reason} -- {violation.description}" - error_body = { - "jsonrpc": body["jsonrpc"], - "id": body["id"], - "error": {"code": -32000, "message": error_message}, - } - body_resp = ep.ProcessingResponse( - immediate_response=ep.ImmediateResponse( - # ok for stream, with error in body - status=http_status_pb2.HttpStatus(code=200), - headers=ep.HeaderMutation( - set_headers=[ - core.HeaderValueOption( - header=core.HeaderValue( - key="content-type", - raw_value="application/json".encode("utf-8"), - ) - ), - core.HeaderValueOption( - header=core.HeaderValue( - key="x-mcp-denied", raw_value="True".encode("utf-8") - ) - ), - ], - ), - body=(json.dumps(error_body)).encode("utf-8"), - ) + body_resp = create_mcp_immediate_error_response( + body, + error_message="No go - Tool args forbidden", + violation=result.violation, ) else: logger.debug("continue_processing true") @@ -109,7 +134,9 @@ async def getToolPreInvokeResponse(body): body["params"]["arguments"] = result_payload.args["tool_args"] body_mutation = ep.BodyResponse( response=ep.CommonResponse( - body_mutation=ep.BodyMutation(body=json.dumps(body).encode("utf-8")) + body_mutation=ep.BodyMutation( + body=json.dumps(body).encode("utf-8") + ) ) ) else: @@ -126,6 +153,10 @@ async def getToolPostInvokeResponse(body): Invokes plugins after a tool has been called, allowing for result validation, modification, or filtering of the tool output. + + Note: In STREAMED mode, blocking responses may fail if headers are already sent. + This implementation uses immediate_response to attempt early termination, but + it may not always succeed due to streaming constraints. """ # FIXME: size of content array is expected to be 1 # for content in body["result"]["content"]: @@ -138,27 +169,33 @@ async def getToolPostInvokeResponse(body): result, _ = await manager.invoke_hook( ToolHookType.TOOL_POST_INVOKE, payload, global_context=global_context ) - logger.info(result) + logger.debug(f"**** Tool Post Invoke result {result}") if not result.continue_processing: - body_resp = ep.ProcessingResponse( - immediate_response=ep.ImmediateResponse( - # TODO: hard-coded error reason - status=http_status_pb2.HttpStatus(code=http_status_pb2.Forbidden), - details="No go", - ) + # In STREAMED mode, we attempt to use immediate_response to terminate early + # This may fail if response headers have already been sent + body_resp = create_mcp_immediate_error_response( + body, + error_message="Tool response forbidden", + violation=result.violation, ) - else: - result_payload = result.modified_payload - if result_payload is not None: - body["result"] = result_payload.result - body_mutation = ep.BodyResponse( - response=ep.CommonResponse( - body_mutation=ep.BodyMutation(body=json.dumps(body).encode("utf-8")) + logger.info(f"****Tool Post Invoke Return body: {body_resp}****") + return body_resp + + # Continue processing - allow or modify the response + result_payload = result.modified_payload + if result_payload is not None: + body["result"] = result_payload.result + body_mutation = ep.BodyResponse( + response=ep.CommonResponse( + body_mutation=ep.BodyMutation( + body=json.dumps(body).encode("utf-8") ) ) - else: - body_mutation = ep.BodyResponse(response=ep.CommonResponse()) - body_resp = ep.ProcessingResponse(request_body=body_mutation) + ) + else: + body_mutation = ep.BodyResponse(response=ep.CommonResponse()) + body_resp = ep.ProcessingResponse(response_body=body_mutation) + logger.info(f"****Tool Post Invoke Return body: {body_resp}****") return body_resp @@ -179,18 +216,19 @@ async def getPromptPreFetchResponse(body): ) logger.info(result) if not result.continue_processing: - body_resp = ep.ProcessingResponse( - immediate_response=ep.ImmediateResponse( - status=http_status_pb2.HttpStatus(code=http_status_pb2.Forbidden), - details="No go", - ) + body_resp = create_mcp_immediate_error_response( + body, + error_message="Tool response forbidden", + violation=result.violation, ) else: body["params"]["arguments"] = result.modified_payload.args body_resp = ep.ProcessingResponse( request_body=ep.BodyResponse( response=ep.CommonResponse( - body_mutation=ep.BodyMutation(body=json.dumps(body).encode("utf-8")) + body_mutation=ep.BodyMutation( + body=json.dumps(body).encode("utf-8") + ) ) ) ) @@ -198,6 +236,85 @@ async def getPromptPreFetchResponse(body): return body_resp +# ============================================================================ +# RESPONSE BODY PROCESSING HELPER +# ============================================================================ + + +async def process_response_body_buffer(buffer: bytearray): + """Process buffered response body content. + + Parses the buffered content (supporting both SSE and plain JSON-RPC formats), + and invokes the tool post-invoke hook if it's a tool result. + + Args: + buffer: The accumulated response body bytes + + Returns: + ProcessingResponse to send back to Envoy + """ + if not buffer: + # Empty buffer at end of stream + logger.debug("End of stream with empty buffer") + return ep.ProcessingResponse( + response_body=ep.BodyResponse(response=ep.CommonResponse()) + ) + + try: + text = buffer.decode("utf-8") + except UnicodeDecodeError: + logger.debug("Response body not UTF-8; skipping") + return ep.ProcessingResponse( + response_body=ep.BodyResponse(response=ep.CommonResponse()) + ) + + lines = text.split("\n") + logger.debug(f"Response body text: {lines}") + + # Handle both SSE format and plain JSON-RPC format + data = None + + # Check if this is SSE format (starts with "event:" or "data:") + if text.strip().startswith(("event:", "data:")): + # Parse SSE format + lines = text.split("\n") + for line in lines: + line = line.strip() + if line.startswith("data:"): + json_str = line[5:].strip() # Remove "data:" prefix + logger.debug(f"Extracted JSON from SSE: {json_str}") + try: + data = json.loads(json_str) + break + except json.JSONDecodeError: + continue + else: + # Parse plain JSON-RPC format + lines = [line.strip() for line in text.split("\n") if line.strip()] + if lines: + try: + data = json.loads(lines[0]) + except json.JSONDecodeError as e: + logger.error(f"Failed to parse JSON: {e}") + + if data: + logger.debug(f"Parsed response data: {data}") + + # Check if this is a tool result response + if "result" in data and "content" in data["result"]: + logger.info("Invoking tool post-invoke hook") + return await getToolPostInvokeResponse(data) + else: + return ep.ProcessingResponse( + response_body=ep.BodyResponse(response=ep.CommonResponse()) + ) + else: + logger.warning("No data parsed from response body") + return ep.ProcessingResponse( + response_body=ep.BodyResponse(response=ep.CommonResponse()) + ) + + # ============================================================================ # ENVOY EXTERNAL PROCESSOR SERVICER # ============================================================================ @@ -293,7 +410,9 @@ async def Process( body = json.loads(text) if "method" in body and body["method"] == "tools/call": body_resp = await getToolPreInvokeResponse(body) - elif "method" in body and body["method"] == "prompts/get": + elif ( + "method" in body and body["method"] == "prompts/get" + ): body_resp = await getPromptPreFetchResponse(body) else: body_resp = ep.ProcessingResponse( @@ -308,43 +427,37 @@ async def Process( # ---------------------------------------------------------------- # Response Body Processing (MCP Tool Results) # ---------------------------------------------------------------- - elif request.HasField("response_body") and request.response_body.body: - chunk = request.response_body.body - resp_body_buf.extend(chunk) + elif request.HasField("response_body"): + logger.debug(f"Processing response body: {request}") + # Buffer content if present in this chunk + if request.response_body.body: + chunk = request.response_body.body + resp_body_buf.extend(chunk) + logger.debug(f"Buffered chunk ({len(chunk)} bytes)") + + # Check for end of stream (regardless of whether this chunk has content) if getattr(request.response_body, "end_of_stream", False): - try: - text = resp_body_buf.decode("utf-8") - except UnicodeDecodeError: - logger.debug("Response body not UTF-8; skipping") - else: - logger.info(text.split("\n")) - # find data key - data = [d for d in text.split("\n") if d.startswith("data:")] - # logger.info(json.loads(data[0].strip("data:"))) - if data: # List can be empty - data = json.loads(data[0].strip("data:")) - # TODO: check for tool call - if "result" in data and "content" in data["result"]: - body_resp = await getToolPostInvokeResponse(data) - else: - body_resp = ep.ProcessingResponse( - response_body=ep.BodyResponse( - response=ep.CommonResponse() - ) - ) - yield body_resp + logger.debug( + "End of stream reached, processing complete buffered response" + ) + + # Process the buffered content + body_resp = await process_response_body_buffer( + resp_body_buf + ) + yield body_resp resp_body_buf.clear() - # ---------------------------------------------------------------- - # Response Body Processing (No body field) - # ---------------------------------------------------------------- - elif request.HasField("response_body"): - logger.warning("On Response, no body.") - logger.warning(request) - body_resp = ep.ProcessingResponse( - response_body=ep.BodyResponse(response=ep.CommonResponse()) - ) - yield body_resp + else: + # Intermediate chunk - acknowledge but don't process yet + logger.debug( + "Buffering intermediate chunk, waiting for end_of_stream" + ) + yield ep.ProcessingResponse( + response_body=ep.BodyResponse( + response=ep.CommonResponse() + ) + ) else: # Unhandled request types diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/pytest.ini b/tests/pytest.ini new file mode 100644 index 0000000..9a43301 --- /dev/null +++ b/tests/pytest.ini @@ -0,0 +1,11 @@ +[pytest] +log_cli = false +log_cli_level = INFO +log_cli_format = %(asctime)s [%(module)s] [%(levelname)s] %(message)s +log_cli_date_format = %Y-%m-%d %H:%M:%S +log_level = INFO +log_format = %(asctime)s [%(module)s] [%(levelname)s] %(message)s +log_date_format = %Y-%m-%d %H:%M:%S +pythonpath = . src +filterwarnings = + ignore::DeprecationWarning \ No newline at end of file diff --git a/tests/test_server.py b/tests/test_server.py new file mode 100644 index 0000000..ab425bb --- /dev/null +++ b/tests/test_server.py @@ -0,0 +1,458 @@ +"""Unit tests for ext-proc server functions + +These tests use dynamic import and mocking to avoid proto dependencies. +""" + +# Standard +from unittest.mock import AsyncMock, Mock, MagicMock +import sys +import json + +# Third-Party +import pytest + +# First-Party +from mcpgateway.plugins.framework import ( + ToolPostInvokeResult, + ToolPostInvokePayload, + PluginViolation, +) + + +@pytest.fixture +def mock_envoy_modules(): + """Mock envoy protobuf modules to avoid proto dependencies.""" + # Create mock modules + mock_ep = MagicMock() + mock_ep_grpc = MagicMock() + mock_core = MagicMock() + mock_http_status = MagicMock() + + # Add to sys.modules before importing server + sys.modules["envoy"] = MagicMock() + sys.modules["envoy.service"] = MagicMock() + sys.modules["envoy.service.ext_proc"] = MagicMock() + sys.modules["envoy.service.ext_proc.v3"] = MagicMock() + sys.modules["envoy.service.ext_proc.v3.external_processor_pb2"] = mock_ep + sys.modules["envoy.service.ext_proc.v3.external_processor_pb2_grpc"] = ( + mock_ep_grpc + ) + sys.modules["envoy.config"] = MagicMock() + sys.modules["envoy.config.core"] = MagicMock() + sys.modules["envoy.config.core.v3"] = MagicMock() + sys.modules["envoy.config.core.v3.base_pb2"] = mock_core + sys.modules["envoy.type"] = MagicMock() + sys.modules["envoy.type.v3"] = MagicMock() + sys.modules["envoy.type.v3.http_status_pb2"] = mock_http_status + + yield { + "ep": mock_ep, + "ep_grpc": mock_ep_grpc, + "core": mock_core, + "http_status": mock_http_status, + } + + # Cleanup + for key in list(sys.modules.keys()): + if key.startswith("envoy"): + del sys.modules[key] + if "src.server" in sys.modules: + del sys.modules["src.server"] + + +@pytest.fixture +def mock_manager(): + """Create a mock PluginManager.""" + mock = Mock() + mock.invoke_hook = AsyncMock() + return mock + + +@pytest.fixture +def sample_tool_result_body(): + """Create a sample tool result body.""" + return { + "jsonrpc": "2.0", + "id": "test-123", + "result": { + "content": [{"type": "text", "text": "Tool execution result"}] + }, + } + + +def setup_response_mocks(mock_envoy_modules): + """Setup common response mocks.""" + mock_envoy_modules["ep"].ProcessingResponse.return_value = MagicMock() + mock_envoy_modules["ep"].BodyResponse.return_value = MagicMock() + mock_envoy_modules["ep"].CommonResponse.return_value = MagicMock() + + +def setup_manager_with_result(mock_manager, continue_processing=True): + """Setup mock manager with a tool post-invoke result.""" + mock_result = ToolPostInvokeResult(continue_processing=continue_processing) + mock_manager.invoke_hook.return_value = (mock_result, None) + return mock_manager + + +def verify_payload_content(payload, expected_result, expected_text): + """Verify payload contains expected content.""" + assert isinstance(payload, ToolPostInvokePayload) + assert payload.result == expected_result + assert payload.result["content"][0]["type"] == "text" + assert payload.result["content"][0]["text"] == expected_text + + +# ============================================================================ +# Tool Post-Invoke Hook Tests +# ============================================================================ + + +@pytest.mark.asyncio +async def test_getToolPostInvokeResponse_continue_processing( + mock_envoy_modules, mock_manager, sample_tool_result_body +): + """Test getToolPostInvokeResponse when plugin allows processing to continue.""" + # Setup mock response objects + mock_response = MagicMock() + mock_response.HasField.return_value = True + mock_response.response_body.response.HasField.return_value = False + mock_envoy_modules["ep"].ProcessingResponse.return_value = mock_response + + # Import server after mocking + import src.server + + # Setup mock to return continue_processing=True + mock_result = ToolPostInvokeResult( + continue_processing=True, + modified_payload=None, + ) + mock_manager.invoke_hook.return_value = (mock_result, None) + + # Inject mock manager + src.server.manager = mock_manager + + # Call the function + _ = await src.server.getToolPostInvokeResponse(sample_tool_result_body) + + # Verify the hook was called + assert mock_manager.invoke_hook.called + call_args = mock_manager.invoke_hook.call_args[0] + payload = call_args[1] + assert isinstance(payload, ToolPostInvokePayload) + assert payload.result == sample_tool_result_body["result"] + # assert payload.name == "replaceme" # Replace this after better naming + + +@pytest.mark.asyncio +async def test_getToolPostInvokeResponse_blocked( + mock_envoy_modules, mock_manager, sample_tool_result_body +): + """Test getToolPostInvokeResponse when plugin blocks the response. + + This test verifies that when continue_processing=False, the function + uses immediate_response (not response_body) and includes violation details. + """ + # Setup mocks for immediate_response path + setup_response_mocks(mock_envoy_modules) + + # Import server after mocking + import src.server + + # Setup mock to return continue_processing=False with violation + violation = PluginViolation( + reason="Sensitive content detected", + description="Tool response contains forbidden content", + code="CONTENT_VIOLATION", + ) + mock_result = ToolPostInvokeResult( + continue_processing=False, + violation=violation, + ) + mock_manager.invoke_hook.return_value = (mock_result, None) + + # Inject mock manager + src.server.manager = mock_manager + + # Capture json.dumps calls to verify error body content + original_dumps = json.dumps + captured_bodies = [] + + def spy_dumps(obj, **kwargs): + if isinstance(obj, dict) and "error" in obj: + captured_bodies.append(obj) + return original_dumps(obj, **kwargs) + + json.dumps = spy_dumps + try: + # Call the function + response = await src.server.getToolPostInvokeResponse( + sample_tool_result_body + ) + finally: + json.dumps = original_dumps + + # Verify the hook was called with correct payload + assert mock_manager.invoke_hook.called + call_args = mock_manager.invoke_hook.call_args[0] + payload = call_args[1] + assert isinstance(payload, ToolPostInvokePayload) + assert payload.result == sample_tool_result_body["result"] + + # Verify response was created (error path taken) + assert response is not None + + # Verify error body was created with violation details + assert len(captured_bodies) > 0 + error_body = captured_bodies[0] + assert "error" in error_body + assert error_body["error"]["code"] == -32000 + # Verify violation message is included + assert "Sensitive content detected" in error_body["error"]["message"] + assert ( + "Tool response contains forbidden content" + in error_body["error"]["message"] + ) + + +@pytest.mark.asyncio +async def test_getToolPostInvokeResponse_modified_payload( + mock_envoy_modules, mock_manager, sample_tool_result_body +): + """Test getToolPostInvokeResponse when plugin modifies the payload.""" + # Import server after mocking + import src.server + + # Setup mock to return modified payload + modified_result = { + "content": [{"type": "text", "text": "Modified tool result"}] + } + modified_payload = ToolPostInvokePayload( + name="test_tool", result=modified_result + ) + mock_result = ToolPostInvokeResult( + continue_processing=True, + modified_payload=modified_payload, + ) + mock_manager.invoke_hook.return_value = (mock_result, None) + + # Inject mock manager + src.server.manager = mock_manager + + # Spy on json.dumps to capture what body is being serialized + original_dumps = json.dumps + captured_body = None + + def spy_dumps(obj, **kwargs): + nonlocal captured_body + # Capture the body dict that's being serialized + if isinstance(obj, dict) and "result" in obj and "jsonrpc" in obj: + captured_body = obj + return original_dumps(obj, **kwargs) + + json.dumps = spy_dumps + try: + # Call the function + response = await src.server.getToolPostInvokeResponse( + sample_tool_result_body + ) + finally: + json.dumps = original_dumps + + # Verify the hook was called + assert mock_manager.invoke_hook.called + + # Verify response was created + assert response is not None + + # Verify the body was modified with the new result + assert captured_body is not None, ( + "json.dumps should have been called with the modified body" + ) + assert captured_body["result"] == modified_result + assert ( + captured_body["result"]["content"][0]["text"] == "Modified tool result" + ) + # Verify original metadata (jsonrpc, id) is preserved + assert captured_body["jsonrpc"] == sample_tool_result_body["jsonrpc"] + assert captured_body["id"] == sample_tool_result_body["id"] + + +@pytest.mark.asyncio +async def test_getToolPostInvokeResponse_multiple_content_items( + mock_envoy_modules, mock_manager +): + """Test getToolPostInvokeResponse with multiple content items.""" + # Setup mock response + mock_response = MagicMock() + mock_envoy_modules["ep"].ProcessingResponse.return_value = mock_response + + # Import server after mocking + import src.server + + body = { + "jsonrpc": "2.0", + "id": "test-789", + "result": { + "content": [ + {"type": "text", "text": "First item"}, + {"type": "text", "text": "Second item"}, + {"type": "image", "url": "http://example.com/img.png"}, + ] + }, + } + + mock_result = ToolPostInvokeResult(continue_processing=True) + mock_manager.invoke_hook.return_value = (mock_result, None) + + # Inject mock manager + src.server.manager = mock_manager + + # Call the function + _ = await src.server.getToolPostInvokeResponse(body) + + # Verify the payload passed to the hook contains all content + call_args = mock_manager.invoke_hook.call_args[0] + payload = call_args[1] + assert len(payload.result["content"]) == 3 + assert payload.result["content"][0]["text"] == "First item" + assert payload.result["content"][1]["text"] == "Second item" + assert payload.result["content"][2]["url"] == "http://example.com/img.png" + + +# ============================================================================ +# Response Body Processing Tests +# ============================================================================ + + +@pytest.mark.asyncio +async def test_process_response_body_buffer_with_tool_result( + mock_envoy_modules, mock_manager +): + """Test process_response_body_buffer with a tool result.""" + setup_response_mocks(mock_envoy_modules) + import src.server + + setup_manager_with_result(mock_manager) + src.server.manager = mock_manager + + tool_result = { + "jsonrpc": "2.0", + "id": "test-123", + "result": {"content": [{"type": "text", "text": "Result"}]}, + } + buffer = bytearray(json.dumps(tool_result).encode("utf-8")) + response = await src.server.process_response_body_buffer(buffer) + + assert mock_manager.invoke_hook.called + payload = mock_manager.invoke_hook.call_args[0][1] + verify_payload_content(payload, tool_result["result"], "Result") + # Verify ProcessingResponse was returned + assert response is not None + + +@pytest.mark.asyncio +async def test_process_response_body_buffer_with_sse_format( + mock_envoy_modules, mock_manager +): + """Test process_response_body_buffer with SSE formatted content.""" + setup_response_mocks(mock_envoy_modules) + import src.server + + setup_manager_with_result(mock_manager) + src.server.manager = mock_manager + + tool_result = { + "jsonrpc": "2.0", + "id": "test-sse", + "result": {"content": [{"type": "text", "text": "SSE data"}]}, + } + sse_body = f"event: message\ndata: {json.dumps(tool_result)}\n\n" + buffer = bytearray(sse_body.encode("utf-8")) + response = await src.server.process_response_body_buffer(buffer) + + assert mock_manager.invoke_hook.called + payload = mock_manager.invoke_hook.call_args[0][1] + verify_payload_content(payload, tool_result["result"], "SSE data") + # Verify ProcessingResponse was returned + assert response is not None + + +@pytest.mark.asyncio +async def test_process_response_body_buffer_multiple_chunks_scenario( + mock_envoy_modules, mock_manager +): + """Test buffering: content in chunks, then empty end_of_stream chunk. + + Simulates: chunk1 (content) + chunk2 (content) + chunk3 (empty, end_of_stream). + """ + setup_response_mocks(mock_envoy_modules) + import src.server + + setup_manager_with_result(mock_manager) + src.server.manager = mock_manager + + tool_result = { + "jsonrpc": "2.0", + "id": "test-multi-chunk", + "result": {"content": [{"type": "text", "text": "Multi chunk data"}]}, + } + body_bytes = json.dumps(tool_result).encode("utf-8") + + # Simulate buffering: chunk1 + chunk2 + empty chunk + buffer = bytearray() + buffer.extend(body_bytes[:25]) # Chunk 1 + buffer.extend(body_bytes[25:]) # Chunk 2 + buffer.extend(b"") # Chunk 3 (empty, triggers processing) + + response = await src.server.process_response_body_buffer(buffer) + + assert mock_manager.invoke_hook.called + payload = mock_manager.invoke_hook.call_args[0][1] + verify_payload_content(payload, tool_result["result"], "Multi chunk data") + # Verify ProcessingResponse was returned + assert response is not None + + +@pytest.mark.asyncio +async def test_process_response_body_buffer_empty( + mock_envoy_modules, mock_manager +): + """Test process_response_body_buffer with empty buffer.""" + setup_response_mocks(mock_envoy_modules) + import src.server + + src.server.manager = mock_manager + response = await src.server.process_response_body_buffer(bytearray()) + + # Verify hook is NOT called for empty buffer + assert not mock_manager.invoke_hook.called, ( + "Tool post-invoke hook should not be called for empty buffer" + ) + # Verify response is returned (function doesn't crash on empty buffer) + assert response is not None + + +@pytest.mark.asyncio +async def test_process_response_body_buffer_non_tool_result( + mock_envoy_modules, mock_manager +): + """Test process_response_body_buffer with non-tool result (error response).""" + setup_response_mocks(mock_envoy_modules) + import src.server + + src.server.manager = mock_manager + + error_response = { + "jsonrpc": "2.0", + "id": "test-error", + "error": {"code": -32000, "message": "Error"}, + } + buffer = bytearray(json.dumps(error_response).encode("utf-8")) + response = await src.server.process_response_body_buffer(buffer) + + # Verify hook is NOT called for error responses + assert not mock_manager.invoke_hook.called, ( + "Tool post-invoke hook should not be called for error responses" + ) + # Verify response is returned (function handles error responses gracefully) + assert response is not None