fix(nemo-relay): preserve managed interceptor outcomes

Signed-off-by: Bryan Bednarski <bbednarski@nvidia.com>

plugins/observability/nemo_relay/__init__.py

@@ -20,6 +20,11 @@ logger = logging.getLogger(__name__)
 _INIT_FAILED = object()
 _LOCK = threading.RLock()
 _RUNTIME: "_Runtime | object | None" = None
+_RELAY_LLM_SURFACE_BY_API_MODE = {
+    "anthropic_messages": "anthropic.messages",
+    "chat_completions": "openai.chat_completions",
+    "codex_responses": "openai.responses",
+}
 
 
 @dataclass
@@ -358,6 +363,8 @@ class _Runtime:
         normalize_payload: Callable[[Any], Any],
         shape_response: Callable[[Any], Any],
         make_managed_execute: Callable[[Callable[[Any], Any]], Any],
+        *,
+        preserve_raw_response: bool,
     ) -> Any:
         # NeMo Relay's native managed execution may wrap a failing callback as an
         # internal runtime error, hiding the real downstream provider/tool
@@ -387,7 +394,9 @@ class _Runtime:
             if downstream_error is not None and _is_relay_wrapped_callback_error(exc, callback_error):
                 raise downstream_error
             raise
-        return raw_response["value"] if raw_response["set"] else managed_result
+        if preserve_raw_response and raw_response["set"]:
+            return raw_response["value"]
+        return managed_result
 
     def execute_llm(self, kwargs: dict[str, Any]) -> Any:
         state = self.ensure_session(kwargs)
@@ -404,7 +413,7 @@ class _Runtime:
         def _make_managed(impl: Callable[[Any], Any]) -> Any:
             async def _managed_execute() -> Any:
                 result = self.nemo_relay.llm.execute(
-                    str(kwargs.get("provider") or "llm"),
+                    _relay_llm_surface(kwargs),
                     request,
                     impl,
                     handle=state.handle,
@@ -426,7 +435,7 @@ class _Runtime:
             return _managed_execute()
 
         return self._run_managed_with_downstream_preservation(
-            next_call, _normalize, _llm_response_payload, _make_managed
+            next_call, _normalize, _llm_response_payload, _make_managed, preserve_raw_response=True
         )
 
     def execute_tool(self, kwargs: dict[str, Any]) -> Any:
@@ -464,7 +473,7 @@ class _Runtime:
             return _managed_execute()
 
         return self._run_managed_with_downstream_preservation(
-            next_call, _normalize, _jsonable, _make_managed
+            next_call, _normalize, _jsonable, _make_managed, preserve_raw_response=False
         )
 
 
@@ -923,6 +932,14 @@ def _tool_key(kwargs: dict[str, Any]) -> str:
     )
 
 
+def _relay_llm_surface(kwargs: dict[str, Any]) -> str:
+    api_mode = str(kwargs.get("api_mode") or "").strip().lower()
+    return _RELAY_LLM_SURFACE_BY_API_MODE.get(
+        api_mode,
+        str(kwargs.get("provider") or "llm"),
+    )
+
+
 def _metadata(kwargs: dict[str, Any]) -> dict[str, Any]:
     keys = (
         "telemetry_schema_version",

tests/plugins/test_nemo_relay_plugin.py

@@ -603,6 +603,92 @@ def test_nemo_relay_plugin_activates_and_owns_dynamic_plugins(tmp_path, monkeypa
     assert event_names.index("atif.deregister") < event_names.index("plugin.activation.close")
 
 
+@pytest.mark.parametrize(
+    ("provider", "api_mode", "expected_surface", "should_rewrite"),
+    [
+        ("custom", "chat_completions", "openai.chat_completions", True),
+        ("openai-codex", "codex_responses", "openai.responses", True),
+        ("anthropic", "anthropic_messages", "anthropic.messages", True),
+        ("custom", "anthropic_messages", "anthropic.messages", True),
+        ("bedrock", "bedrock_converse", "bedrock", False),
+    ],
+)
+def test_nemo_relay_managed_llm_uses_wire_protocol_for_interceptor_dispatch(
+    tmp_path,
+    monkeypatch,
+    provider,
+    api_mode,
+    expected_surface,
+    should_rewrite,
+):
+    fake = _FakeNemoRelay()
+    supported_surfaces = {
+        "anthropic.messages",
+        "openai.chat_completions",
+        "openai.responses",
+    }
+
+    def execute(name, request, func, **kwargs):
+        fake.events.append(("llm.execute.start", name, request.content, kwargs))
+        content = dict(request.content)
+        if name in supported_surfaces:
+            content["rewritten_for"] = name
+        result = func(_FakeLLMRequest(request.headers, content))
+        fake.events.append(("llm.execute.end", name, result, kwargs))
+        return result
+
+    fake.llm.execute = execute
+    plugin = _fresh_plugin(monkeypatch, fake)
+    _enable_dynamic_plugin(tmp_path, monkeypatch)
+
+    result = plugin.on_llm_execution_middleware(
+        session_id="s1",
+        provider=provider,
+        api_mode=api_mode,
+        model="fixture",
+        request={"messages": [{"role": "user", "content": "hi"}]},
+        next_call=lambda request: request,
+    )
+
+    execute_start = next(
+        event for event in fake.events if event[0] == "llm.execute.start"
+    )
+    assert execute_start[1] == expected_surface
+    assert execute_start[3]["metadata"]["provider"] == provider
+    assert execute_start[3]["metadata"]["api_mode"] == api_mode
+    if should_rewrite:
+        assert result["rewritten_for"] == expected_surface
+    else:
+        assert "rewritten_for" not in result
+
+
+def test_nemo_relay_managed_tool_returns_post_interceptor_result(tmp_path, monkeypatch):
+    fake = _FakeNemoRelay()
+
+    def execute(name, args, func, **kwargs):
+        fake.events.append(("tool.execute.start", name, args, kwargs))
+        raw = func({"intercepted": True, **args})
+        result = {"compressed": True, "raw": raw}
+        fake.events.append(("tool.execute.end", name, result, kwargs))
+        return result
+
+    fake.tools.execute = execute
+    plugin = _fresh_plugin(monkeypatch, fake)
+    _enable_dynamic_plugin(tmp_path, monkeypatch)
+
+    result = plugin.on_tool_execution_middleware(
+        session_id="s1",
+        tool_name="fixture-tool",
+        args={"value": 1},
+        next_call=lambda args: {"tool_output": args},
+    )
+
+    assert result == {
+        "compressed": True,
+        "raw": {"tool_output": {"intercepted": True, "value": 1}},
+    }
+
+
 def test_nemo_relay_plugin_activates_before_registering_managed_middleware(tmp_path, monkeypatch):
     fake = _FakeNemoRelay()
     plugin = _fresh_plugin(monkeypatch, fake)