mirror of
https://github.com/NousResearch/hermes-agent.git
synced 2026-07-20 15:33:54 +00:00
fix(nemo-relay): harden dynamic plugin lifecycle
Signed-off-by: Bryan Bednarski <bbednarski@nvidia.com>
This commit is contained in:
parent
2f74e29e3b
commit
f846a350c7
2 changed files with 235 additions and 37 deletions
|
|
@ -130,21 +130,43 @@ class _Runtime:
|
|||
def _clear_plugins_toml(self) -> None:
|
||||
if not self._plugin_config_initialized:
|
||||
return
|
||||
try:
|
||||
if self._plugin_activation is not None:
|
||||
failures: list[str] = []
|
||||
if self._plugin_activation is not None:
|
||||
activation = self._plugin_activation
|
||||
try:
|
||||
_flush_relay_subscribers(self.nemo_relay)
|
||||
close = getattr(self._plugin_activation, "close", None)
|
||||
if callable(close):
|
||||
except Exception as exc:
|
||||
failures.append(f"subscriber flush failed: {exc}")
|
||||
|
||||
close = getattr(activation, "close", None)
|
||||
if callable(close):
|
||||
try:
|
||||
_resolve_awaitable(close())
|
||||
except Exception as exc:
|
||||
failures.append(f"dynamic plugin activation close failed: {exc}")
|
||||
finally:
|
||||
# Retain the owned activation through the complete close
|
||||
# attempt. The binding transitions it to a terminal state
|
||||
# before its awaitable resolves, including error results.
|
||||
self._plugin_activation = None
|
||||
self._plugin_config_initialized = False
|
||||
self._plugin_config_needs_reinit = bool(self.settings.plugins_config)
|
||||
else:
|
||||
failures.append("dynamic plugin activation has no close method")
|
||||
else:
|
||||
try:
|
||||
plugin_mod = getattr(self.nemo_relay, "plugin", None)
|
||||
clear = getattr(plugin_mod, "clear", None)
|
||||
if callable(clear):
|
||||
_resolve_awaitable(clear())
|
||||
finally:
|
||||
self._plugin_activation = None
|
||||
self._plugin_config_initialized = False
|
||||
self._plugin_config_needs_reinit = bool(self.settings.plugins_config)
|
||||
except Exception as exc:
|
||||
failures.append(f"static plugin configuration clear failed: {exc}")
|
||||
finally:
|
||||
self._plugin_config_initialized = False
|
||||
self._plugin_config_needs_reinit = bool(self.settings.plugins_config)
|
||||
|
||||
if failures:
|
||||
raise RuntimeError("; ".join(failures))
|
||||
|
||||
def _activate_direct_fallbacks(self) -> None:
|
||||
self._plugin_config_needs_reinit = False
|
||||
|
|
@ -267,21 +289,25 @@ class _Runtime:
|
|||
state = self.sessions.pop(session_id, None)
|
||||
if state is None:
|
||||
return
|
||||
failures: list[str] = []
|
||||
if state.handle is not None:
|
||||
try:
|
||||
self.nemo_relay.scope.pop(state.handle, output=_jsonable(kwargs))
|
||||
except Exception:
|
||||
logger.debug("NeMo Relay session pop failed", exc_info=True)
|
||||
except Exception as exc:
|
||||
failures.append(f"session scope pop failed: {exc}")
|
||||
try:
|
||||
_flush_relay_subscribers(self.nemo_relay)
|
||||
except Exception:
|
||||
logger.debug("NeMo Relay subscriber flush failed", exc_info=True)
|
||||
self.export_atif(state)
|
||||
except Exception as exc:
|
||||
failures.append(f"subscriber flush failed: {exc}")
|
||||
try:
|
||||
self.export_atif(state)
|
||||
except Exception as exc:
|
||||
failures.append(f"ATIF export failed: {exc}")
|
||||
if state.atif_exporter is not None and state.atif_subscriber_name:
|
||||
try:
|
||||
state.atif_exporter.deregister(state.atif_subscriber_name)
|
||||
except Exception:
|
||||
logger.debug("NeMo Relay ATIF deregister failed", exc_info=True)
|
||||
except Exception as exc:
|
||||
failures.append(f"ATIF deregister failed: {exc}")
|
||||
if (
|
||||
self._plugin_config_initialized
|
||||
and self._plugin_activation is None
|
||||
|
|
@ -289,28 +315,43 @@ class _Runtime:
|
|||
):
|
||||
try:
|
||||
self._clear_plugins_toml()
|
||||
except Exception:
|
||||
logger.debug("NeMo Relay plugins.toml clear failed", exc_info=True)
|
||||
except Exception as exc:
|
||||
failures.append(f"plugin configuration clear failed: {exc}")
|
||||
elif (
|
||||
self.settings.plugins_config
|
||||
and self._plugin_activation is None
|
||||
and not self.sessions
|
||||
):
|
||||
self._plugin_config_needs_reinit = True
|
||||
if failures:
|
||||
logger.warning(
|
||||
"NeMo Relay session %s teardown completed with errors: %s",
|
||||
session_id,
|
||||
"; ".join(failures),
|
||||
)
|
||||
|
||||
def shutdown(self) -> None:
|
||||
"""Close active sessions and the process-lifetime plugin activation."""
|
||||
failures: list[str] = []
|
||||
for session_id in list(self.sessions):
|
||||
self.close_session({"session_id": session_id, "reason": "runtime_shutdown"})
|
||||
try:
|
||||
self.close_session({"session_id": session_id, "reason": "runtime_shutdown"})
|
||||
except Exception as exc:
|
||||
failures.append(f"session {session_id} close failed: {exc}")
|
||||
if self._plugin_config_initialized:
|
||||
try:
|
||||
self._clear_plugins_toml()
|
||||
except Exception:
|
||||
logger.debug("NeMo Relay plugin runtime shutdown failed", exc_info=True)
|
||||
except Exception as exc:
|
||||
failures.append(f"plugin runtime close failed: {exc}")
|
||||
self._clear_atof()
|
||||
if self._shutdown_registered:
|
||||
if self._shutdown_registered and self._plugin_activation is None:
|
||||
atexit.unregister(self.shutdown)
|
||||
self._shutdown_registered = False
|
||||
if failures:
|
||||
logger.warning(
|
||||
"NeMo Relay runtime shutdown completed with errors: %s",
|
||||
"; ".join(failures),
|
||||
)
|
||||
|
||||
def mark(self, name: str, kwargs: dict[str, Any]) -> None:
|
||||
state = self.ensure_session(kwargs)
|
||||
|
|
@ -372,7 +413,7 @@ class _Runtime:
|
|||
# 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}
|
||||
raw_response: dict[str, Any] = {"set": False, "value": None, "normalized": None}
|
||||
callback_error: Exception | None = None
|
||||
downstream_error: BaseException | None = None
|
||||
|
||||
|
|
@ -386,7 +427,8 @@ class _Runtime:
|
|||
raise
|
||||
raw_response["set"] = True
|
||||
raw_response["value"] = raw
|
||||
return shape_response(raw)
|
||||
raw_response["normalized"] = shape_response(raw)
|
||||
return raw_response["normalized"]
|
||||
|
||||
try:
|
||||
managed_result = _resolve_awaitable(make_managed_execute(_impl))
|
||||
|
|
@ -394,7 +436,11 @@ class _Runtime:
|
|||
if downstream_error is not None and _is_relay_wrapped_callback_error(exc, callback_error):
|
||||
raise downstream_error
|
||||
raise
|
||||
if preserve_raw_response and raw_response["set"]:
|
||||
if (
|
||||
preserve_raw_response
|
||||
and raw_response["set"]
|
||||
and _json_semantically_equal(managed_result, raw_response["normalized"])
|
||||
):
|
||||
return raw_response["value"]
|
||||
return managed_result
|
||||
|
||||
|
|
@ -763,11 +809,13 @@ def _dynamic_plugin_specs(plugins_config: dict[str, Any] | None) -> list[dict[st
|
|||
return []
|
||||
|
||||
specs: list[dict[str, Any]] = []
|
||||
invalid = False
|
||||
for index, raw_spec in enumerate(raw_specs):
|
||||
if not isinstance(raw_spec, dict):
|
||||
logger.warning(
|
||||
"Ignoring invalid NeMo Relay dynamic_plugins[%d]: expected an object", index
|
||||
"Invalid NeMo Relay dynamic_plugins[%d]: expected an object", index
|
||||
)
|
||||
invalid = True
|
||||
continue
|
||||
plugin_id = raw_spec.get("plugin_id")
|
||||
kind = raw_spec.get("kind")
|
||||
|
|
@ -776,30 +824,35 @@ def _dynamic_plugin_specs(plugins_config: dict[str, Any] | None) -> list[dict[st
|
|||
environment_ref = raw_spec.get("environment_ref")
|
||||
if not isinstance(plugin_id, str) or not plugin_id.strip():
|
||||
logger.warning(
|
||||
"Ignoring invalid NeMo Relay dynamic_plugins[%d]: plugin_id is required", index
|
||||
"Invalid NeMo Relay dynamic_plugins[%d]: plugin_id is required", index
|
||||
)
|
||||
invalid = True
|
||||
continue
|
||||
if kind not in {"rust_dynamic", "worker"}:
|
||||
logger.warning(
|
||||
"Ignoring invalid NeMo Relay dynamic_plugins[%d]: kind must be rust_dynamic or worker",
|
||||
"Invalid NeMo Relay dynamic_plugins[%d]: kind must be rust_dynamic or worker",
|
||||
index,
|
||||
)
|
||||
invalid = True
|
||||
continue
|
||||
if not isinstance(manifest_ref, str) or not manifest_ref.strip():
|
||||
logger.warning(
|
||||
"Ignoring invalid NeMo Relay dynamic_plugins[%d]: manifest_ref is required", index
|
||||
"Invalid NeMo Relay dynamic_plugins[%d]: manifest_ref is required", index
|
||||
)
|
||||
invalid = True
|
||||
continue
|
||||
if not isinstance(config, dict):
|
||||
logger.warning(
|
||||
"Ignoring invalid NeMo Relay dynamic_plugins[%d]: config must be an object", index
|
||||
"Invalid NeMo Relay dynamic_plugins[%d]: config must be an object", index
|
||||
)
|
||||
invalid = True
|
||||
continue
|
||||
if environment_ref is not None and not isinstance(environment_ref, str):
|
||||
logger.warning(
|
||||
"Ignoring invalid NeMo Relay dynamic_plugins[%d]: environment_ref must be a string",
|
||||
"Invalid NeMo Relay dynamic_plugins[%d]: environment_ref must be a string",
|
||||
index,
|
||||
)
|
||||
invalid = True
|
||||
continue
|
||||
spec: dict[str, Any] = {
|
||||
"plugin_id": plugin_id,
|
||||
|
|
@ -810,6 +863,12 @@ def _dynamic_plugin_specs(plugins_config: dict[str, Any] | None) -> list[dict[st
|
|||
if environment_ref is not None:
|
||||
spec["environment_ref"] = environment_ref
|
||||
specs.append(spec)
|
||||
if invalid:
|
||||
logger.error(
|
||||
"NeMo Relay dynamic plugin configuration is invalid; no dynamic plugins "
|
||||
"will be activated. Continuing with static observability only."
|
||||
)
|
||||
return []
|
||||
return specs
|
||||
|
||||
|
||||
|
|
@ -999,6 +1058,15 @@ def _jsonable(value: Any) -> Any:
|
|||
return str(value)
|
||||
|
||||
|
||||
def _json_semantically_equal(left: Any, right: Any) -> bool:
|
||||
"""Compare JSON-compatible values without conflating booleans and numbers."""
|
||||
try:
|
||||
options = {"ensure_ascii": False, "sort_keys": True, "separators": (",", ":")}
|
||||
return json.dumps(_jsonable(left), **options) == json.dumps(_jsonable(right), **options)
|
||||
except (TypeError, ValueError):
|
||||
return False
|
||||
|
||||
|
||||
def _value(obj: Any, key: str, default: Any = None) -> Any:
|
||||
if isinstance(obj, dict):
|
||||
return obj.get(key, default)
|
||||
|
|
|
|||
|
|
@ -662,6 +662,42 @@ def test_nemo_relay_managed_llm_uses_wire_protocol_for_interceptor_dispatch(
|
|||
assert "rewritten_for" not in result
|
||||
|
||||
|
||||
def test_nemo_relay_managed_llm_returns_post_next_interceptor_result(tmp_path, monkeypatch):
|
||||
fake = _FakeNemoRelay()
|
||||
raw_response = SimpleNamespace(
|
||||
model="fixture",
|
||||
choices=[
|
||||
SimpleNamespace(
|
||||
message=SimpleNamespace(role="assistant", content="raw", tool_calls=[]),
|
||||
finish_reason="stop",
|
||||
)
|
||||
],
|
||||
usage=None,
|
||||
)
|
||||
|
||||
def execute(name, request, func, **kwargs):
|
||||
del name, kwargs
|
||||
normalized = func(_FakeLLMRequest(request.headers, request.content))
|
||||
return {**normalized, "post_next_interceptor": True}
|
||||
|
||||
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="openai",
|
||||
api_mode="chat_completions",
|
||||
model="fixture",
|
||||
request={"messages": []},
|
||||
next_call=lambda request: raw_response,
|
||||
)
|
||||
|
||||
assert result["post_next_interceptor"] is True
|
||||
assert result["assistant_message"]["content"] == "raw"
|
||||
assert result is not raw_response
|
||||
|
||||
|
||||
def test_nemo_relay_managed_tool_returns_post_interceptor_result(tmp_path, monkeypatch):
|
||||
fake = _FakeNemoRelay()
|
||||
|
||||
|
|
@ -730,7 +766,7 @@ def test_nemo_relay_plugin_degrades_to_static_config_on_relay_0_5(
|
|||
assert "require nemo-relay>=0.6,<0.7" in caplog.text
|
||||
|
||||
|
||||
def test_nemo_relay_plugin_ignores_invalid_dynamic_specs(tmp_path, monkeypatch, caplog):
|
||||
def test_nemo_relay_plugin_rejects_invalid_dynamic_specs(tmp_path, monkeypatch, caplog):
|
||||
fake = _FakeNemoRelay()
|
||||
plugin = _fresh_plugin(monkeypatch, fake)
|
||||
plugins_toml = tmp_path / "plugins.toml"
|
||||
|
|
@ -750,6 +786,38 @@ dynamic_plugins = [{ kind = "rust_dynamic", manifest_ref = "missing-id" }]
|
|||
assert "plugin_id is required" in caplog.text
|
||||
|
||||
|
||||
def test_nemo_relay_plugin_rejects_entire_mixed_valid_invalid_dynamic_request(
|
||||
tmp_path, monkeypatch, caplog
|
||||
):
|
||||
fake = _FakeNemoRelay()
|
||||
plugin = _fresh_plugin(monkeypatch, fake)
|
||||
plugins_toml = tmp_path / "plugins.toml"
|
||||
plugins_toml.write_text(
|
||||
f"""
|
||||
version = 1
|
||||
|
||||
[[dynamic_plugins]]
|
||||
plugin_id = "valid-fixture"
|
||||
kind = "rust_dynamic"
|
||||
manifest_ref = "{(tmp_path / "valid" / "relay-plugin.toml").as_posix()}"
|
||||
|
||||
[[dynamic_plugins]]
|
||||
kind = "worker"
|
||||
manifest_ref = "{(tmp_path / "invalid" / "relay-plugin.toml").as_posix()}"
|
||||
""",
|
||||
encoding="utf-8",
|
||||
)
|
||||
monkeypatch.setenv("HERMES_NEMO_RELAY_PLUGINS_TOML", str(plugins_toml))
|
||||
|
||||
with caplog.at_level("WARNING"):
|
||||
plugin.on_session_start(session_id="s1")
|
||||
|
||||
assert not any(event[0] == "plugin.activate_dynamic" for event in fake.events)
|
||||
initialize = next(event for event in fake.events if event[0] == "plugin.initialize")
|
||||
assert "dynamic_plugins" not in initialize[1]
|
||||
assert "no dynamic plugins will be activated" in caplog.text
|
||||
|
||||
|
||||
def test_nemo_relay_plugin_registers_shutdown_after_dynamic_retry(tmp_path, monkeypatch):
|
||||
fake = _FakeNemoRelay()
|
||||
activation_attempts = 0
|
||||
|
|
@ -781,6 +849,66 @@ def test_nemo_relay_plugin_registers_shutdown_after_dynamic_retry(tmp_path, monk
|
|||
assert any(event[0] == "plugin.activation.close" for event in fake.events)
|
||||
|
||||
|
||||
def test_nemo_relay_plugin_attempts_activation_close_after_subscriber_flush_failure(
|
||||
tmp_path, monkeypatch, caplog
|
||||
):
|
||||
fake = _FakeNemoRelay()
|
||||
|
||||
def _failing_flush():
|
||||
fake.events.append(("subscribers.flush.failed",))
|
||||
raise RuntimeError("flush boom")
|
||||
|
||||
fake.subscribers.flush = _failing_flush
|
||||
plugin = _fresh_plugin(monkeypatch, fake)
|
||||
_enable_dynamic_plugin(tmp_path, monkeypatch)
|
||||
plugin.on_session_start(session_id="s1")
|
||||
runtime = plugin._get_runtime()
|
||||
assert runtime is not None
|
||||
|
||||
with caplog.at_level("WARNING"):
|
||||
runtime.shutdown()
|
||||
|
||||
event_names = [event[0] for event in fake.events]
|
||||
assert event_names.count("subscribers.flush.failed") == 2
|
||||
flush_indices = [
|
||||
index for index, name in enumerate(event_names) if name == "subscribers.flush.failed"
|
||||
]
|
||||
assert max(flush_indices) < event_names.index("plugin.activation.close")
|
||||
assert runtime._plugin_activation is None
|
||||
assert "subscriber flush failed: flush boom" in caplog.text
|
||||
|
||||
|
||||
def test_nemo_relay_plugin_continues_shutdown_after_atif_export_failure(
|
||||
tmp_path, monkeypatch, caplog
|
||||
):
|
||||
fake = _FakeNemoRelay()
|
||||
|
||||
class _FailingAtifExporter(_FakeAtifExporter):
|
||||
def export_json(self):
|
||||
self.events.append(("atif.export.failed", self.session_id))
|
||||
raise OSError("disk full")
|
||||
|
||||
fake.AtifExporter = lambda session_id, agent_name, agent_version, **kwargs: (
|
||||
_FailingAtifExporter(fake.events, session_id, agent_name, agent_version, kwargs)
|
||||
)
|
||||
plugin = _fresh_plugin(monkeypatch, fake)
|
||||
_enable_dynamic_plugin(tmp_path, monkeypatch)
|
||||
monkeypatch.setenv("HERMES_NEMO_RELAY_ATIF_ENABLED", "1")
|
||||
monkeypatch.setenv("HERMES_NEMO_RELAY_ATIF_OUTPUT_DIRECTORY", str(tmp_path / "atif"))
|
||||
plugin.on_session_start(session_id="s1")
|
||||
runtime = plugin._get_runtime()
|
||||
assert runtime is not None
|
||||
|
||||
with caplog.at_level("WARNING"):
|
||||
runtime.shutdown()
|
||||
|
||||
event_names = [event[0] for event in fake.events]
|
||||
assert event_names.index("atif.export.failed") < event_names.index("atif.deregister")
|
||||
assert event_names.index("atif.deregister") < event_names.index("plugin.activation.close")
|
||||
assert runtime._plugin_activation is None
|
||||
assert "ATIF export failed: disk full" in caplog.text
|
||||
|
||||
|
||||
def test_nemo_relay_plugin_keeps_plugins_toml_active_while_other_sessions_remain(tmp_path, monkeypatch):
|
||||
fake = _FakeNemoRelay()
|
||||
plugin = _fresh_plugin(monkeypatch, fake)
|
||||
|
|
@ -1037,15 +1165,16 @@ mode = "observe_only"
|
|||
),
|
||||
finish_reason="tool_calls",
|
||||
)
|
||||
raw_response = SimpleNamespace(
|
||||
id="resp-1",
|
||||
model="demo-model",
|
||||
choices=[raw_choice],
|
||||
usage=SimpleNamespace(prompt_tokens=3, completion_tokens=5, total_tokens=8),
|
||||
)
|
||||
|
||||
def next_call(request):
|
||||
seen_request.update(request)
|
||||
return SimpleNamespace(
|
||||
id="resp-1",
|
||||
model="demo-model",
|
||||
choices=[raw_choice],
|
||||
usage=SimpleNamespace(prompt_tokens=3, completion_tokens=5, total_tokens=8),
|
||||
)
|
||||
return raw_response
|
||||
|
||||
response = plugin.on_llm_execution_middleware(
|
||||
session_id="s1",
|
||||
|
|
@ -1059,6 +1188,7 @@ mode = "observe_only"
|
|||
next_call=next_call,
|
||||
)
|
||||
|
||||
assert response is raw_response
|
||||
assert response.model == "demo-model"
|
||||
assert response.choices == [raw_choice]
|
||||
assert seen_request["intercepted"] is True
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue