diff --git a/hermes_cli/middleware.py b/hermes_cli/middleware.py index 891d07e26f..27e15e54a0 100644 --- a/hermes_cli/middleware.py +++ b/hermes_cli/middleware.py @@ -7,7 +7,7 @@ wrapping the actual execution callback. Agent-loop call sites and plugins share from __future__ import annotations import logging -from copy import deepcopy +from copy import copy, deepcopy from dataclasses import dataclass, field from typing import Any, Callable, Dict, List @@ -48,58 +48,107 @@ def middleware_payload(**kwargs: Any) -> Dict[str, Any]: def _safe_copy(payload: Any) -> Any: - """Deep-copy a request payload, tolerating non-deepcopyable members. + """Copy a request graph, tolerating non-deepcopyable members. An LLM request can carry clients/callbacks/file handles; a hard ``deepcopy`` failure would - otherwise abort the whole request-middleware pass. + otherwise abort the whole request-middleware pass. The fallback isolates surrounding + containers and preserves their aliases and cycles; opaque leaves retain their identity. """ try: return deepcopy(payload) except Exception as exc: logger.debug("deepcopy failed for request payload (%s); retaining opaque leaves", exc) - # A client or file handle must not make unrelated request lists/dicts shared - # between callbacks. Preserve opaque leaves, while still copying containers. + # Walk containers without Python recursion: an opaque leaf or a deeply nested + # request must not share surrounding containers or bypass the middleware pass. Shallow + # copies retain container subclasses (including defaultdict's factory). memo: Dict[int, Any] = {} + pending: Dict[int, List[Callable[[Any], None]]] = {} + tasks: List[Callable[[], None]] = [] + root: List[Any] = [] - def copy_value(value: Any) -> Any: + def copy_state(value: Any, result: Any) -> None: + state = getattr(value, "__dict__", None) + if state is not None and hasattr(result, "__dict__"): + tasks.append(lambda: copy_value(state, lambda copied: setattr(result, "__dict__", copied))) + + def copy_value(value: Any, assign: Callable[[Any], None]) -> None: if type(value) in (str, bytes, int, float, bool, complex, type(None)): - return value + assign(value) + return identity = id(value) if identity in memo: - return memo[identity] + assign(memo[identity]) + return + if identity in pending: + pending[identity].append(assign) + return result: Any - if type(value) is dict: - result = {} - memo[identity] = result - for key, item in value.items(): - result[copy_value(key)] = copy_value(item) - elif type(value) is list: - result = [] - memo[identity] = result - result.extend(copy_value(item) for item in value) - elif type(value) is set: - result = set() - memo[identity] = result - result.update(copy_value(item) for item in value) - elif type(value) is tuple: - items = [copy_value(item) for item in value] - # Recursion through a mutable child may already have copied this tuple. - result = memo.get(identity, tuple(items)) - else: - # A failing __deepcopy__ may leave partial containers in its memo. - # Publish that memo only on success, never after an opaque leaf fails. - leaf_memo = memo.copy() + if isinstance(value, (dict, list, set)): + factory = dict if isinstance(value, dict) else list if isinstance(value, list) else set try: - result = deepcopy(value, leaf_memo) + result = copy(value) + if result is value or not isinstance(result, factory): + result = factory() + result.clear() except Exception: - result = value + result = factory() + memo[identity] = result + assign(result) + copy_state(value, result) + if isinstance(result, dict): + for key, item in reversed(list(value.items())): + pair: List[Any] = [None, None] + tasks.append(lambda pair=pair, target=result: target.__setitem__(pair[0], pair[1])) + tasks.append(lambda item=item, pair=pair: copy_value( + item, lambda copied: pair.__setitem__(1, copied))) + tasks.append(lambda key=key, pair=pair: copy_value( + key, lambda copied: pair.__setitem__(0, copied))) + elif isinstance(result, list): + list.extend(result, [None] * len(value)) + for index, item in reversed(list(enumerate(value))): + tasks.append(lambda index=index, item=item, target=result: copy_value( + item, lambda copied: list.__setitem__(target, index, copied))) else: - memo.update(leaf_memo) - memo[identity] = result - return result + for item in value: + tasks.append(lambda item=item, target=result: copy_value( + item, lambda copied: set.add(target, copied))) + return + if isinstance(value, (tuple, frozenset)): + pending[identity] = [assign] + items: List[Any] = [None] * len(value) - return copy_value(payload) + def finish() -> None: + factory = tuple if isinstance(value, tuple) else frozenset + result = factory.__new__(type(value), items) + memo[identity] = result + for setter in pending.pop(identity): + setter(result) + copy_state(value, result) + + # Mutable children are allocated before their contents, so tuple/list + # cycles can be completed after the tuple itself becomes available. + tasks.append(finish) + for index, item in reversed(list(enumerate(value))): + tasks.append(lambda index=index, item=item: copy_value( + item, lambda copied: items.__setitem__(index, copied))) + return + # A failing __deepcopy__ may leave partial containers in its memo. + # Publish that memo only on success, never after an opaque leaf fails. + leaf_memo = memo.copy() + try: + result = deepcopy(value, leaf_memo) + except Exception: + result = value + else: + memo.update(leaf_memo) + memo[identity] = result + assign(result) + + tasks.append(lambda: copy_value(payload, root.append)) + while tasks: + tasks.pop()() + return root[0] def _apply_request_chain( @@ -209,6 +258,11 @@ def _run_execution_chain(kind: str, terminal_call: Callable[[Any], Any], **kwarg if not callbacks: return terminal_call(kwargs[payload_key]) + # Execution-only plugins may rewrite the caller's payload in place, so freeze + # its original before entering the first frame, then isolate each frame's view. + original_key = "original_" + payload_key + kwargs[original_key] = _safe_copy(kwargs[original_key]) + def call_at(index: int, payload: Any) -> Any: if index >= len(callbacks): return terminal_call(payload) @@ -238,6 +292,7 @@ def _run_execution_chain(kind: str, terminal_call: Callable[[Any], Any], **kwarg call_kwargs = middleware_payload(**kwargs) call_kwargs[payload_key] = payload + call_kwargs[original_key] = _safe_copy(kwargs[original_key]) call_kwargs["next_call"] = next_call try: return callback(**call_kwargs) diff --git a/hermes_cli/plugins_dispatch.py b/hermes_cli/plugins_dispatch.py index b7a74ba461..84634a59cb 100644 --- a/hermes_cli/plugins_dispatch.py +++ b/hermes_cli/plugins_dispatch.py @@ -593,18 +593,22 @@ class PluginDispatchMixin: ) -> List[Any]: """Dispatch middleware; request rewrites compose only after a valid return. - Execution middleware remains a flat fan-out. Request callbacks receive a - copy of the last accepted payload, so a failing/observer callback cannot - mutate an earlier accepted rewrite or leak changes into later callbacks. + Request callbacks receive a copy of the last accepted payload, so a + failing/observer callback cannot mutate an earlier accepted rewrite. + Each callback also receives its own pre-middleware snapshot. """ from hermes_cli.middleware import _safe_copy + copy_keys = {"original_request", "original_args"} + if _payload_key is not None: + copy_keys.add(_payload_key) results: List[Any] = [] for cb in self._middleware.get(kind, []): try: - callback_kwargs = kwargs - if _payload_key is not None: - callback_kwargs = {**kwargs, _payload_key: _safe_copy(kwargs[_payload_key])} + callback_kwargs = { + key: _safe_copy(value) if key in copy_keys else value + for key, value in kwargs.items() + } ret = cb(**callback_kwargs) if ret is not None: if (_payload_key is not None and isinstance(ret, dict) diff --git a/tests/hermes_cli/test_middleware_isolation.py b/tests/hermes_cli/test_middleware_isolation.py new file mode 100644 index 0000000000..f7bf082459 --- /dev/null +++ b/tests/hermes_cli/test_middleware_isolation.py @@ -0,0 +1,264 @@ +"""Request rewrites and original snapshots stay isolated across discovered middleware.""" + +from collections import OrderedDict, defaultdict +from typing import Any, NamedTuple, cast + +import pytest + +from hermes_cli import plugins +from hermes_cli.config import atomic_config_write +from hermes_cli.middleware import ( + apply_llm_request_middleware, + apply_tool_request_middleware, + run_llm_execution_middleware, + run_tool_execution_middleware, +) + + +class RequestDict(dict): + pass + + +class RequestList(list): + pass + + +class RequestTuple(NamedTuple): + data: Any + file: Any + + +class OpaqueMember: + def __init__(self, related): + self.related = related + + def __deepcopy__(self, memo): + memo[id(self.related)] = {"partial": True} + raise TypeError("opaque member cannot be copied") + + +def _load_plugin(tmp_path, monkeypatch, source): + home = tmp_path / "home" + plugin = home / "plugins" / "copy_probe" + plugin.mkdir(parents=True) + monkeypatch.setenv("HERMES_HOME", str(home)) + plugin.joinpath("plugin.yaml").write_text( + "name: copy_probe\nversion: 0.1.0\ndescription: Middleware isolation probe\n", + encoding="utf-8", + ) + plugin.joinpath("__init__.py").write_text(source, encoding="utf-8") + atomic_config_write(home / "config.yaml", {"plugins": {"enabled": ["copy_probe"]}}) + manager = plugins.PluginManager() + manager.discover_and_load() + monkeypatch.setattr(plugins, "get_plugin_manager", lambda: manager) + assert any(manager.has_middleware(kind) for kind in ( + "llm_request", "tool_request", "llm_execution", "tool_execution", + )) + + +def _apply(kind, payload): + if kind == "llm": + return apply_llm_request_middleware(payload) + return apply_tool_request_middleware("probe", payload) + + +@pytest.mark.parametrize("kind", ["llm", "tool"]) +@pytest.mark.parametrize("stage", ["request", "execution", "execution_only"]) +@pytest.mark.parametrize("outcome", ["observe", "raise", "rewrite"]) +@pytest.mark.parametrize("opaque", [False, True]) +def test_original_snapshot_survives_each_middleware_callback( + tmp_path, monkeypatch, kind, stage, outcome, opaque, +): + key = "request" if kind == "llm" else "args" + request_listener = stage != "execution_only" + if not request_listener: + stage = "execution" + _load_plugin(tmp_path, monkeypatch, f""" +KIND = {kind!r} +KEY = {key!r} +STAGE = {stage!r} +OUTCOME = {outcome!r} +REQUEST_LISTENER = {request_listener!r} + +def mutate(**kw): + original = kw['original_' + KEY] + original['input'].append('snapshot-write') + original['injected'] = True + if not REQUEST_LISTENER: + kw[KEY]['input'].append('execution-write') + if OUTCOME == 'raise': + raise RuntimeError('snapshot probe failure') + if OUTCOME == 'rewrite': + payload = {{**kw[KEY], 'rewritten': True}} + if STAGE == 'request': + return {{KEY: payload, 'source': 'mutate'}} + return kw['next_call'](payload) + if STAGE == 'execution': + return kw['next_call']() + +def verify(**kw): + original = kw['original_' + KEY] + stage = 'execution' if 'next_call' in kw else 'request' + payload = {{**kw[KEY], stage + '_snapshot': list(original['input'])}} + payload[stage + '_injected'] = 'injected' in original + if 'next_call' in kw: + return kw['next_call'](payload) + return {{KEY: payload, 'source': 'verify'}} + +def register(ctx): + ctx.register_middleware(KIND + '_' + STAGE, mutate) + if REQUEST_LISTENER: + ctx.register_middleware(KIND + '_request', verify) + ctx.register_middleware(KIND + '_execution', verify) +""") + with (tmp_path / "upload.txt").open("w+", encoding="utf-8") as handle: + original = {"input": ["original"], **({"file": handle} if opaque else {})} + result = _apply(kind, original) if request_listener else None + payload = result.payload if result else original + context = {"original_" + key: result.original_payload} if result else {} + calls = [] + + def terminal(payload): + calls.append(payload) + return payload + + if kind == "llm": + executed = run_llm_execution_middleware( + payload, terminal, **context, + ) + else: + executed = run_tool_execution_middleware( + "probe", payload, terminal, **context, + ) + + expected_input = ["original"] if request_listener else ["original", "execution-write"] + assert original["input"] == expected_input + assert "injected" not in original + if result: + assert result.original_payload["input"] == ["original"] + assert "injected" not in result.original_payload + for phase in ("request", "execution") if request_listener else ("execution",): + assert executed[phase + "_snapshot"] == ["original"] + assert executed[phase + "_injected"] is False + assert executed["input"] == expected_input + assert bool(executed.get("rewritten")) == (outcome == "rewrite") + assert calls == [executed] + expected_sources = ["verify"] + if stage == "request" and outcome == "rewrite": + expected_sources.insert(0, "mutate") + if result: + assert [entry["source"] for entry in result.trace] == expected_sources + if opaque: + assert executed["file"] is original["file"] is handle + if result: + assert result.original_payload["file"] is handle + assert not handle.closed + + +def _leaf(node): + while isinstance(node, dict) and "child" in node: + node = node["child"] + while isinstance(node, (list, tuple)): + node = node[0] + return node + + +@pytest.mark.parametrize("kind", ["llm", "tool"]) +@pytest.mark.parametrize("container", [ + dict, OrderedDict, defaultdict, RequestDict, list, RequestList, tuple, RequestTuple, +]) +@pytest.mark.parametrize("depth,wrapper", [ + (0, dict), (1500, dict), (6000, dict), (6000, list), (6000, tuple), +]) +def test_opaque_request_graph_preserves_isolation_aliases_and_cycles( + tmp_path, monkeypatch, kind, container, depth, wrapper, +): + key = "request" if kind == "llm" else "args" + _load_plugin(tmp_path, monkeypatch, f""" +KIND = {kind!r} +KEY = {key!r} + +def leaf(node): + while isinstance(node, dict) and 'child' in node: + node = node['child'] + while isinstance(node, (list, tuple)): + node = node[0] + return node + +def first(**kw): + return {{KEY: {{**kw[KEY], 'accepted': True}}, 'source': 'first'}} + +def broken(**kw): + leaf(kw[KEY]['node'])['steps'].append('broken') + raise RuntimeError('graph probe failure') + +def observer(**kw): + leaf(kw[KEY]['node'])['steps'].append('observer') + +def second(**kw): + leaf(kw[KEY]['node'])['steps'].append('second') + return {{KEY: kw[KEY], 'source': 'second'}} + +def register(ctx): + for callback in (first, broken, observer, second): + ctx.register_middleware(KIND + '_request', callback) +""") + with (tmp_path / "upload.txt").open("w+", encoding="utf-8") as handle: + leaf = {"steps": ["original"], "file": handle} + if issubclass(container, dict): + node = container(list, leaf) if container is defaultdict else container(leaf) + elif container is RequestTuple: + node = container(leaf, handle) + else: + node = container([leaf, handle]) + if container in (RequestDict, RequestList): + node.metadata = ["original"] + for _ in range(depth): + node = {"child": node} if wrapper is dict else wrapper([node]) + cycle = [] + pair = (cycle, handle) + cycle.append(pair) + opaque = OpaqueMember(node) + original: dict[str, Any] = { + "opaque": opaque, "node": node, "alias": node, "cycle": pair, + "cycle_alias": pair, "tags": {"original"}, + "frozen": frozenset({RequestTuple("original", handle)}), + } + try: + result = _apply(kind, original) + except RecursionError: + # Keep pytest from comparing deeply nested frame locals while locating + # recursion: that can hang its traceback formatter on this graph. + raise AssertionError("Deep requests must still run middleware") from None + + assert result.payload["accepted"] is True + assert _leaf(result.payload["node"])["steps"] == ["original", "second"] + assert _leaf(result.original_payload["node"])["steps"] == ["original"] + assert _leaf(original["node"])["steps"] == ["original"] + assert result.payload["node"] is result.payload["alias"] + assert result.original_payload["node"] is result.original_payload["alias"] + assert result.payload["cycle"][0][0] is result.payload["cycle"] + assert result.original_payload["cycle"][0][0] is result.original_payload["cycle"] + assert result.payload["cycle_alias"] is result.payload["cycle"] + assert result.original_payload["cycle_alias"] is result.original_payload["cycle"] + assert result.payload["opaque"] is result.original_payload["opaque"] is opaque + assert result.payload["frozen"] == result.original_payload["frozen"] == original["frozen"] + assert _leaf(result.payload["node"])["file"] is handle + assert not handle.closed + assert [entry["source"] for entry in result.trace] == ["first", "second"] + copied_node = result.payload["node"] + snapshot_node = result.original_payload["node"] + source_node = cast(Any, original["node"]) + for _ in range(depth): + index = "child" if wrapper is dict else 0 + copied_node = copied_node[index] + snapshot_node = snapshot_node[index] + source_node = source_node[index] + assert type(copied_node) is type(snapshot_node) is type(source_node) is container + if container is defaultdict: + assert copied_node.default_factory is snapshot_node.default_factory is list + if container in (RequestDict, RequestList): + copied_node.metadata.append("accepted-only") + assert source_node.metadata == snapshot_node.metadata == ["original"] + result.payload["tags"].add("accepted-only") + assert original["tags"] == result.original_payload["tags"] == {"original"}