diff --git a/EvoScientist/langgraph_dev/main_graph.py b/EvoScientist/langgraph_dev/main_graph.py index fc1bc41..7921621 100644 --- a/EvoScientist/langgraph_dev/main_graph.py +++ b/EvoScientist/langgraph_dev/main_graph.py @@ -5,8 +5,244 @@ in ``EvoScientist/EvoScientist.py`` so it doesn't construct on plain ``import EvoScientist``. ``langgraph dev`` 's symbol resolver inspects module attributes directly and doesn't trigger ``__getattr__``, so we re-export here to make it visible. + +Before re-export we upgrade the compiled graph's class in place to +``_EvoFilteredGraph``, which strips ``PrivateStateAttr``-marked fields +(currently just ``_quickjs_snapshot_payload``) from ``get_state`` / +``get_state_history`` responses. Upstream ``langchain_quickjs`` annotates +the field ``PrivateStateAttr = OmitFromSchema(input=True, output=True)``, +but LangGraph's ``_prepare_state_snapshot`` doesn't honor that on +checkpoint reads — every ``getState`` materializes the delta chain back +into a full ~1.4 MB blob, which the WebUI then downloads. The subclass +closes the gap without touching the middleware's write path, preserving +cross-turn REPL persistence as ``langchain-ai/deepagents#3064`` shipped it. """ -from EvoScientist.EvoScientist import EvoScientist_agent +from langgraph.graph.state import CompiledStateGraph +from langgraph.types import PregelTask, StateSnapshot + +from EvoScientist.EvoScientist import EvoScientist_agent as _agent + +_PRIVATE_STATE_FIELDS = frozenset({"_quickjs_snapshot_payload"}) + +# Sanity check on the LangGraph internals ``_strip_private`` scrubs. If any +# of these attributes disappear or get renamed in a future upstream bump, +# the assertion fires at import time — the deployment refuses to start, +# instead of silently degrading (the filter would ``.get()`` its way to a +# no-op and the private-field payload would come back on the wire without +# anyone noticing until a user reports slow thread switches again). +# +# Doesn't cover every internal we depend on — ``metadata["writes"]`` / +# ``metadata["counters_since_delta_snapshot"]`` dict keys aren't a canary +# target because ``dict.get`` already tolerates their absence. What we +# canary here is the ``NamedTuple`` field set: renames there would be the +# highest-impact silent regression. +_EXPECTED_SNAPSHOT_FIELDS = frozenset({"values", "metadata", "tasks"}) +_EXPECTED_TASK_FIELDS = frozenset({"result", "state"}) + +_missing_snap = _EXPECTED_SNAPSHOT_FIELDS - set(StateSnapshot._fields) +_missing_task = _EXPECTED_TASK_FIELDS - set(PregelTask._fields) +if _missing_snap or _missing_task: + raise RuntimeError( + "LangGraph state shape drifted from the version _strip_private was " + f"written against. Missing StateSnapshot fields: {_missing_snap or set()}. " + f"Missing PregelTask fields: {_missing_task or set()}. Review " + "_strip_private and re-verify against the current upstream shape " + "before removing this assertion." + ) + + +def _strip_private(snap): + """Strip ``PrivateStateAttr``-marked fields from a ``StateSnapshot``. + + Empirically verified against a live history response for a thread with + a single touched turn: the private field leaks on four surfaces — three + trivial, one heavy: + + * ``snap.values`` — the materialized channel state exposed as the main + payload. For DeltaChannels this is the delta chain replayed into full + bytes (~1.4 MB for the quickjs snapshot). ``get_state`` and every + history entry. + * ``snap.metadata['writes']`` — ``{node_name: {channel: value}}`` map of + the raw writes that produced each checkpoint. On the ``after_agent`` + step that first snapshots the REPL, ``value`` is the encoded write + record ``("snap", full_bytes)`` ≈ 1.4 MB. + * ``snap.tasks[*].result`` — the return dict of each completed + ``PregelTask``. ``after_agent`` returns + ``{"_quickjs_snapshot_payload": ("snap", bytes)}``; this dict becomes + the task's ``result`` field, which the API surfaces verbatim under + ``tasks[*].result`` (``langgraph_api.state:106``). This is the + dominant leak: 1.7 MB in the last history entry of any thread whose + most-recent-in-window checkpoint had a snapshot anchor. + * ``snap.metadata['counters_since_delta_snapshot']`` — DeltaChannel's + snapshot cadence bookkeeping ``{channel: [count, superstep]}``. Tiny + (~20 B) but exposes the private field name; strip for cleanliness. + * ``snap.tasks[*].state`` (nested ``StateSnapshot``) — populated when the + caller passes ``subgraphs=True``. Repeats all of the above surfaces + for each subgraph task, so recurse into it. Not exercised by the + current WebUI (which doesn't pass ``subgraphs=True`` on REST reads), + but SDK / curl / gRPC callers can. + """ + if snap is None: + return snap + values = {k: v for k, v in snap.values.items() if k not in _PRIVATE_STATE_FIELDS} + metadata = snap.metadata + if metadata: + new_metadata = metadata + if new_metadata.get("writes"): + scrubbed_writes = { + node: { + k: v for k, v in ch_writes.items() if k not in _PRIVATE_STATE_FIELDS + } + for node, ch_writes in new_metadata["writes"].items() + } + new_metadata = {**new_metadata, "writes": scrubbed_writes} + if new_metadata.get("counters_since_delta_snapshot"): + scrubbed_counters = { + k: v + for k, v in new_metadata["counters_since_delta_snapshot"].items() + if k not in _PRIVATE_STATE_FIELDS + } + new_metadata = { + **new_metadata, + "counters_since_delta_snapshot": scrubbed_counters, + } + metadata = new_metadata + tasks = snap.tasks + if tasks: + new_tasks = [] + changed = False + for t in tasks: + replace_kwargs: dict = {} + result = getattr(t, "result", None) + if isinstance(result, dict) and any( + k in result for k in _PRIVATE_STATE_FIELDS + ): + replace_kwargs["result"] = { + k: v for k, v in result.items() if k not in _PRIVATE_STATE_FIELDS + } + # ``t.state`` is a ``RunnableConfig | StateSnapshot | None`` per + # ``PregelTask``'s typing. When ``subgraphs=True`` on the caller, + # this holds the subgraph's fully-materialized ``StateSnapshot`` — + # which repeats the same four leak surfaces (``values``, + # ``metadata.writes``, ``metadata.counters_since_delta_snapshot``, + # ``tasks[*].result/state``). Recurse so the whole tree is clean. + nested_state = getattr(t, "state", None) + if isinstance(nested_state, StateSnapshot): + scrubbed_state = _strip_private(nested_state) + if scrubbed_state is not nested_state: + replace_kwargs["state"] = scrubbed_state + if replace_kwargs: + new_tasks.append(t._replace(**replace_kwargs)) + changed = True + else: + new_tasks.append(t) + if changed: + tasks = tuple(new_tasks) + return snap._replace(values=values, metadata=metadata, tasks=tasks) + + +class _EvoFilteredGraph(CompiledStateGraph): + """Filters ``PrivateStateAttr``-marked state fields from checkpoint reads. + + ``Pregel.copy`` uses ``self.__class__(**attrs)`` so this subclass + survives the ``graph_obj.copy(update=...)`` call in + ``langgraph_api.graph.get_graph`` that binds the checkpointer / store + before yielding to endpoint handlers. + + **Known gap — streaming paths.** The overrides only cover ``get_state`` + / ``get_state_history``. On this compiled graph, + ``self.output_channels`` correctly excludes ``_quickjs_snapshot_payload`` + (respects ``OmitFromSchema(output=True)``), but + ``self.stream_channels_asis`` includes it alongside other private + fields (``jump_to``, ``_summarization_event``) — the two lists are + built by ``langgraph.graph.state``'s graph builder and only the first + checks the output schema. So a client streaming with + ``stream_mode="values"`` or ``stream_mode="events"`` (which fall back + to ``stream_channels_asis`` when ``output_keys`` is ``None``) can pull + the anchor blob in per-run event data. Empirically the WebUI's + ``stream_mode=["updates"]`` path is clean, so this is transient per-run + rather than the persistent per-getState download this PR targets. + Filter here first; extend into the stream layer if a client relying on + ``values`` / ``events`` reports it. + """ + + async def aget_state(self, config, *, subgraphs=False): + return _strip_private(await super().aget_state(config, subgraphs=subgraphs)) + + def get_state(self, config, *, subgraphs=False): + return _strip_private(super().get_state(config, subgraphs=subgraphs)) + + async def aget_state_history(self, config, **kw): + async for snap in super().aget_state_history(config, **kw): + yield _strip_private(snap) + + def get_state_history(self, config, **kw): + for snap in super().get_state_history(config, **kw): + yield _strip_private(snap) + + +# In-place ``__class__`` swap: the subclass adds only methods (no new +# instance attributes) so the memory layout is identical and the swap is +# safe. Constructing a fresh ``_EvoFilteredGraph`` via ``.copy()`` would +# require reproducing the deep-agent build pipeline; the swap avoids that. +_agent.__class__ = _EvoFilteredGraph +EvoScientist_agent = _agent + + +def _apply_filter_to_all_registered_graphs() -> None: + """Extend the class swap to every graph registered in ``langgraph.json``. + + ``EvoScientist.py:_build_middleware_stack`` installs + ``create_code_interpreter_middleware`` unconditionally — it's not gated + on the ``for_async_subagent`` flag — so every subagent (sync ``task`` + dispatch and async ``start_async_task``) carries the QuickJS REPL and + can produce ``_quickjs_snapshot_payload`` writes on its own checkpoint + namespace. + + Async subagents get their own ``thread_id`` and their ``/threads/{id}/state`` + endpoint is served by their own compiled graph. Without swapping the + class on those graphs, the filter we applied to ``EvoScientist_agent`` + doesn't reach that endpoint and any real code_interpreter touch inside + a subagent leaks the anchor snapshot verbatim. + + Reads the graph registry straight from ``langgraph.json`` so a new + subagent added to the config picks up the swap automatically — no + hardcoded list to keep in sync. + + Idempotent (skips graphs already swapped) and safe on graphs that don't + use the middleware — ``_strip_private`` returns snapshots unchanged when + the private field is absent. Best-effort: if the config is unreadable + or an entry can't be resolved, the deployment still starts — only the + unresolvable subagents remain unfiltered. + """ + import json + from importlib import import_module + from pathlib import Path + + config_path = Path(__file__).parent / "langgraph.json" + try: + config = json.loads(config_path.read_text()) + except (OSError, json.JSONDecodeError): + return + + for path in config.get("graphs", {}).values(): + # Format: "module.dotted.path:attr_name" + if ":" not in path: + continue + module_path, attr = path.rsplit(":", 1) + try: + module = import_module(module_path) + except ImportError: + continue + graph = getattr(module, attr, None) + if isinstance(graph, CompiledStateGraph) and not isinstance( + graph, _EvoFilteredGraph + ): + graph.__class__ = _EvoFilteredGraph + + +_apply_filter_to_all_registered_graphs() + __all__ = ["EvoScientist_agent"] diff --git a/EvoScientist/middleware/code_interpreter.py b/EvoScientist/middleware/code_interpreter.py index 1f54c50..467cb77 100644 --- a/EvoScientist/middleware/code_interpreter.py +++ b/EvoScientist/middleware/code_interpreter.py @@ -45,7 +45,21 @@ _MEMORY_FIRST_INTERPRETER_PROMPT = ( class EvoCodeInterpreterMiddleware(CodeInterpreterMiddleware): - """Code interpreter middleware with EvoScientist's memory preflight hint.""" + """Code interpreter middleware with EvoScientist's memory preflight hint. + + ``after_agent`` / ``aafter_agent`` are intentionally NOT overridden. An + earlier "conditional snapshot" gate that skipped ``after_agent`` on turns + where ``code_interpreter`` wasn't called saved ~50 ms/turn of + ``create_snapshot()`` work, but also skipped the slot eviction upstream + performs in the same hook (``finally: self._registry.evict(thread_id)`` + in ``langchain_quickjs.middleware.CodeInterpreterMiddleware.after_agent``). + ``before_agent`` restores the REPL on every turn that follows a touched + one via ``self._registry.get(thread_id)`` (get-or-create), so skipping + eviction leaked one ``ThreadWorker`` + QuickJS Runtime per persistent + ``thread_id`` that ever went touched → quiet. The regression test + ``test_after_agent_evicts_slot_on_untouched_turn`` guards against + reintroducing the gate. + """ def _prepare_for_call(self, request: ModelRequest) -> str: return super()._prepare_for_call(request) + _MEMORY_FIRST_INTERPRETER_PROMPT diff --git a/tests/test_code_interpreter_middleware.py b/tests/test_code_interpreter_middleware.py index 8d5103f..8baa89d 100644 --- a/tests/test_code_interpreter_middleware.py +++ b/tests/test_code_interpreter_middleware.py @@ -1,4 +1,5 @@ -"""Regression tests for the code_interpreter PTC allowlist. +"""Regression tests for the code_interpreter PTC allowlist and the +``EvoCodeInterpreterMiddleware`` subclass shape. langchain-quickjs >=0.3 reserves the ``task`` sub-agent dispatch tool as the top-level REPL global and raises ``ValueError`` if ``task`` appears in the @@ -8,7 +9,10 @@ allowlist (``task()`` stays reachable as the REPL global, with responseSchema). from __future__ import annotations +from unittest.mock import MagicMock + import pytest +from langchain_core.messages import AIMessage, HumanMessage from EvoScientist.middleware.code_interpreter import ( _DEFAULT_PTC_ALLOWLIST, @@ -45,3 +49,290 @@ def test_filter_tools_for_ptc_accepts_default_allowlist(): def test_create_code_interpreter_middleware_builds(): assert create_code_interpreter_middleware() is not None + + +def test_middleware_uses_thread_mode(): + """Upstream ``mode="thread"`` (the default) preserves cross-turn REPL + state as ``langchain-ai/deepagents#3064`` shipped it. The wire-cost + bloat that motivated the earlier ``mode="turn"`` regression guard is + fixed at the API serialization layer (``EvoFilteredGraph`` in + ``EvoScientist/langgraph_dev/main_graph.py``), not by revoking the + persistence feature. + """ + mw = create_code_interpreter_middleware() + assert mw._mode == "thread" + + +def test_after_agent_evicts_slot_on_untouched_turn(): + """Regression guard against reintroducing a conditional-snapshot gate + that skips ``after_agent`` on untouched turns. + + Upstream ``after_agent`` in ``langchain_quickjs/middleware.py`` performs + two things: snapshot the REPL AND evict the slot (``finally: + self._registry.evict(thread_id)``). ``before_agent`` restores the REPL + on any turn that follows a touched one via ``self._registry.get`` — + which is get-or-create. So if ``after_agent`` returns early without + evicting, one ``ThreadWorker`` + QuickJS Runtime leaks per persistent + ``thread_id`` that ever went touched → quiet. + + Fix: don't override ``after_agent`` / ``aafter_agent`` at all — inherit + upstream's unconditional snapshot+evict behavior. This test creates a + slot the way ``before_agent`` would, calls ``after_agent`` with an + untouched-state input, and asserts the slot was evicted. + """ + mw = create_code_interpreter_middleware() + tid = mw._fallback_thread_id + + # Simulate the slot creation that ``before_agent`` performs when it sees + # a prior turn's snapshot payload in state. + mw._registry.get(tid) + assert len(mw._registry._slots) == 1 + + # Untouched-turn state: no ``code_interpreter`` tool call between the + # last ``HumanMessage`` and end. Under the earlier buggy gate this + # returned ``{}`` without evicting — leaking the slot created above. + untouched_state = { + "_quickjs_snapshot_payload": b"payload-from-prior-turn", + "messages": [ + HumanMessage(content="thanks"), + AIMessage(content="you're welcome"), + ], + } + mw.after_agent(untouched_state, runtime=None) + + assert len(mw._registry._slots) == 0, ( + "after_agent must evict the slot even on untouched turns, because " + "before_agent already restored a REPL that owns a ThreadWorker + " + "QuickJS Runtime. Skipping eviction leaks those resources." + ) + + +def test_evo_filtered_graph_strips_private_snapshot_field(): + """The ``StateSnapshot`` returned by ``EvoScientist_agent.get_state`` must + not contain ``_quickjs_snapshot_payload`` in either ``values`` (the + materialized channel payload) or ``metadata['writes']`` (the raw write + records surfaced by ``get_state_history``). + """ + from EvoScientist.langgraph_dev.main_graph import _EvoFilteredGraph, _strip_private + + snap = MagicMock() + snap.values = { + "messages": ["m1"], + "_quickjs_snapshot_payload": b"x" * 100, + "skills_metadata": [], + } + snap.metadata = { + "source": "loop", + "step": 42, + "writes": { + "CodeInterpreterMiddleware.after_agent": { + "_quickjs_snapshot_payload": ("snap", b"y" * 1_400_000), + "messages": [], + }, + "model": {"messages": ["m1"]}, + }, + "parents": {}, + } + _strip_private(snap) + snap._replace.assert_called_once() + kwargs = snap._replace.call_args.kwargs + assert "_quickjs_snapshot_payload" not in kwargs["values"] + assert "messages" in kwargs["values"] + assert "skills_metadata" in kwargs["values"] + scrubbed_writes = kwargs["metadata"]["writes"] + assert ( + "_quickjs_snapshot_payload" + not in scrubbed_writes["CodeInterpreterMiddleware.after_agent"] + ) + assert "messages" in scrubbed_writes["CodeInterpreterMiddleware.after_agent"] + assert scrubbed_writes["model"] == {"messages": ["m1"]} + # Non-writes metadata keys are preserved. + assert kwargs["metadata"]["source"] == "loop" + assert kwargs["metadata"]["step"] == 42 + # Sanity: the class exists and inherits from CompiledStateGraph. + from langgraph.graph.state import CompiledStateGraph + + assert issubclass(_EvoFilteredGraph, CompiledStateGraph) + + +def test_strip_private_handles_missing_metadata_writes(): + """``metadata['writes']`` can be missing or ``None`` on some snapshots + (e.g. initial state). The filter must not crash and must still strip + values. + """ + from EvoScientist.langgraph_dev.main_graph import _strip_private + + snap = MagicMock() + snap.values = {"_quickjs_snapshot_payload": b"x", "messages": []} + snap.metadata = {"source": "input", "step": -1, "writes": None} + snap.tasks = () + _strip_private(snap) + kwargs = snap._replace.call_args.kwargs + assert "_quickjs_snapshot_payload" not in kwargs["values"] + # writes was None, metadata passes through unchanged. + assert kwargs["metadata"]["writes"] is None + + +def test_strip_private_scrubs_task_result_snapshot_blob(): + """``tasks[*].result`` is where ``after_agent``'s return dict lands. + When the middleware snapshots, ``result`` carries + ``{"_quickjs_snapshot_payload": ("snap", ~1.4 MB bytes)}``. Verified + on live history: this is the dominant per-response leak, larger than + ``values`` and ``metadata.writes`` combined for anchor checkpoints. + """ + from EvoScientist.langgraph_dev.main_graph import _strip_private + + class FakeTask: + def __init__(self, id_, result): + self.id = id_ + self.name = "CodeInterpreterMiddleware.after_agent" + self.result = result + + def _replace(self, **kwargs): + for k, v in kwargs.items(): + setattr(self, k, v) + return self + + leaking_task = FakeTask( + "t1", {"_quickjs_snapshot_payload": ("snap", b"z" * 1_400_000), "messages": []} + ) + clean_task = FakeTask("t2", {"messages": ["hi"]}) + snap = MagicMock() + snap.values = {} + snap.metadata = {"source": "loop", "step": 5} + snap.tasks = (leaking_task, clean_task) + _strip_private(snap) + kwargs = snap._replace.call_args.kwargs + tasks_after = kwargs["tasks"] + assert "_quickjs_snapshot_payload" not in tasks_after[0].result + assert "messages" in tasks_after[0].result + # Clean task is passed through untouched. + assert tasks_after[1] is clean_task + + +def test_agent_uses_filtered_graph_class(): + """The ``__class__`` swap in ``main_graph.py`` is the load-bearing wiring + that makes ``_strip_private`` reach the langgraph-api endpoints. + ``_strip_private`` and ``_EvoFilteredGraph`` in isolation don't prove the + swap ran; every other test in this file passes even if someone drops the + swap line. This asserts the compiled agent is actually the filtered + subclass at module-load time, and that the subclass survives + ``Pregel.copy(update=...)`` — the call langgraph-api makes in + ``get_graph`` before yielding the graph to endpoint handlers. + """ + from EvoScientist.langgraph_dev.main_graph import ( + EvoScientist_agent, + _EvoFilteredGraph, + ) + + assert isinstance(EvoScientist_agent, _EvoFilteredGraph) + assert isinstance(EvoScientist_agent.copy(update={}), _EvoFilteredGraph) + + +def test_all_registered_graphs_use_filtered_graph_class(): + """Every graph registered in ``langgraph.json`` (main + all subagents) + gets the ``__class__`` swap via ``_apply_filter_to_all_registered_graphs``. + Iterating the config directly matches the auto-detect refactor: adding + a new subagent to ``langgraph.json`` should not require a corresponding + test update. + + Subagents get ``create_code_interpreter_middleware`` unconditionally + (``EvoScientist.py:_build_middleware_stack``), so they can touch the + QuickJS REPL and write ``_quickjs_snapshot_payload`` on their own + checkpoint namespace. Async subagents also get their own ``thread_id`` + and their ``/threads/{id}/state`` endpoint runs on their own compiled + graph — without the swap on those graphs, our filter would miss that + endpoint entirely. + """ + import json + from importlib import import_module + from pathlib import Path + + # Import triggers ``main_graph``'s swap loop. + from EvoScientist.langgraph_dev import main_graph + from EvoScientist.langgraph_dev.main_graph import _EvoFilteredGraph + + config_path = Path(main_graph.__file__).parent / "langgraph.json" + config = json.loads(config_path.read_text()) + for name, path in config["graphs"].items(): + module_path, attr = path.rsplit(":", 1) + graph = getattr(import_module(module_path), attr) + assert isinstance(graph, _EvoFilteredGraph), ( + f"graph {name!r} ({path}) did not receive the class swap" + ) + + +def test_strip_private_recurses_into_nested_subgraph_state(): + """When ``subgraphs=True``, ``PregelTask.state`` holds a nested + ``StateSnapshot`` for the subgraph. Its ``values`` (and its own nested + tasks) can carry ``_quickjs_snapshot_payload`` just like the parent. + Recursion covers the compound leak path CodeRabbit flagged. + """ + from langgraph.types import StateSnapshot + + from EvoScientist.langgraph_dev.main_graph import _strip_private + + nested_snap = StateSnapshot( + values={"_quickjs_snapshot_payload": b"n" * 1_400_000, "messages": []}, + next=(), + config={}, + metadata={"source": "loop", "step": 3}, + created_at="2026-07-01T12:00:00Z", + parent_config=None, + tasks=(), + interrupts=(), + ) + + class FakeTask: + def __init__(self, state): + self.id = "sub-1" + self.name = "subgraph" + self.result = None + self.state = state + + def _replace(self, **kwargs): + for k, v in kwargs.items(): + setattr(self, k, v) + return self + + task_with_nested = FakeTask(nested_snap) + task_with_config_state = FakeTask({"configurable": {"thread_id": "t"}}) + snap = MagicMock() + snap.values = {} + snap.metadata = {"source": "loop", "step": 5} + snap.tasks = (task_with_nested, task_with_config_state) + _strip_private(snap) + kwargs = snap._replace.call_args.kwargs + tasks_after = kwargs["tasks"] + # Nested StateSnapshot got recursively scrubbed. + assert "_quickjs_snapshot_payload" not in tasks_after[0].state.values + assert "messages" in tasks_after[0].state.values + # A dict (RunnableConfig-shaped) state passes through unchanged — we only + # recurse into ``StateSnapshot`` instances. + assert tasks_after[1].state == {"configurable": {"thread_id": "t"}} + + +def test_strip_private_scrubs_delta_counters(): + """``metadata['counters_since_delta_snapshot']`` is a small + ``{channel: [count, superstep]}`` bookkeeping map. Not a size problem, + but leaks the channel name — strip for consistency with the private + annotation. + """ + from EvoScientist.langgraph_dev.main_graph import _strip_private + + snap = MagicMock() + snap.values = {} + snap.metadata = { + "source": "loop", + "step": 5, + "counters_since_delta_snapshot": { + "_quickjs_snapshot_payload": [1, 14], + "messages": [3, 14], + }, + } + snap.tasks = () + _strip_private(snap) + kwargs = snap._replace.call_args.kwargs + counters = kwargs["metadata"]["counters_since_delta_snapshot"] + assert "_quickjs_snapshot_payload" not in counters + assert "messages" in counters