fix(nemo-relay): preserve downstream errors in adaptive execution (#42691)

Based on #42658 by @mnajafian-nv.

Preserves the real downstream provider/tool exception when NeMo Relay's
managed adaptive execution wraps a failing callback as an internal runtime
error. Without this, the original exception (and its retry-classification
signal, e.g. status_code) is lost behind Relay's wrapper.

Salvage changes on top of the original PR:

- Tolerant Relay-wrapper match: _is_relay_wrapped_callback_error now uses
  str.startswith on the "internal error: <cls>: <msg>" prefix instead of
  exact equality, so a future Relay version appending a traceback/suffix
  doesn't silently defeat the unwrap. On a total format change it returns
  False and falls back to the pre-fix behavior (surfacing Relay's error)
  rather than masking it.
- Deduplicated the LLM and tool execute paths into a shared
  _run_managed_with_downstream_preservation helper, removing ~20 lines of
  copy-pasted nonlocal/try-except scaffolding that could drift out of sync.
- Added a real-middleware regression guard
  (test_nemo_relay_downstream_unwrap_matches_real_middleware_wrapper_shape)
  that drives hermes_cli.middleware._run_execution_chain and asserts the
  plugin's _original_downstream_error unwraps the actual private
  _DownstreamExecutionError wrapper. The original synthetic tests modeled the
  wrapper with a local class, so a rename or shape change in core middleware
  would not have been caught; this test fails loudly if that contract drifts.

Co-authored-by: mnajafian-nv <mnajafian@nvidia.com>

plugins/observability/nemo_relay/__init__.py

@@ -9,6 +9,7 @@ import logging
 import os
 import threading
 import tomllib
+from collections.abc import Callable
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any, Optional
@@ -284,6 +285,43 @@ class _Runtime:
             and callable(getattr(getattr(self.nemo_relay, "tools", None), "execute", None))
         )
 
+    def _run_managed_with_downstream_preservation(
+        self,
+        next_call: Callable[[Any], Any],
+        normalize_payload: Callable[[Any], Any],
+        shape_response: Callable[[Any], Any],
+        make_managed_execute: Callable[[Callable[[Any], Any]], Any],
+    ) -> Any:
+        # NeMo Relay's native managed execution may wrap a failing callback as an
+        # internal runtime error, hiding the real downstream provider/tool
+        # exception. Capture the original here and re-raise it after managed
+        # execution so Hermes retry classification still sees it. The LLM and tool
+        # paths share this scaffolding; they differ only in payload normalization,
+        # response shaping, and the Relay call itself.
+        raw_response: dict[str, Any] = {"set": False, "value": None}
+        callback_error: Exception | None = None
+        downstream_error: BaseException | None = None
+
+        def _impl(next_payload: Any) -> Any:
+            nonlocal callback_error, downstream_error
+            try:
+                raw = next_call(normalize_payload(next_payload))
+            except Exception as exc:
+                callback_error = exc
+                downstream_error = _original_downstream_error(exc)
+                raise
+            raw_response["set"] = True
+            raw_response["value"] = raw
+            return shape_response(raw)
+
+        try:
+            managed_result = _resolve_awaitable(make_managed_execute(_impl))
+        except Exception as exc:
+            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
+
     def execute_llm(self, kwargs: dict[str, Any]) -> Any:
         state = self.ensure_session(kwargs)
         request_body = _jsonable(kwargs.get("request") or {})
@@ -292,38 +330,37 @@ class _Runtime:
         if not callable(next_call):
             return request_body
 
-        raw_response: dict[str, Any] = {"set": False, "value": None}
-
-        def _impl(next_request: Any) -> Any:
+        def _normalize(next_request: Any) -> Any:
             next_body = getattr(next_request, "content", next_request)
-            raw = next_call(next_body if isinstance(next_body, dict) else request_body)
-            raw_response["set"] = True
-            raw_response["value"] = raw
-            return _llm_response_payload(raw)
+            return next_body if isinstance(next_body, dict) else request_body
 
-        async def _managed_execute() -> Any:
-            result = self.nemo_relay.llm.execute(
-                str(kwargs.get("provider") or "llm"),
-                request,
-                _impl,
-                handle=state.handle,
-                data=_jsonable(
-                    {
-                        "turn_id": kwargs.get("turn_id"),
-                        "api_request_id": kwargs.get("api_request_id"),
-                        "api_call_count": kwargs.get("api_call_count"),
-                        "mode": self.settings.adaptive_mode,
-                    }
-                ),
-                metadata=_metadata(kwargs),
-                model_name=str(kwargs.get("model") or ""),
-            )
-            if inspect.isawaitable(result):
-                return await result
-            return result
+        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"),
+                    request,
+                    impl,
+                    handle=state.handle,
+                    data=_jsonable(
+                        {
+                            "turn_id": kwargs.get("turn_id"),
+                            "api_request_id": kwargs.get("api_request_id"),
+                            "api_call_count": kwargs.get("api_call_count"),
+                            "mode": self.settings.adaptive_mode,
+                        }
+                    ),
+                    metadata=_metadata(kwargs),
+                    model_name=str(kwargs.get("model") or ""),
+                )
+                if inspect.isawaitable(result):
+                    return await result
+                return result
 
-        managed_result = _resolve_awaitable(_managed_execute())
-        return raw_response["value"] if raw_response["set"] else managed_result
+            return _managed_execute()
+
+        return self._run_managed_with_downstream_preservation(
+            next_call, _normalize, _llm_response_payload, _make_managed
+        )
 
     def execute_tool(self, kwargs: dict[str, Any]) -> Any:
         state = self.ensure_session(kwargs)
@@ -333,37 +370,35 @@ class _Runtime:
         if not callable(next_call):
             return args
 
-        raw_response: dict[str, Any] = {"set": False, "value": None}
+        def _normalize(next_args: Any) -> Any:
+            return next_args if isinstance(next_args, dict) else args
 
-        def _impl(next_args: Any) -> Any:
-            effective_args = next_args if isinstance(next_args, dict) else args
-            raw = next_call(effective_args)
-            raw_response["set"] = True
-            raw_response["value"] = raw
-            return _jsonable(raw)
+        def _make_managed(impl: Callable[[Any], Any]) -> Any:
+            async def _managed_execute() -> Any:
+                result = self.nemo_relay.tools.execute(
+                    tool_name,
+                    args,
+                    impl,
+                    handle=state.handle,
+                    data=_jsonable(
+                        {
+                            "turn_id": kwargs.get("turn_id"),
+                            "api_request_id": kwargs.get("api_request_id"),
+                            "tool_call_id": kwargs.get("tool_call_id"),
+                            "mode": self.settings.adaptive_mode,
+                        }
+                    ),
+                    metadata=_metadata(kwargs),
+                )
+                if inspect.isawaitable(result):
+                    return await result
+                return result
 
-        async def _managed_execute() -> Any:
-            result = self.nemo_relay.tools.execute(
-                tool_name,
-                args,
-                _impl,
-                handle=state.handle,
-                data=_jsonable(
-                    {
-                        "turn_id": kwargs.get("turn_id"),
-                        "api_request_id": kwargs.get("api_request_id"),
-                        "tool_call_id": kwargs.get("tool_call_id"),
-                        "mode": self.settings.adaptive_mode,
-                    }
-                ),
-                metadata=_metadata(kwargs),
-            )
-            if inspect.isawaitable(result):
-                return await result
-            return result
+            return _managed_execute()
 
-        managed_result = _resolve_awaitable(_managed_execute())
-        return raw_response["value"] if raw_response["set"] else managed_result
+        return self._run_managed_with_downstream_preservation(
+            next_call, _normalize, _jsonable, _make_managed
+        )
 
 
 def register(ctx) -> None:
@@ -806,6 +841,30 @@ def _value(obj: Any, key: str, default: Any = None) -> Any:
     return getattr(obj, key, default)
 
 
+def _original_downstream_error(exc: Exception) -> BaseException:
+    # Hermes wraps downstream execution failures in a local/private exception
+    # class, so detect the wrapper by shape instead of importing it here.
+    original = getattr(exc, "original", None)
+    if exc.__class__.__name__ == "_DownstreamExecutionError" and isinstance(original, BaseException):
+        return original
+    return exc
+
+
+def _is_relay_wrapped_callback_error(exc: Exception, callback_error: Exception | None) -> bool:
+    # NeMo Relay re-wraps a failing callback as ``RuntimeError("internal error:
+    # <ClassName>: <message>")``. Match by prefix rather than exact equality so a
+    # trailing traceback/suffix in a future Relay version doesn't silently defeat
+    # the unwrap; the class-name + message prefix still discriminates the real
+    # downstream failure from unrelated Relay-internal errors. If Relay drops the
+    # leading ``internal error:`` shape entirely, this returns False and Hermes
+    # falls back to surfacing Relay's error (the pre-fix behavior) rather than
+    # masking it.
+    if callback_error is None or not isinstance(exc, RuntimeError):
+        return False
+    expected = f"internal error: {callback_error.__class__.__name__}: {callback_error}"
+    return str(exc).startswith(expected)
+
+
 def _llm_response_payload(response: Any) -> Any:
     """Return the LLM response shape NeMo Relay's ATIF conversion expects."""
     payload = _jsonable(response)