perf: eliminate full message-history deep clone in add_messages reducer - #30
Open
MuFengMuXue wants to merge 5 commits into
Open
perf: eliminate full message-history deep clone in add_messages reducer#30MuFengMuXue wants to merge 5 commits into
MuFengMuXue wants to merge 5 commits into
Conversation
The messages reducer is the default for the `messages` channel in every agent state. Each LLM turn previously deep-cloned the entire accumulated history (all message content + tool args) — O(n^2) over a run. Change the reducer signature from `fn(&JsonValue, &JsonValue) -> JsonValue` to `fn(JsonValue, &JsonValue) -> JsonValue` and have BinaryOperatorAggregate::update hand ownership of the accumulated value to the reducer via `guard.take()`, so existing messages are moved instead of cloned. `update` stays borrowed (usually a single new message). add_messages now: - consumes `current` and moves existing messages into the result - collects remove-ids as borrowed `&str` (no per-message String alloc) - takes a fast path when there are no removals A benchmark is added (tests/bench_messages.rs, #[ignore]d) measuring the reducer hot path: ~100x faster at 250 steps up to ~330x at 1000 steps, and scaling drops from quadratic to near-linear. This is a breaking API change for reducers: they now receive `current` by value.
Checkpoint saves re-encoded and re-inserted a blob row for every channel on every super-step. With a large static channel (e.g. embedded context) that meant re-serializing and re-writing the same value each step — pure write amplification. save_checkpoint now takes the channel versions as of the start of the run and derives new_versions = channels whose version moved, passing only the delta to the saver. SqliteSaver::put writes blob rows only for new_versions; reads merge version-joined blob values over the checkpoint row body (extend, not replace) so every channel resolves through the version join regardless of which checkpoint wrote the blob. Adds test_incremental_blob_writes (a second checkpoint with one updated channel writes one new blob row; both new and old checkpoints read back fully) and a benchmark suite (tests/bench_pregel.rs, #[ignore]d) covering checkpointed growth and a static-context scenario. Measured (release, 800 steps, per-step): sqlite linear 12.651->12.230ms (~3%), static 200KB context 12.723->12.200ms (~4%). The win is small on these micro-benches because the dominant per-step cost is in-process state handling (read_channels deep clone), not blob I/O; write amplification is eliminated, which matters far more for remote savers and large static channels.
…he multi-step loop measurement - LatestOnlySaver retains only the newest checkpoint per thread. InMemorySaver keeps every checkpoint, so the growth benches retained O(steps²) serialized state and OOM'd after a few hundred steps; retention is now O(latest state). - The multi-step loop's conditional routing only sees the node output, so routing on messages.len() always saw a single-element array and the loop ran to the recursion limit (25 default) — silently truncating every run. Route on an explicit count channel so the loop runs the real target super-steps. - Same bench content as PR Onelevenvy#31's correction commit (6006a43) so either PR merging first lands the identical corrected bench.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
The
messagesreducer (add_messages/add_messages_ref) is the defaultreducer for the
messageschannel in every agent state. Previously every LLMturn deep-cloned the entire accumulated message history (all content +
tool-call args) just to append a new message — O(n²) over a run.
Changes
ReducerFnchanged fromfn(&JsonValue, &JsonValue) -> JsonValuetofn(JsonValue, &JsonValue) -> JsonValue.BinaryOperatorAggregate::updatenow hands ownership of the accumulatedvalue to the reducer via
guard.take(), so existing messages are movedinstead of deep-cloned (
updatestays borrowed — usually one new message).add_messages: consumescurrent, moves existing messages into theresult, collects remove-ids as borrowed
&str(no per-messageStringalloc), and takes a fast path when there are no removals.
add_messages_refkept as a thin forwarder for the derive macro's#[channel(messages)]and existing#[channel(reducer = "add_messages_ref")].Breaking change
Reducers now receive
currentby value. Custom reducers written asfn(&JsonValue, &JsonValue) -> JsonValuemust change the first parameter toJsonValue.Benchmark
Added
crates/langgraph-prebuilt/tests/bench_messages.rs(#[ignore]d),reducer hot path, best-of-3, release:
Scaling: quadratic → near-linear.
Testing
cargo build --workspace✅cargo test --workspace✅ (one pre-existing unrelated failure inlanggraph-checkpoint-sqlite::test_put_writes_and_pending_writes_round_trip,reproduced on clean
main)cargo clippy --all-targets --all-features -D warnings✅cargo fmt --all --check✅Note (2026-08-02) — bench corrected
tests/bench_pregel.rswas corrected in2c5d99c:InMemorySaver, which keeps every checkpoint and retained O(steps²) serialized state (OOM after a few hundred steps). They now useLatestOnlySaver(retains only the newest checkpoint per thread).messages.len()always saw a single-element array and silently truncated every run at the recursion limit. It now routes on an explicitcountchannel and runs the real super-steps.6006a43), so whichever PR merges first lands the same file.