fix(plugins): give dual-kind memory hooks a single owner
This commit is contained in:
Joey
Teknium
@@ -1127,6 +1127,8 @@ class PluginManager(PluginLoaderMixin, PluginDispatchMixin, PluginLedgerMixin):
|
||||
# (matcher, callback, plugin_name), platform handler factories (lowercase platform -> list).
|
||||
self._plugins: Dict[str, LoadedPlugin] = {}
|
||||
self._hooks: Dict[str, List[Callable]] = {}
|
||||
# Fallback hooks registered by a memory provider before general discovery.
|
||||
self._memory_hook_registrations: Dict[Tuple[str, str], List[PluginRegistration]] = {}
|
||||
self._middleware: Dict[str, List[Callable]] = {}
|
||||
self._plugin_tool_names: Set[str] = set()
|
||||
self._plugin_platform_names: Set[str] = set()
|
||||
|
||||
@@ -161,12 +161,13 @@ class PluginLedgerMixin:
|
||||
return
|
||||
ids = {id(r) for r in registrations}
|
||||
self._registration_order = [r for r in self._registration_order if id(r) not in ids]
|
||||
for plugin_key, owned in list(self._ownership_ledger.items()):
|
||||
remaining = [r for r in owned if id(r) not in ids]
|
||||
if remaining:
|
||||
self._ownership_ledger[plugin_key] = remaining
|
||||
else:
|
||||
self._ownership_ledger.pop(plugin_key, None)
|
||||
for ledger in (self._ownership_ledger, self._memory_hook_registrations):
|
||||
for plugin_key, owned in list(ledger.items()):
|
||||
remaining = [r for r in owned if id(r) not in ids]
|
||||
if remaining:
|
||||
ledger[plugin_key] = remaining
|
||||
else:
|
||||
ledger.pop(plugin_key, None)
|
||||
|
||||
def _dispose_registrations(self, registrations: List[PluginRegistration]) -> None:
|
||||
"""Dispose registrations in reverse acquisition order, best effort."""
|
||||
|
||||
@@ -307,6 +307,11 @@ class PluginLoaderMixin:
|
||||
register_fn(PluginContext(manifest, self))
|
||||
self._attribute_registrations(loaded, plugin_key, registration_start)
|
||||
loaded.enabled = True
|
||||
source = getattr(module, "__file__", None)
|
||||
if source:
|
||||
hook_source = (manifest.name, str(Path(source).resolve()))
|
||||
for handle in self._memory_hook_registrations.pop(hook_source, []):
|
||||
handle.dispose()
|
||||
except Exception as exc:
|
||||
owned = [r for r in self._registration_order if r.plugin_key == plugin_key]
|
||||
self._dispose_registrations(owned)
|
||||
|
||||
@@ -67,7 +67,13 @@ def _is_bundled(provider_dir: Path) -> bool:
|
||||
|
||||
def _module_name(provider_dir: Path, name: str) -> str:
|
||||
"""``plugins.memory.<name>`` for bundled providers, else under the synthetic user namespace."""
|
||||
return f"plugins.memory.{name}" if _is_bundled(provider_dir) else f"{_USER_NAMESPACE}.{name}"
|
||||
if _is_bundled(provider_dir):
|
||||
return f"plugins.memory.{name}"
|
||||
# Separate package trees, including relative imports, across homes/sources.
|
||||
import hashlib
|
||||
|
||||
digest = hashlib.sha256(str(provider_dir.resolve()).encode()).hexdigest()[:16]
|
||||
return f"{_USER_NAMESPACE}.{name}__source_{digest}"
|
||||
|
||||
|
||||
def _external_source_dirs() -> List[Path]:
|
||||
@@ -224,7 +230,7 @@ def _load_provider_from_entry_point(entry_point, *, register_skills: bool = True
|
||||
pass
|
||||
if hasattr(loaded, "register"):
|
||||
collector = _ProviderCollector(entry_point.name, register_skills=register_skills)
|
||||
loaded.register(collector)
|
||||
collector.collect(loaded.register, source=getattr(loaded, "__file__", None))
|
||||
if collector.provider:
|
||||
return collector.provider
|
||||
if callable(loaded):
|
||||
@@ -235,7 +241,7 @@ def _load_provider_from_entry_point(entry_point, *, register_skills: bool = True
|
||||
except TypeError:
|
||||
pass
|
||||
collector = _ProviderCollector(entry_point.name, register_skills=register_skills)
|
||||
loaded(collector)
|
||||
collector.collect(loaded)
|
||||
return collector.provider
|
||||
|
||||
provider = _instantiate_subclass(loaded)
|
||||
@@ -259,7 +265,7 @@ def _load_provider_from_dir(provider_dir: Path, *, register_skills: bool = True)
|
||||
if hasattr(mod, "register"):
|
||||
collector = _ProviderCollector(name, register_skills=register_skills)
|
||||
try:
|
||||
mod.register(collector)
|
||||
collector.collect(mod.register, source=mod.__file__)
|
||||
except Exception as e:
|
||||
# A raise AFTER register_memory_provider() must not cost us the provider:
|
||||
# falling through to the subclass scan would hand back a bare second
|
||||
@@ -288,6 +294,43 @@ class _ProviderCollector:
|
||||
self.provider = None
|
||||
self._register_skills = register_skills
|
||||
self._context = None
|
||||
self._hook_source = None
|
||||
|
||||
def collect(self, register, *, source=None):
|
||||
"""General discovery owns hooks; memory-only activation supplies a fallback.
|
||||
|
||||
Replace a whole registration group, not callbacks by name or identity:
|
||||
a register function may deliberately create several distinct closures.
|
||||
"""
|
||||
module = sys.modules.get(getattr(register, "__module__", ""))
|
||||
source = source or getattr(module, "__file__", None)
|
||||
if source:
|
||||
self._hook_source = (self.name, str(Path(source).resolve()))
|
||||
manager = self._plugin_context()._manager
|
||||
with manager._discovery_lock:
|
||||
if self._hook_source is not None:
|
||||
for handle in manager._memory_hook_registrations.pop(self._hook_source, []):
|
||||
handle.dispose()
|
||||
register(self)
|
||||
|
||||
def register_hook(self, hook_name, callback):
|
||||
context = self._plugin_context()
|
||||
manager = context._manager
|
||||
with manager._discovery_lock:
|
||||
if self._hook_source is not None:
|
||||
for loaded in manager._plugins.values():
|
||||
source = getattr(loaded.module, "__file__", None)
|
||||
if (loaded.enabled and source and
|
||||
(loaded.manifest.name, str(Path(source).resolve())) == self._hook_source):
|
||||
# Preserve the disposable-handle contract without leasing
|
||||
# the general plugin's callback to this provider instance.
|
||||
handle = context._track("hook", hook_name, lambda: None)
|
||||
manager._memory_hook_registrations.setdefault(self._hook_source, []).append(handle)
|
||||
return handle
|
||||
handle = context.register_hook(hook_name, callback)
|
||||
if self._hook_source is not None:
|
||||
manager._memory_hook_registrations.setdefault(self._hook_source, []).append(handle)
|
||||
return handle
|
||||
|
||||
def register_memory_provider(self, provider):
|
||||
self.provider = provider
|
||||
@@ -386,7 +429,7 @@ def discover_plugin_cli_commands() -> List[dict]:
|
||||
# resolve without executing the plugin's __init__.py (the shell has no
|
||||
# __file__, so _load_provider_from_dir() still loads the real module).
|
||||
_register_synthetic_package(_USER_NAMESPACE, [])
|
||||
_register_synthetic_package(f"{_USER_NAMESPACE}.{active_provider}", [str(plugin_dir)])
|
||||
_register_synthetic_package(_module_name(plugin_dir, active_provider), [str(plugin_dir)])
|
||||
spec = importlib.util.spec_from_file_location(module_name, str(plugin_dir / "cli.py"))
|
||||
if not spec or not spec.loader:
|
||||
return []
|
||||
|
||||
184
tests/plugins/test_memory_hook_registration.py
Normal file
184
tests/plugins/test_memory_hook_registration.py
Normal file
@@ -0,0 +1,184 @@
|
||||
"""Exercise dual-kind plugins through the real general and memory loaders."""
|
||||
|
||||
import textwrap
|
||||
|
||||
import pytest
|
||||
|
||||
from hermes_cli.plugins import get_plugin_manager
|
||||
from plugins.memory import load_memory_provider
|
||||
|
||||
|
||||
def _install(home, monkeypatch, *, label="first", enabled=True):
|
||||
home.mkdir(parents=True, exist_ok=True)
|
||||
monkeypatch.setenv("HERMES_HOME", str(home))
|
||||
monkeypatch.setenv("HERMES_BUNDLED_PLUGINS", str(home / "empty"))
|
||||
monkeypatch.delenv("HERMES_ENABLE_PROJECT_PLUGINS", raising=False)
|
||||
monkeypatch.chdir(home)
|
||||
(home / "config.yaml").write_text(
|
||||
f"plugins:\n enabled: {'[dual]' if enabled else '[]'}\nmemory:\n provider: dual\n"
|
||||
)
|
||||
plugin = home / "plugins" / "dual"
|
||||
plugin.mkdir(parents=True)
|
||||
(plugin / "plugin.yaml").write_text("name: dual\nversion: 1.0.0\nkind: standalone\n")
|
||||
(plugin / "values.py").write_text(f"LABEL = {label!r}\n")
|
||||
(plugin / "__init__.py").write_text(textwrap.dedent('''\
|
||||
from agent.memory_provider import MemoryProvider
|
||||
from .values import LABEL
|
||||
|
||||
class Provider(MemoryProvider):
|
||||
name = "dual"
|
||||
def is_available(self): return True
|
||||
def initialize(self, session_id, **kwargs): pass
|
||||
def get_tool_schemas(self): return []
|
||||
|
||||
def make_hook(label):
|
||||
def callback(**kwargs):
|
||||
return {"context": label}
|
||||
return callback
|
||||
|
||||
def register(ctx):
|
||||
ctx.register_memory_provider(Provider())
|
||||
ctx.register_hook("pre_llm_call", make_hook(LABEL))
|
||||
ctx.register_hook("pre_llm_call", make_hook("second"))
|
||||
'''))
|
||||
return get_plugin_manager()
|
||||
|
||||
|
||||
def _contexts(manager):
|
||||
return manager.invoke_hook("pre_llm_call", session_id="")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("order", ["plugin-first", "memory-first", "memory-only"])
|
||||
def test_dual_kind_plugin_hooks_run_once(tmp_path, monkeypatch, order):
|
||||
manager = _install(tmp_path, monkeypatch, enabled=order != "memory-only")
|
||||
try:
|
||||
if order == "plugin-first":
|
||||
manager.discover_and_load()
|
||||
provider = load_memory_provider("dual")
|
||||
assert provider is not None and provider.name == "dual"
|
||||
if order == "memory-first":
|
||||
manager.discover_and_load()
|
||||
expected = [{"context": "first"}, {"context": "second"}]
|
||||
assert _contexts(manager) == expected
|
||||
# New provider instances must not append another fallback hook group.
|
||||
assert load_memory_provider("dual") is not provider
|
||||
assert _contexts(manager) == expected
|
||||
manager.unload()
|
||||
assert _contexts(manager) == []
|
||||
assert not manager._memory_hook_registrations
|
||||
if order != "memory-only":
|
||||
manager.discover_and_load(force=True)
|
||||
assert load_memory_provider("dual") is not None
|
||||
assert _contexts(manager) == expected
|
||||
finally:
|
||||
manager.unload()
|
||||
|
||||
|
||||
def test_equal_names_in_different_homes_keep_their_own_imports(tmp_path, monkeypatch):
|
||||
managers = []
|
||||
try:
|
||||
for label in ("alpha", "beta"):
|
||||
manager = _install(tmp_path / label, monkeypatch, label=label)
|
||||
managers.append(manager)
|
||||
assert load_memory_provider("dual") is not None
|
||||
manager.discover_and_load()
|
||||
assert _contexts(manager) == [{"context": label}, {"context": "second"}]
|
||||
assert _contexts(managers[0]) == [{"context": "alpha"}, {"context": "second"}]
|
||||
managers[1].unload()
|
||||
assert _contexts(managers[0]) == [{"context": "alpha"}, {"context": "second"}]
|
||||
finally:
|
||||
for manager in managers:
|
||||
manager.unload()
|
||||
|
||||
|
||||
def test_failed_general_registration_leaves_no_callable(tmp_path, monkeypatch):
|
||||
manager = _install(tmp_path, monkeypatch)
|
||||
try:
|
||||
assert load_memory_provider("dual") is not None
|
||||
source = tmp_path / "plugins" / "dual" / "__init__.py"
|
||||
source.write_text(source.read_text() + '\n raise RuntimeError("broken registration")\n')
|
||||
manager.discover_and_load()
|
||||
assert _contexts(manager) == []
|
||||
assert not manager._plugins["dual"].enabled
|
||||
finally:
|
||||
manager.unload()
|
||||
|
||||
|
||||
def test_same_name_different_sources_are_not_suppressed(tmp_path, monkeypatch):
|
||||
import shutil
|
||||
|
||||
home = tmp_path / "home"
|
||||
manager = _install(home, monkeypatch)
|
||||
project = tmp_path / "project"
|
||||
source = project / ".hermes" / "plugins" / "dual"
|
||||
shutil.copytree(home / "plugins" / "dual", source)
|
||||
(source / "values.py").write_text('LABEL = "project"\n')
|
||||
monkeypatch.chdir(project)
|
||||
monkeypatch.setenv("HERMES_ENABLE_PROJECT_PLUGINS", "1")
|
||||
try:
|
||||
assert load_memory_provider("dual") is not None
|
||||
manager.discover_and_load()
|
||||
assert _contexts(manager) == [
|
||||
{"context": "first"}, {"context": "second"},
|
||||
{"context": "project"}, {"context": "second"},
|
||||
]
|
||||
finally:
|
||||
manager.unload()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("entry_kind", ["module", "function"])
|
||||
def test_entry_point_registration_uses_same_hook_ownership(tmp_path, monkeypatch, entry_kind):
|
||||
from types import SimpleNamespace
|
||||
from plugins.memory import _load_provider_from_entry_point
|
||||
|
||||
manager = _install(tmp_path, monkeypatch)
|
||||
try:
|
||||
manager.discover_and_load()
|
||||
module = manager._plugins["dual"].module
|
||||
target = module if entry_kind == "module" else module.register
|
||||
entry = SimpleNamespace(name="dual", load=lambda: target)
|
||||
assert _load_provider_from_entry_point(entry) is not None
|
||||
assert _contexts(manager) == [{"context": "first"}, {"context": "second"}]
|
||||
finally:
|
||||
manager.unload()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("order", ["plugin-first", "memory-first"])
|
||||
def test_reexported_register_uses_plugin_source(tmp_path, monkeypatch, order):
|
||||
manager = _install(tmp_path, monkeypatch)
|
||||
plugin = tmp_path / "plugins" / "dual"
|
||||
original = plugin / "__init__.py"
|
||||
(plugin / "implementation.py").write_text(original.read_text())
|
||||
original.write_text("from .implementation import register, Provider # MemoryProvider\n")
|
||||
try:
|
||||
if order == "plugin-first":
|
||||
manager.discover_and_load()
|
||||
assert load_memory_provider("dual") is not None
|
||||
if order == "memory-first":
|
||||
manager.discover_and_load()
|
||||
assert _contexts(manager) == [{"context": "first"}, {"context": "second"}]
|
||||
finally:
|
||||
manager.unload()
|
||||
|
||||
|
||||
def test_suppressed_hook_still_returns_disposable_handle(tmp_path, monkeypatch):
|
||||
manager = _install(tmp_path, monkeypatch)
|
||||
source = tmp_path / "plugins" / "dual" / "__init__.py"
|
||||
source.write_text(source.read_text().split("def register(ctx):")[0] + textwrap.dedent('''\
|
||||
def register(ctx):
|
||||
handle = ctx.register_hook("pre_llm_call", make_hook("discard"))
|
||||
assert handle.active
|
||||
handle.dispose()
|
||||
assert not handle.active
|
||||
provider = Provider()
|
||||
provider.configured = True
|
||||
ctx.register_memory_provider(provider)
|
||||
ctx.register_hook("pre_llm_call", make_hook("kept"))
|
||||
'''))
|
||||
try:
|
||||
manager.discover_and_load()
|
||||
provider = load_memory_provider("dual")
|
||||
assert provider is not None and provider.configured
|
||||
assert _contexts(manager) == [{"context": "kept"}]
|
||||
finally:
|
||||
manager.unload()
|
||||
@@ -1511,6 +1511,8 @@ def register(ctx):
|
||||
|
||||
Memory providers are single-select — only one is active at a time, chosen via `memory.provider` in `config.yaml`.
|
||||
|
||||
If a provider also loads as a general plugin, general discovery owns its lifecycle hooks. The memory loader supplies hooks only as a fallback until that same plugin source loads successfully through general discovery. Repeated provider loads replace the fallback hook group; distinct callbacks within the group are preserved. This does not deduplicate hooks from different plugin sources or change provider activation.
|
||||
|
||||
**Full guide:** [Memory Provider Plugins](/developer-guide/memory-provider-plugin) — full `MemoryProvider` ABC, threading contract, profile isolation, CLI command registration via `cli.py`.
|
||||
|
||||
### Context engine plugins — replace the context compressor
|
||||
|
||||
Reference in New Issue
Block a user