fix: scope quickjs snapshot to turn to keep checkpoints small (#316)
* fix: scope quickjs snapshot to turn to keep checkpoints small * fix: strip _quickjs_snapshot_payload from state/history responses instead of dropping mode=thread * fix: recurse strip into nested subgraph StateSnapshot * fix: drop conditional-snapshot gate that leaked repl slots * fix: assert LangGraph state-shape invariants at import time * refactor: discover graphs to filter from langgraph.json * test: assert copy() preserves subclass; iterate langgraph.json for subagent coverage --------- Co-authored-by: Xi Zhang <106144707+X-iZhang@users.noreply.github.com>
This commit is contained in:
jfilipiuk
@@ -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"]
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user