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>
2026-06-09 02:31:10 -07:00
parent 8d99b5bc4f
commit 85852b71d8
2 changed files with 479 additions and 57 deletions
@@ -12,6 +12,7 @@ import warnings
 from pathlib import Path
 from types import SimpleNamespace

+import pytest
 import yaml

 from hermes_cli.plugins import PluginManager
@@ -153,6 +154,33 @@ def _fresh_plugin(monkeypatch, fake):
    return plugin


+def _wrapped_downstream_error(original):
+    class _DownstreamExecutionError(Exception):
+        def __init__(self, original):
+            super().__init__(str(original))
+            self.original = original
+
+    return _DownstreamExecutionError(original)
+
+
+def _enable_adaptive_plugin(tmp_path, monkeypatch) -> None:
+    plugins_toml = tmp_path / "plugins.toml"
+    plugins_toml.write_text(
+        """
+version = 1
+
+[[components]]
+kind = "adaptive"
+enabled = true
+
+[components.config.tool_parallelism]
+mode = "observe_only"
+""",
+        encoding="utf-8",
+    )
+    monkeypatch.setenv("HERMES_NEMO_RELAY_PLUGINS_TOML", str(plugins_toml))
+
+
 def test_manifest_fields():
    data = yaml.safe_load((PLUGIN_DIR / "plugin.yaml").read_text())
    assert data["name"] == "nemo_relay"
@@ -783,6 +811,220 @@ mode = "observe_only"
    }


+def test_nemo_relay_adaptive_llm_execution_preserves_downstream_error(tmp_path, monkeypatch):
+    fake = _FakeNemoRelay()
+
+    def native_like_execute(name, request, func, **kwargs):
+        fake.events.append(("llm.execute.start", name, request.content, kwargs))
+        try:
+            return func(_FakeLLMRequest(request.headers, {"intercepted": True, **request.content}))
+        except Exception as exc:
+            raise RuntimeError(f"internal error: {type(exc).__name__}: {exc}") from None
+
+    fake.llm.execute = native_like_execute
+    plugin = _fresh_plugin(monkeypatch, fake)
+    _enable_adaptive_plugin(tmp_path, monkeypatch)
+
+    class ProviderAuthError(Exception):
+        status_code = 403
+
+    provider_error = ProviderAuthError("provider auth failed")
+
+    def next_call(request):
+        raise _wrapped_downstream_error(provider_error)
+
+    with pytest.raises(ProviderAuthError) as caught:
+        plugin.on_llm_execution_middleware(
+            session_id="s1",
+            provider="anthropic",
+            model="demo-model",
+            request={"messages": [{"role": "user", "content": "hi"}]},
+            next_call=next_call,
+        )
+
+    assert caught.value is provider_error
+    assert caught.value.status_code == 403
+
+
+def test_nemo_relay_adaptive_llm_execution_preserves_downstream_error_with_relay_suffix(
+    tmp_path, monkeypatch
+):
+    # Guards the startswith (vs exact ==) match in _is_relay_wrapped_callback_error:
+    # Relay re-wraps the callback failure with its canonical prefix but APPENDS a
+    # trailing suffix. Exact equality would miss this and surface Relay's wrapper;
+    # prefix matching must still recover the original downstream error.
+    fake = _FakeNemoRelay()
+
+    def native_like_execute(name, request, func, **kwargs):
+        try:
+            return func(_FakeLLMRequest(request.headers, {"intercepted": True, **request.content}))
+        except Exception as exc:
+            raise RuntimeError(f"internal error: {type(exc).__name__}: {exc} (retried 3x)") from None
+
+    fake.llm.execute = native_like_execute
+    plugin = _fresh_plugin(monkeypatch, fake)
+    _enable_adaptive_plugin(tmp_path, monkeypatch)
+
+    class ProviderAuthError(Exception):
+        status_code = 403
+
+    provider_error = ProviderAuthError("provider auth failed")
+
+    def next_call(request):
+        raise _wrapped_downstream_error(provider_error)
+
+    with pytest.raises(ProviderAuthError) as caught:
+        plugin.on_llm_execution_middleware(
+            session_id="s1",
+            provider="anthropic",
+            model="demo-model",
+            request={"messages": [{"role": "user", "content": "hi"}]},
+            next_call=next_call,
+        )
+
+    assert caught.value is provider_error
+    assert caught.value.status_code == 403
+
+
+def test_nemo_relay_adaptive_llm_execution_keeps_unrelated_internal_error(tmp_path, monkeypatch):
+    fake = _FakeNemoRelay()
+
+    relay_error = RuntimeError("internal error: relay setup failed")
+
+    def internal_error_execute(name, request, func, **kwargs):
+        raise relay_error
+
+    fake.llm.execute = internal_error_execute
+    plugin = _fresh_plugin(monkeypatch, fake)
+    _enable_adaptive_plugin(tmp_path, monkeypatch)
+
+    with pytest.raises(RuntimeError) as caught:
+        plugin.on_llm_execution_middleware(
+            session_id="s1",
+            provider="anthropic",
+            model="demo-model",
+            request={"messages": [{"role": "user", "content": "hi"}]},
+            next_call=lambda request: {"raw": request},
+        )
+
+    assert caught.value is relay_error
+
+
+def test_nemo_relay_adaptive_llm_execution_keeps_wrapped_relay_error_after_downstream_failure(
+    tmp_path, monkeypatch
+):
+    fake = _FakeNemoRelay()
+    relay_error = RuntimeError("internal error: RuntimeError: relay policy blocked after downstream")
+
+    def translated_execute(name, request, func, **kwargs):
+        try:
+            return func(_FakeLLMRequest(request.headers, {"intercepted": True, **request.content}))
+        except Exception:
+            raise relay_error
+
+    fake.llm.execute = translated_execute
+    plugin = _fresh_plugin(monkeypatch, fake)
+    _enable_adaptive_plugin(tmp_path, monkeypatch)
+
+    def next_call(request):
+        raise _wrapped_downstream_error(RuntimeError("provider failed"))
+
+    with pytest.raises(RuntimeError) as caught:
+        plugin.on_llm_execution_middleware(
+            session_id="s1",
+            provider="anthropic",
+            model="demo-model",
+            request={"messages": [{"role": "user", "content": "hi"}]},
+            next_call=next_call,
+        )
+
+    assert caught.value is relay_error
+
+
+def test_nemo_relay_adaptive_llm_execution_keeps_relay_translated_error(tmp_path, monkeypatch):
+    fake = _FakeNemoRelay()
+
+    class RelayPolicyError(Exception):
+        pass
+
+    relay_error = RelayPolicyError("relay policy blocked")
+
+    def translated_execute(name, request, func, **kwargs):
+        try:
+            return func(_FakeLLMRequest(request.headers, {"intercepted": True, **request.content}))
+        except Exception:
+            raise relay_error
+
+    fake.llm.execute = translated_execute
+    plugin = _fresh_plugin(monkeypatch, fake)
+    _enable_adaptive_plugin(tmp_path, monkeypatch)
+
+    provider_error = RuntimeError("provider failed")
+
+    def next_call(request):
+        raise _wrapped_downstream_error(provider_error)
+
+    with pytest.raises(RelayPolicyError) as caught:
+        plugin.on_llm_execution_middleware(
+            session_id="s1",
+            provider="anthropic",
+            model="demo-model",
+            request={"messages": [{"role": "user", "content": "hi"}]},
+            next_call=next_call,
+        )
+
+    assert caught.value is relay_error
+
+
+def test_nemo_relay_downstream_unwrap_matches_real_middleware_wrapper_shape(monkeypatch):
+    # Regression guard against core/plugin drift. The synthetic tests above model
+    # the downstream-error wrapper with a local class, so they keep passing even
+    # if core middleware renames its private ``_DownstreamExecutionError`` or drops
+    # ``.original`` -- the exact shape the plugin matches by name at
+    # ``_original_downstream_error``. Capture the wrapper the REAL
+    # ``hermes_cli.middleware._run_execution_chain`` hands to a middleware
+    # callback's ``next_call`` and assert the plugin's detector unwraps it to the
+    # original exception. If core middleware changes the wrapper shape, this fails
+    # here instead of silently defeating the unwrap in production.
+    from hermes_cli import middleware
+
+    from plugins.observability.nemo_relay import _original_downstream_error
+
+    class ProviderError(Exception):
+        status_code = 403
+
+    provider_error = ProviderError("provider auth failed")
+    captured: dict[str, Exception] = {}
+
+    def terminal_call(payload):
+        raise provider_error
+
+    def capturing_callback(**kwargs):
+        next_call = kwargs["next_call"]
+        try:
+            return next_call(kwargs.get("request"))
+        except Exception as exc:
+            captured["wrapper"] = exc
+            # Surface the original so the chain unwinds without re-wrapping noise.
+            raise _original_downstream_error(exc) from None
+
+    with pytest.raises(ProviderError) as caught:
+        middleware._run_execution_chain(
+            "llm",
+            [capturing_callback],
+            terminal_call,
+            request={"messages": []},
+        )
+
+    wrapper = captured["wrapper"]
+    # The wrapper the plugin sees must match what _original_downstream_error keys on.
+    assert wrapper.__class__.__name__ == "_DownstreamExecutionError"
+    assert isinstance(getattr(wrapper, "original", None), BaseException)
+    assert _original_downstream_error(wrapper) is provider_error
+    assert caught.value is provider_error
+    assert caught.value.status_code == 403
+
+
 def _adaptive_llm_execute_mode(tmp_path, monkeypatch, plugins_toml_text: str) -> str:
    fake = _FakeNemoRelay()
    plugin = _fresh_plugin(monkeypatch, fake)
@@ -920,6 +1162,127 @@ mode = "observe_only"
    assert execute_start[3]["data"]["tool_call_id"] == "tool-1"


+def test_nemo_relay_adaptive_tool_execution_preserves_downstream_error(tmp_path, monkeypatch):
+    fake = _FakeNemoRelay()
+
+    def native_like_execute(name, args, func, **kwargs):
+        fake.events.append(("tool.execute.start", name, args, kwargs))
+        try:
+            return func({"intercepted": True, **args})
+        except Exception as exc:
+            raise RuntimeError(f"internal error: {type(exc).__name__}: {exc}") from None
+
+    fake.tools.execute = native_like_execute
+    plugin = _fresh_plugin(monkeypatch, fake)
+    _enable_adaptive_plugin(tmp_path, monkeypatch)
+
+    class ToolAuthError(Exception):
+        status_code = 403
+
+    tool_error = ToolAuthError("tool auth failed")
+
+    def next_call(args):
+        raise _wrapped_downstream_error(tool_error)
+
+    with pytest.raises(ToolAuthError) as caught:
+        plugin.on_tool_execution_middleware(
+            session_id="s1",
+            tool_name="terminal",
+            args={"command": "pwd"},
+            next_call=next_call,
+        )
+
+    assert caught.value is tool_error
+    assert caught.value.status_code == 403
+
+
+def test_nemo_relay_adaptive_tool_execution_keeps_unrelated_internal_error(tmp_path, monkeypatch):
+    fake = _FakeNemoRelay()
+
+    relay_error = RuntimeError("internal error: relay setup failed")
+
+    def internal_error_execute(name, args, func, **kwargs):
+        raise relay_error
+
+    fake.tools.execute = internal_error_execute
+    plugin = _fresh_plugin(monkeypatch, fake)
+    _enable_adaptive_plugin(tmp_path, monkeypatch)
+
+    with pytest.raises(RuntimeError) as caught:
+        plugin.on_tool_execution_middleware(
+            session_id="s1",
+            tool_name="terminal",
+            args={"command": "pwd"},
+            next_call=lambda args: {"raw": args},
+        )
+
+    assert caught.value is relay_error
+
+
+def test_nemo_relay_adaptive_tool_execution_keeps_wrapped_relay_error_after_downstream_failure(
+    tmp_path, monkeypatch
+):
+    fake = _FakeNemoRelay()
+    relay_error = RuntimeError("internal error: RuntimeError: relay policy blocked after downstream")
+
+    def translated_execute(name, args, func, **kwargs):
+        try:
+            return func({"intercepted": True, **args})
+        except Exception:
+            raise relay_error
+
+    fake.tools.execute = translated_execute
+    plugin = _fresh_plugin(monkeypatch, fake)
+    _enable_adaptive_plugin(tmp_path, monkeypatch)
+
+    def next_call(args):
+        raise _wrapped_downstream_error(RuntimeError("tool failed"))
+
+    with pytest.raises(RuntimeError) as caught:
+        plugin.on_tool_execution_middleware(
+            session_id="s1",
+            tool_name="terminal",
+            args={"command": "pwd"},
+            next_call=next_call,
+        )
+
+    assert caught.value is relay_error
+
+
+def test_nemo_relay_adaptive_tool_execution_keeps_relay_translated_error(tmp_path, monkeypatch):
+    fake = _FakeNemoRelay()
+
+    class RelayPolicyError(Exception):
+        pass
+
+    relay_error = RelayPolicyError("relay policy blocked")
+
+    def translated_execute(name, args, func, **kwargs):
+        try:
+            return func({"intercepted": True, **args})
+        except Exception:
+            raise relay_error
+
+    fake.tools.execute = translated_execute
+    plugin = _fresh_plugin(monkeypatch, fake)
+    _enable_adaptive_plugin(tmp_path, monkeypatch)
+
+    tool_error = RuntimeError("tool failed")
+
+    def next_call(args):
+        raise _wrapped_downstream_error(tool_error)
+
+    with pytest.raises(RelayPolicyError) as caught:
+        plugin.on_tool_execution_middleware(
+            session_id="s1",
+            tool_name="terminal",
+            args={"command": "pwd"},
+            next_call=next_call,
+        )
+
+    assert caught.value is relay_error
+
+
 def test_nemo_relay_tool_execution_middleware_calls_through_without_adaptive(monkeypatch):
    fake = _FakeNemoRelay()
    plugin = _fresh_plugin(monkeypatch, fake)