fix(observability): coordinate Relay plugin lifecycle

Signed-off-by: Alex Fournier <afournier@nvidia.com>

plugins/observability/nemo_relay/__init__.py

@@ -60,6 +60,140 @@ class _Settings:
     atif_model_name: str = "unknown"
 
 
+class _ProcessPluginConfiguration:
+    """Own Relay's process-global plugin configuration across profile runtimes."""
+
+    def __init__(self) -> None:
+        self._lock = threading.RLock()
+        self._key: str | None = None
+        self._plugin_mod: Any = None
+        self._activation: Any = None
+        self._owners: set[int] = set()
+
+    def acquire(
+        self,
+        owner: "_Runtime",
+        plugin_mod: Any,
+        plugin_config: dict[str, Any],
+        dynamic_plugins: list[dict[str, Any]],
+    ) -> tuple[bool, Any]:
+        owner_id = id(owner)
+        key = _plugin_configuration_key(plugin_config, dynamic_plugins)
+        with self._lock:
+            if owner_id in self._owners:
+                return True, self._activation
+            if self._owners:
+                if self._plugin_mod is plugin_mod and self._key == key:
+                    self._owners.add(owner_id)
+                    return True, self._activation
+                logger.warning(
+                    "NeMo Relay plugin configuration is already active for another "
+                    "Hermes profile; keeping the existing process-global configuration "
+                    "and using direct observability for this profile."
+                )
+                return False, None
+
+            activation = None
+            if dynamic_plugins:
+                initialize_dynamic = getattr(
+                    plugin_mod,
+                    "initialize_with_dynamic_plugins",
+                    None,
+                )
+                if callable(initialize_dynamic):
+                    try:
+                        activation = _resolve_awaitable(
+                            initialize_dynamic(plugin_config, dynamic_plugins)
+                        )
+                    except Exception as exc:
+                        logger.warning(
+                            "NeMo Relay dynamic plugin activation failed; continuing "
+                            "with static observability only: %s",
+                            exc,
+                        )
+                else:
+                    logger.warning(
+                        "NeMo Relay dynamic plugins require a binding that exposes "
+                        "plugin.initialize_with_dynamic_plugins (available in NeMo "
+                        "Relay 0.6+). Continuing with static observability only."
+                    )
+
+            if activation is None:
+                initialize = getattr(plugin_mod, "initialize", None)
+                if not callable(initialize):
+                    return False, None
+                try:
+                    _resolve_awaitable(initialize(plugin_config))
+                except Exception as exc:
+                    logger.debug(
+                        "NeMo Relay plugins.toml init failed: %s",
+                        exc,
+                        exc_info=True,
+                    )
+                    return False, None
+
+            self._key = key
+            self._plugin_mod = plugin_mod
+            self._activation = activation
+            self._owners.add(owner_id)
+            return True, activation
+
+    def release(self, owner: "_Runtime", nemo_relay: Any) -> None:
+        owner_id = id(owner)
+        with self._lock:
+            if owner_id not in self._owners:
+                return
+            self._owners.remove(owner_id)
+            if self._owners:
+                return
+
+            failures: list[str] = []
+            activation = self._activation
+            plugin_mod = self._plugin_mod
+            try:
+                if activation is not None:
+                    try:
+                        _flush_relay_subscribers(nemo_relay)
+                    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}"
+                            )
+                    else:
+                        failures.append("dynamic plugin activation has no close method")
+                else:
+                    clear = getattr(plugin_mod, "clear", None)
+                    if callable(clear):
+                        try:
+                            _resolve_awaitable(clear())
+                        except Exception as exc:
+                            failures.append(
+                                f"static plugin configuration clear failed: {exc}"
+                            )
+            finally:
+                self._key = None
+                self._plugin_mod = None
+                self._activation = None
+
+            if failures:
+                raise RuntimeError("; ".join(failures))
+
+    def reset_for_tests(self) -> None:
+        with self._lock:
+            self._key = None
+            self._plugin_mod = None
+            self._activation = None
+            self._owners.clear()
+
+
+_PLUGIN_CONFIGURATION = _ProcessPluginConfiguration()
+
+
 class _Runtime:
     def __init__(
         self,
@@ -107,38 +241,17 @@ class _Runtime:
         if plugin_mod is None:
             return False
         plugin_config = _static_plugin_config(self.settings.plugins_config)
-        if self.settings.dynamic_plugins:
-            activate_dynamic = getattr(plugin_mod, "activate_dynamic_plugins", None)
-            if callable(activate_dynamic):
-                try:
-                    self._ensure_plugin_config_output_dirs(plugin_config)
-                    self._plugin_activation = _resolve_awaitable(
-                        activate_dynamic(plugin_config, self.settings.dynamic_plugins)
-                    )
-                    self._ensure_shutdown_registered()
-                    return True
-                except Exception as exc:
-                    logger.warning(
-                        "NeMo Relay dynamic plugin activation failed; continuing with static "
-                        "observability only: %s",
-                        exc,
-                    )
-            else:
-                logger.warning(
-                    "NeMo Relay dynamic plugins require a binding that exposes "
-                    "plugin.activate_dynamic_plugins (available in NeMo Relay 0.6+). "
-                    "Continuing with static observability only."
-                )
-        initialize = getattr(plugin_mod, "initialize", None)
-        if not callable(initialize):
-            return False
-        try:
-            self._ensure_plugin_config_output_dirs(plugin_config)
-            _resolve_awaitable(initialize(plugin_config))
-            return True
-        except Exception as exc:
-            logger.debug("NeMo Relay plugins.toml init failed: %s", exc, exc_info=True)
-            return False
+        self._ensure_plugin_config_output_dirs(plugin_config)
+        initialized, activation = _PLUGIN_CONFIGURATION.acquire(
+            self,
+            plugin_mod,
+            plugin_config,
+            self.settings.dynamic_plugins,
+        )
+        self._plugin_activation = activation
+        if activation is not None:
+            self._ensure_shutdown_registered()
+        return initialized
 
     def _ensure_shutdown_registered(self) -> None:
         if self._shutdown_registered:
@@ -149,43 +262,12 @@ class _Runtime:
     def _clear_plugins_toml(self) -> None:
         if not self._plugin_config_initialized:
             return
-        failures: list[str] = []
-        if self._plugin_activation is not None:
-            activation = self._plugin_activation
-            try:
-                _flush_relay_subscribers(self.nemo_relay)
-            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())
-            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))
+        try:
+            _PLUGIN_CONFIGURATION.release(self, self.nemo_relay)
+        finally:
+            self._plugin_activation = None
+            self._plugin_config_initialized = False
+            self._plugin_config_needs_reinit = bool(self.settings.plugins_config)
 
     def _activate_direct_fallbacks(self) -> None:
         self._plugin_config_needs_reinit = False
@@ -610,6 +692,18 @@ def _static_plugin_config(plugins_config: dict[str, Any]) -> dict[str, Any]:
     }
 
 
+def _plugin_configuration_key(
+    plugin_config: dict[str, Any],
+    dynamic_plugins: list[dict[str, Any]],
+) -> str:
+    return json.dumps(
+        {"config": plugin_config, "dynamic_plugins": dynamic_plugins},
+        sort_keys=True,
+        separators=(",", ":"),
+        default=str,
+    )
+
+
 def _dynamic_plugin_specs(
     plugins_config: dict[str, Any] | None,
     plugins_toml_path: str = "",
@@ -926,3 +1020,4 @@ def reset_for_tests() -> None:
     for runtime in runtimes:
         if isinstance(runtime, _Runtime):
             runtime.shutdown()
+    _PLUGIN_CONFIGURATION.reset_for_tests()

tests/plugins/test_nemo_relay_plugin.py

@@ -53,7 +53,7 @@ class _FakeNemoRelay:
         self.plugin = SimpleNamespace(
             initialize=self._plugin_initialize,
             clear=self._plugin_clear,
-            activate_dynamic_plugins=self._plugin_activate_dynamic,
+            initialize_with_dynamic_plugins=self._plugin_initialize_with_dynamic,
         )
         self.subscribers = SimpleNamespace(
             register=self._register_subscriber,
@@ -159,7 +159,7 @@ class _FakeNemoRelay:
     async def _plugin_clear(self):
         self.events.append(("plugin.clear",))
 
-    async def _plugin_activate_dynamic(self, config, dynamic_plugins):
+    async def _plugin_initialize_with_dynamic(self, config, dynamic_plugins):
         self.events.append(("plugin.activate_dynamic", config, dynamic_plugins))
         return _FakePluginActivation(self.events)
 
@@ -770,6 +770,190 @@ 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")
 
 
+def test_nemo_relay_plugin_uses_real_0_6_dynamic_activation_api(
+    tmp_path,
+    monkeypatch,
+):
+    relay = pytest.importorskip("nemo_relay")
+    if getattr(relay, "_native", None) is None:
+        pytest.skip("NeMo Relay native binding is unavailable on this platform")
+    plugin = _fresh_plugin(monkeypatch, relay)
+    _enable_dynamic_plugin(tmp_path, monkeypatch)
+    calls = []
+
+    class _NativeActivation:
+        def __init__(self):
+            self.report = {"diagnostics": []}
+            self.is_active = True
+
+        async def close(self):
+            self.is_active = False
+
+    async def _initialize_with_dynamic_plugins(config, dynamic_plugins):
+        calls.append((config, dynamic_plugins))
+        return _NativeActivation()
+
+    monkeypatch.setattr(
+        relay.plugin,
+        "_initialize_with_dynamic_plugins",
+        _initialize_with_dynamic_plugins,
+    )
+
+    plugin.on_session_start(session_id="s1")
+    runtime = plugin._get_runtime()
+
+    assert runtime is not None
+    assert isinstance(runtime._plugin_activation, relay.plugin.PluginHostActivation)
+    assert calls == [
+        (
+            {"version": 1},
+            [
+                {
+                    "plugin_id": "fixture",
+                    "kind": "rust_dynamic",
+                    "manifest_ref": str(tmp_path / "fixture" / "relay-plugin.toml"),
+                    "config": {"mode": "test"},
+                }
+            ],
+        )
+    ]
+
+    activation = runtime._plugin_activation
+    runtime.shutdown()
+    assert activation.is_active is False
+
+
+def test_real_binding_shares_plugin_configuration_across_two_profiles(
+    tmp_path,
+    monkeypatch,
+):
+    relay = pytest.importorskip("nemo_relay")
+    if getattr(relay, "_native", None) is None:
+        pytest.skip("NeMo Relay native binding is unavailable on this platform")
+    plugin = _fresh_plugin(monkeypatch, relay)
+    original_initialize = relay.plugin.initialize
+    original_clear = relay.plugin.clear
+    original_clear()
+    initialize_calls = []
+    clear_calls = 0
+
+    async def _initialize(config):
+        initialize_calls.append(config)
+        return await original_initialize(config)
+
+    def _clear():
+        nonlocal clear_calls
+        clear_calls += 1
+        return original_clear()
+
+    monkeypatch.setattr(relay.plugin, "initialize", _initialize)
+    monkeypatch.setattr(relay.plugin, "clear", _clear)
+    monkeypatch.setattr(
+        plugin,
+        "_load_settings",
+        lambda: plugin._Settings(plugins_config={"version": 1}),
+    )
+    profile_a = str(tmp_path / "profile-a")
+    profile_b = str(tmp_path / "profile-b")
+    host_a = relay_runtime.RelayRuntime(relay=relay, profile_key=profile_a)
+    host_b = relay_runtime.RelayRuntime(relay=relay, profile_key=profile_b)
+
+    try:
+        runtime_a = plugin._get_runtime(profile_key=profile_a, host=host_a)
+        runtime_b = plugin._get_runtime(profile_key=profile_b, host=host_b)
+        assert runtime_a is not None
+        assert runtime_b is not None
+        runtime_a.ensure_session({"session_id": "session-a"})
+        runtime_b.ensure_session({"session_id": "session-b"})
+
+        assert initialize_calls == [{"version": 1}]
+        assert relay.plugin.report() is not None
+
+        runtime_a.close_session({"session_id": "session-a"})
+
+        assert clear_calls == 0
+        assert relay.plugin.report() is not None
+        assert runtime_b.host.get_session("session-b") is not None
+
+        runtime_b.close_session({"session_id": "session-b"})
+
+        assert clear_calls == 1
+        assert relay.plugin.report() is None
+    finally:
+        plugin.reset_for_tests()
+        host_a.shutdown()
+        host_b.shutdown()
+        original_clear()
+
+
+def test_real_binding_does_not_replace_another_profiles_plugin_configuration(
+    tmp_path,
+    monkeypatch,
+):
+    relay = pytest.importorskip("nemo_relay")
+    if getattr(relay, "_native", None) is None:
+        pytest.skip("NeMo Relay native binding is unavailable on this platform")
+    plugin = _fresh_plugin(monkeypatch, relay)
+    original_initialize = relay.plugin.initialize
+    original_clear = relay.plugin.clear
+    original_clear()
+    initialize_calls = []
+    clear_calls = 0
+
+    async def _initialize(config):
+        initialize_calls.append(config)
+        return await original_initialize(config)
+
+    def _clear():
+        nonlocal clear_calls
+        clear_calls += 1
+        return original_clear()
+
+    settings = iter((
+        plugin._Settings(
+            plugins_config={"version": 1, "policy": {"unsupported": "warn"}}
+        ),
+        plugin._Settings(
+            plugins_config={"version": 1, "policy": {"unsupported": "ignore"}}
+        ),
+    ))
+    monkeypatch.setattr(relay.plugin, "initialize", _initialize)
+    monkeypatch.setattr(relay.plugin, "clear", _clear)
+    monkeypatch.setattr(plugin, "_load_settings", lambda: next(settings))
+    profile_a = str(tmp_path / "profile-a")
+    profile_b = str(tmp_path / "profile-b")
+    host_a = relay_runtime.RelayRuntime(relay=relay, profile_key=profile_a)
+    host_b = relay_runtime.RelayRuntime(relay=relay, profile_key=profile_b)
+
+    try:
+        runtime_a = plugin._get_runtime(profile_key=profile_a, host=host_a)
+        runtime_b = plugin._get_runtime(profile_key=profile_b, host=host_b)
+        assert runtime_a is not None
+        assert runtime_b is not None
+
+        assert initialize_calls == [
+            {"version": 1, "policy": {"unsupported": "warn"}}
+        ]
+        assert runtime_a._plugin_config_initialized is True
+        assert runtime_b._plugin_config_initialized is False
+        assert relay.plugin.report() is not None
+
+        runtime_b.shutdown()
+
+        assert clear_calls == 0
+        assert relay.plugin.report() is not None
+
+        runtime_a.shutdown()
+
+        assert clear_calls == 1
+        assert relay.plugin.report() is None
+    finally:
+        plugin.reset_for_tests()
+        host_a.shutdown()
+        host_b.shutdown()
+        original_clear()
+
+
 def test_nemo_relay_rejects_gateway_dynamic_config_with_actionable_diagnostic(
     tmp_path, monkeypatch, caplog
 ):
@@ -897,7 +1081,7 @@ def test_nemo_relay_plugin_degrades_to_static_config_on_relay_0_5(
     tmp_path, monkeypatch, caplog
 ):
     fake = _FakeNemoRelay()
-    delattr(fake.plugin, "activate_dynamic_plugins")
+    delattr(fake.plugin, "initialize_with_dynamic_plugins")
     plugin = _fresh_plugin(monkeypatch, fake)
     _enable_dynamic_plugin(tmp_path, monkeypatch)
 
@@ -976,7 +1160,7 @@ def test_nemo_relay_plugin_registers_shutdown_after_dynamic_retry(tmp_path, monk
             raise RuntimeError("temporary activation failure")
         return _FakePluginActivation(fake.events)
 
-    fake.plugin.activate_dynamic_plugins = _flaky_activate
+    fake.plugin.initialize_with_dynamic_plugins = _flaky_activate
     plugin = _fresh_plugin(monkeypatch, fake)
     _enable_dynamic_plugin(tmp_path, monkeypatch)