diff --git a/agent/agent_init.py b/agent/agent_init.py index 30bb6d8370533..400456050cf81 100644 --- a/agent/agent_init.py +++ b/agent/agent_init.py @@ -1045,6 +1045,10 @@ def init_agent( # Track conversation messages for session logging agent._session_messages: List[Dict[str, Any]] = [] + # Plugin context engines can opt into per-agent cloning. When a clone is + # created, this agent owns its lifecycle and may close it during teardown; + # shared registered engines must stay alive for other cached agents. + agent._owns_context_engine = False # Responses encrypted reasoning replay state. Some OpenAI-compatible # routes accept GPT-5 Responses requests but later reject replayed # encrypted reasoning blobs (HTTP 400 ``invalid_encrypted_content``). @@ -1501,7 +1505,24 @@ def init_agent( # else: config says "compressor" — use built-in, don't auto-activate plugins if _selected_engine is not None: - agent.context_compressor = _selected_engine + _registered_engine = _selected_engine + try: + _agent_engine = _registered_engine.clone_for_agent() + except Exception as _ce_clone_err: + _ra().logger.warning( + "Context engine '%s' failed to clone for agent — using registered instance: %s", + getattr(_registered_engine, "name", _engine_name), + _ce_clone_err, + ) + _agent_engine = _registered_engine + if _agent_engine is None: + _ra().logger.warning( + "Context engine '%s' returned None from clone_for_agent — using registered instance", + getattr(_registered_engine, "name", _engine_name), + ) + _agent_engine = _registered_engine + agent._owns_context_engine = _agent_engine is not _registered_engine + agent.context_compressor = _agent_engine # Resolve context_length for plugin engines — mirrors switch_model() path from agent.model_metadata import get_model_context_length _plugin_ctx_len = get_model_context_length( diff --git a/agent/context_engine.py b/agent/context_engine.py index 79c31fb48e6cd..e971ec6146174 100644 --- a/agent/context_engine.py +++ b/agent/context_engine.py @@ -17,11 +17,13 @@ Lifecycle: 1. Engine is instantiated and registered (plugin register() or default) - 2. on_session_start() called when a conversation begins - 3. update_from_response() called after each API response with usage data - 4. should_compress() checked after each turn - 5. compress() called when should_compress() returns True - 6. on_session_end() called at real session boundaries (CLI exit, /reset, + 2. If the engine implements per-agent cloning, the host clones it for each + AIAgent before binding session state + 3. on_session_start() called when a conversation begins + 4. update_from_response() called after each API response with usage data + 5. should_compress() checked after each turn + 6. compress() called when should_compress() returns True + 7. on_session_end() called at real session boundaries (CLI exit, /reset, gateway session expiry) — NOT per-turn """ @@ -105,6 +107,21 @@ def compress( don't support it may simply ignore this argument. """ + # -- Optional: agent instance lifecycle -------------------------------- + + def clone_for_agent(self) -> "ContextEngine": + """Return the context-engine instance an AIAgent should own. + + Plugin registration is process-wide, but gateway runtimes may keep + multiple cached AIAgent instances alive at the same time (different + platforms, chats, cron jobs, etc.). Engines that keep mutable session + binding or cursor state on ``self`` should override this method and + return a fresh engine instance that shares durable storage/configuration + as needed. Stateless engines can use the default, which preserves the + historical shared-instance behavior. + """ + return self + # -- Optional: pre-flight check ---------------------------------------- def should_compress_preflight(self, messages: List[Dict[str, Any]]) -> bool: diff --git a/run_agent.py b/run_agent.py index 9c720bcbfe091..cb6390c686622 100644 --- a/run_agent.py +++ b/run_agent.py @@ -3089,7 +3089,20 @@ def close(self) -> None: except Exception: pass - # 6. Free conversation history. Mirrors _release_evicted_agent_soft's + # 6. Close context-engine resources only when this agent owns a cloned + # engine instance. Plugin-registered singleton engines are process-wide + # and may still be used by other cached agents. + try: + if getattr(self, "_owns_context_engine", False): + compressor = getattr(self, "context_compressor", None) + shutdown = getattr(compressor, "shutdown", None) + if callable(shutdown): + shutdown() + self._owns_context_engine = False + except Exception: + pass + + # 7. Free conversation history. Mirrors _release_evicted_agent_soft's # soft-eviction clear — close() is the hard teardown for true session # boundaries (/new, /reset, session expiry), so the message list won't # be reused. Drops the reference proactively rather than waiting for diff --git a/tests/agent/test_context_engine_host_contract.py b/tests/agent/test_context_engine_host_contract.py index bb6fb4c410892..fd1165a4b9c7a 100644 --- a/tests/agent/test_context_engine_host_contract.py +++ b/tests/agent/test_context_engine_host_contract.py @@ -28,10 +28,83 @@ from unittest.mock import MagicMock - +from agent.context_engine import ContextEngine from run_agent import AIAgent +class _ContractEngine(ContextEngine): + def __init__(self, name: str = "contract-engine"): + self._name = name + self.clones: list[_ContractEngine] = [] + self.update_model_calls: list[dict[str, object]] = [] + self.session_start_calls: list[tuple[str, dict[str, object]]] = [] + self.shutdown_count = 0 + + @property + def name(self) -> str: + return self._name + + def update_from_response(self, usage): + pass + + def should_compress(self, prompt_tokens=None): + return False + + def compress(self, messages, current_tokens=None, focus_topic=None): + return messages + + def update_model( + self, + model: str, + context_length: int, + base_url: str = "", + api_key: str = "", + provider: str = "", + api_mode: str = "", + ): + self.update_model_calls.append({ + "model": model, + "context_length": context_length, + "base_url": base_url, + "api_key": api_key, + "provider": provider, + "api_mode": api_mode, + }) + self.context_length = context_length + + def on_session_start(self, session_id: str, **kwargs) -> None: + self.session_start_calls.append((session_id, kwargs)) + + def shutdown(self) -> None: + self.shutdown_count += 1 + + +class _CloningContractEngine(_ContractEngine): + def clone_for_agent(self) -> ContextEngine: + clone = _ContractEngine(self.name) + self.clones.append(clone) + return clone + + +def _patch_agent_init_for_plugin_engine(monkeypatch, engine: ContextEngine) -> None: + monkeypatch.setattr( + "hermes_cli.config.load_config", + lambda: { + "context": {"engine": engine.name}, + "model": {"context_length": 200_000}, + }, + ) + monkeypatch.setattr( + "hermes_cli.plugins.get_plugin_context_engine", + lambda: engine, + ) + monkeypatch.setattr( + "agent.model_metadata.get_model_context_length", + lambda *args, **kwargs: 200_000, + ) + monkeypatch.setattr("run_agent.OpenAI", MagicMock(return_value=MagicMock())) + + def _bare_agent() -> AIAgent: agent = object.__new__(AIAgent) agent.session_id = "test-session" @@ -287,3 +360,65 @@ def test_engine_collector_rejects_builtin_command_conflicts(): # Must NOT have overwritten / registered against built-in /help. assert "help" not in manager._plugin_commands or \ manager._plugin_commands["help"].get("plugin") != "context-engine:my-lcm" + + +def test_agent_init_clones_plugin_context_engine_per_agent(monkeypatch): + """Mutable plugin engines can provide isolated per-AIAgent runtime state.""" + registered = _CloningContractEngine() + _patch_agent_init_for_plugin_engine(monkeypatch, registered) + + agent = AIAgent( + api_key="test-key", + base_url="https://example.invalid/v1", + provider="openai", + model="test-model", + quiet_mode=True, + skip_context_files=True, + skip_memory=True, + enabled_toolsets=[], + session_id="agent-s1", + platform="telegram", + gateway_session_key="agent:main:telegram:dm:42", + ) + clone = registered.clones[0] + try: + assert getattr(agent, "context_compressor") is clone + assert getattr(agent, "_owns_context_engine") is True + assert registered.update_model_calls == [] + assert registered.session_start_calls == [] + assert clone.update_model_calls[0]["context_length"] == 200_000 + assert clone.session_start_calls[0][0] == "agent-s1" + assert clone.session_start_calls[0][1]["conversation_id"] == "agent:main:telegram:dm:42" + finally: + agent.close() + + assert clone.shutdown_count == 1 + assert registered.shutdown_count == 0 + assert getattr(agent, "_owns_context_engine") is False + + +def test_agent_close_does_not_shutdown_shared_plugin_context_engine(monkeypatch): + """The default clone_for_agent() keeps backward-compatible shared engines alive.""" + shared = _ContractEngine() + _patch_agent_init_for_plugin_engine(monkeypatch, shared) + + agent = AIAgent( + api_key="test-key", + base_url="https://example.invalid/v1", + provider="openai", + model="test-model", + quiet_mode=True, + skip_context_files=True, + skip_memory=True, + enabled_toolsets=[], + session_id="agent-s2", + platform="weixin", + gateway_session_key="agent:main:weixin:dm:77", + ) + try: + assert getattr(agent, "context_compressor") is shared + assert getattr(agent, "_owns_context_engine") is False + finally: + agent.close() + + assert shared.shutdown_count == 0