feat(agents): add prompt-compaction middleware for McpClient#2055
Open
Mgczacki wants to merge 7 commits into
Open
feat(agents): add prompt-compaction middleware for McpClient#2055Mgczacki wants to merge 7 commits into
Mgczacki wants to merge 7 commits into
Conversation
Caps the prompt the agent sends to its LLM so the conversation history
never grows unbounded. Runs as a langchain AgentMiddleware via
create_agent(middleware=...), so the size bound becomes an invariant of
the agent loop — `before_model` fires before every model call, including
intra-turn re-invocations (model -> tool -> tool result -> model).
Two-stage compaction:
1. Strip image content blocks from older messages (replace with a small
text placeholder).
2. If still over target, summarize older messages into a single
SystemMessage and keep the most recent turns verbatim.
The current turn (latest dimos_turn group + any trailing untagged
messages, i.e. in-flight tool calls) is preserved untouched — never
compacted, never image-stripped.
Configuration via McpClientConfig fields, env-driven by default:
AGENT_COMPACTION_THRESHOLD trigger size (default 40000)
AGENT_COMPACTION_TARGET size after compaction (default 3000)
AGENT_COMPACTION_SUMMARY_SIZE generated summary size (default 1000)
AGENT_COMPACTION_MODEL optional separate summarizer model
Also includes:
- Per-message turn tagging via additional_kwargs["dimos_turn"], stamped
in McpClient._process_message so compaction can group/score by turn.
- McpClient._history mirror updated to honor langgraph's add_messages
reducer semantics (RemoveMessage(id=REMOVE_ALL_MESSAGES) sentinel) so
the local history doesn't accrete pre-compaction state.
- Token counter is a pessimistic placeholder (3 chars/token,
1000/image), memoized on each message for O(new-only) recompute cost.
Designed to be swapped for a real tokenizer later without touching
callers.
- 15 pytest cases (hermetic, no API key needed), including two
integration tests that drive a real create_agent loop and prove
compaction can fire mid-turn between a tool result and the next
model call.
Defaults are intentionally conservative so the feature is on by default
without changing behavior for short sessions.
Contributor
Greptile SummaryThis PR adds
Confidence Score: 5/5Safe to merge. The compaction algorithm, turn-boundary protection, tool-call coherence, and duplicate-publish suppression are all well-handled, and the two full-loop integration tests confirm the middleware integrates correctly with langgraph's add_messages reducer. The core logic is sound: the backward-scan current-turn detection, the _split boundary alignment, and the _apply_messages_update replay logic all hold up under close inspection. The previous review concerns have been addressed. Remaining notes are minor style/edge-case issues that do not affect correctness. No files require special attention. compaction_middleware.py contains one dead method worth removing; mcp_client.py has a minor zero-coercion edge case in the env-var defaults. Important Files Changed
Sequence DiagramsequenceDiagram
participant U as "User / tool-stream"
participant MP as "McpClient._process_message"
participant SG as "LangGraph state_graph"
participant BM as "before_model hook"
participant LLM as "Agent LLM"
participant AU as "_apply_messages_update"
U->>MP: HumanMessage
MP->>MP: "increment turn, tag message"
MP->>MP: "append to history, publish"
MP->>SG: "stream(messages: history)"
loop "agent node execution"
SG->>BM: "before_model(state)"
alt "total <= threshold"
BM-->>SG: "None (no-op)"
else "Stage 1: image strip"
BM-->>SG: "RemoveMessage(ALL) + stripped + current_turn"
else "Stage 2: summarise"
BM->>LLM: "summarise older turns"
LLM-->>BM: "summary_text"
BM-->>SG: "RemoveMessage(ALL) + protected + summary + keep + current_turn"
end
SG->>LLM: "invoke(compacted messages)"
LLM-->>SG: "AIMessage"
SG-->>MP: "stream update {node: messages}"
MP->>AU: "_apply_messages_update(node_messages, turn)"
AU->>AU: "wipe history on RemoveMessage(ALL)"
AU->>AU: "skip-publish replayed objects"
AU->>AU: "append and publish new objects"
end
Reviews (6): Last reviewed commit: "fix: small efficiency rewrite" | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Comment on lines
+49
to
+51
| def _env_int(name: str) -> int | None: | ||
| v = os.environ.get(name) | ||
| return int(v) if v else None |
Contributor
There was a problem hiding this comment.
_env_int calls int(v) without a try/except, so a non-numeric value like AGENT_COMPACTION_THRESHOLD=abc raises a bare ValueError deep inside pydantic's default_factory during config construction, producing an unhelpful traceback with no mention of which env var is at fault.
Suggested change
| def _env_int(name: str) -> int | None: | |
| v = os.environ.get(name) | |
| return int(v) if v else None | |
| def _env_int(name: str) -> int | None: | |
| v = os.environ.get(name) | |
| if not v: | |
| return None | |
| try: | |
| return int(v) | |
| except ValueError: | |
| raise ValueError(f"Environment variable {name!r} must be an integer, got {v!r}") from None |
- McpClient._apply_messages_update: dedupe publish on compaction replay. When the middleware emits [RemoveMessage, protected..., summary, keep..., current_turn...], the protected/keep/current messages are the same Python objects that were already published when they first arrived. Skip publish+print for any iter_msg whose id() was in the pre-wipe history; only the genuinely-new summary (and later AIMessages from the agent node in subsequent stream updates) get republished. Identified by Greptile P1. - McpClient._env_int: re-raise a labeled ValueError when the env var value isn't a valid integer, so misconfiguration surfaces with the offending name instead of a bare pydantic traceback. Identified by Greptile P2. - DimosCompactionMiddleware._static_tokens: drop the per-call hash computation. Inputs (system_prompt, tool_schemas) are bound at __init__ and never mutate, so a simple None-check on the cache is sufficient. Identified by Greptile P2.
❌ 1 Tests Failed:
View the top 1 failed test(s) by shortest run time
To view more test analytics, go to the Test Analytics Dashboard |
Comment on lines
+384
to
+385
| if max_turn is None: | ||
| return len(messages) |
Contributor
There was a problem hiding this comment.
_current_turn_start returns len(messages) when no messages carry a dimos_turn tag, but this places every message into compactable and nothing into current_turn. The downstream split in before_model then treats the latest user query itself as eligible for summarization, silently folding the active request into the summary. The docstring says "the caller will no-op" in this case, which only holds if the function returns 0 (making compactable = []). Any deployment that feeds the middleware a history without turn tags — e.g., a session started before the tagging feature landed, or a standalone use outside McpClient — would have its current-turn messages summarized away on the first threshold crossing.
Suggested change
| if max_turn is None: | |
| return len(messages) | |
| if max_turn is None: | |
| return 0 |
Author
There was a problem hiding this comment.
This turns the middleware off and keeps the problem intact. I have added an alternative split for any agents that don't have turn-tagged messages where the compressible section starts just before the latest human message. It's a good compromise for that edge case.
added 2 commits
May 12, 2026 20:09
…Message When `_current_turn_start` encounters a history with no `dimos_turn` tags at all (a caller wired the middleware in without going through McpClient), it now walks back to find the latest `HumanMessage` and uses that as the boundary. Older messages compactable, latest user input + any in-flight assistant/tool messages after it protected. The previous behavior returned `len(messages)` — making every message compactable — which would silently summarize the active user query the first time compaction crossed threshold. (Returning `0` would protect everything and instead let an oversized prompt reach the LLM, where it would raise on context overflow — worse than the silent path.) Greptile P4. Also: clarify in the module docstring that tool-call coherence is the harness's responsibility. The middleware never introduces orphan tool_calls (same-turn messages always travel together via `_split`'s boundary alignment) but doesn't fix orphans it inherits — those flow through to the LLM in the current-turn case, where the malformed conversation will surface as a provider error. New regression test: test_untagged_history_anchors_current_turn_on_latest_human.
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
2 participants
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
Closes #1899
Caps the prompt the dimos agent sends to its LLM so the conversation history
never grows unbounded. Implemented as a langchain
AgentMiddlewareplugged intocreate_agent(middleware=...). Because the hook (before_model) fires beforeevery model invocation, the input-size bound becomes an invariant of the agent
loop — including intra-turn re-invocations (model → tool → tool result → model).
On long sessions the middleware quietly summarizes older turns once it detects
an oversized prompt. Behavior is unchanged for short sessions.
Concepts
dimos_turnA new integer tag attached to each message's
additional_kwargsdict.Incremented once per
McpClient._process_messagecall — that is, once peruser-facing turn (a human input from
agent-send, or a tool-streamnotification that wakes the agent). Every message that flows through during
that turn — the input
HumanMessage, intermediateAIMessages withtool_calls, the resultingToolMessages, the finalAIMessage— all getstamped with the same turn number.
This is what lets compaction:
together (compaction selects entire turns, never partial ones — no orphan
tool_call_idreferences).untagged in-flight messages from the agent loop) and preserve it untouched
regardless of threshold.
keep-N-most-recent strategies).
dimos_turnis metadata only — it lives inadditional_kwargs, whichproviders ignore but langchain serialization preserves. The compaction
summary itself is tagged with the max turn it covers (plus
dimos_compacted: True), so re-compaction folds the prior summary into thenext one cleanly.
Current turn is sacred
_current_turn_startwalks from the end of the message list to find theboundary of the latest turn. Everything from that boundary forward is never
compacted — no image strip, no summary touch. This protects:
ToolMessageresponsesHow it works
Two-stage compaction inside
before_model:Strip images in messages older than the current turn. Image content
blocks are replaced with a small text placeholder. If this alone gets us
below
target_tokens, we stop here.As to why I decided to strip images: LLM's visual reasoning capabilities are
currently noticeably worse than with text. Additionally, the way in which the
agent loop is set up right now makes it so the model gets to see the image at the
beginning of a new turn, and it tends to give a description of what's in the image.
This description is detailed enough for reasoning about the content of the image,
but it also causes a secondary effect: The model, when considering the image, will
default to anchor its perception (even if the image is available in chat history) to the
comment it gave at the moment. Keeping images that were already observed therefore
seems like a waste of tokens that we can save since we are already going to cause a
cache burst with our compaction process.
SystemMessagewhile keepingthe most recent turns verbatim. The summarizer LLM is configurable;
defaults to reusing the agent's own model. Output is hard-capped via
summarizer.bind(max_tokens=summary_size_tokens).See it in action
A public Langfuse trace captured with deliberately small defaults so
compaction fires after a handful of turns:
https://us.cloud.langfuse.com/project/cmp23t80n09ooad08jnw1lksy/traces/887630cfbf49bb97f1c5b4d2cc980ad1?observation=b73fcf77cb4f2dc5×tamp=2026-05-12T07:54:34.311Z
Use the trace timeline to see the prompt that hits the LLM at each
agent-turn-Nspan — older turns get folded into a single summarySystemMessageand the agent continues with a shrunk prompt.Configuration
All on by default via
McpClientConfig, env-driven:AGENT_COMPACTION_THRESHOLDagent_compaction_threshold40000AGENT_COMPACTION_TARGETagent_compaction_target3000AGENT_COMPACTION_SUMMARY_SIZEagent_compaction_summary_size1000AGENT_COMPACTION_MODELagent_compaction_modelNone(reuses agent's model)Why a middleware
Two reasons, both documented in
compaction_middleware.py's module docstring:_historywould only fire once per user turn, leaving every intra-turn re-invocation
unprotected. Middleware fires before each model call.
before_modelvsafter_model/wrap_model_call.before_modelisthe minimal-intervention hook.
after_modelis too late (the modelalready errored on overflow);
wrap_model_callconflates compaction withthe model-call concerns (retries, error shaping, tool dispatch).
Changes
New files
dimos/agents/compaction_middleware.py—DimosCompactionMiddlewareclass (subclass of
langchain.agents.middleware.AgentMiddleware),placeholder token counter (3 chars/token, 1000 tokens/image; memoized in
additional_kwargs[\"dimos_tokens\"]for O(new-only) recompute), statictoken cache for
system_prompt+ tool schemas, and the algorithm helpers(
_strip_images,_split,_current_turn_start,_summarize).dimos/agents/test_compaction_middleware.py— 15 pytest cases,hermetic (no API key needed). Coverage includes:
before_modelno-op below thresholdcreate_agentloop with aRecordingFakeAgentand assert: (a) the agent node receives a compactedprompt (proves langgraph's
add_messagesreducer interprets theRemoveMessage(REMOVE_ALL_MESSAGES)sentinel correctly), and(b) compaction can fire mid-turn between a tool result and the next
model call.
Modified:
dimos/agents/mcp/mcp_client.pyMcpClientConfigreading the env vars inthe table above.
_env_int/_env_strhelpers loaded via pydanticField(default_factory=...)._turn: intcounter onMcpClient(incremented atthe top of
_process_message), and a new module-level_tag_turn(message, turn)helper that stampsadditional_kwargs[\"dimos_turn\"]. Everymessage flowing through a turn gets stamped — the incoming
HumanMessagefirst, then every message emitted by the state graph.
_apply_messages_updatemethod that mirrorslanggraph's
add_messagesreducer semantics locally — honorsRemoveMessage(id=REMOVE_ALL_MESSAGES)as "wipe history, use whatfollows" and specific-id
RemoveMessageas targeted removal. This keepsMcpClient._historyin sync with the graph's internal state even when themiddleware replaces the entire message list.
on_system_modules, construct the summarizer(either via
init_chat_model(agent_compaction_model), orinit_chat_model(model)if the agent's model is a string, or reuse theagent's model object), build the middleware with the system prompt and
tool JSON schemas (
t.args_schema.model_json_schema()), and pass it ascreate_agent(..., middleware=middleware).middleware no-op updates that yield
{node: None}instead of{node: {\"messages\": [...]}}, which would previously crash with'NoneType' object has no attribute 'get'.Modified:
.gitignoreAdds
MUJOCO_LOG.TXT(MuJoCo runtime artifact written to the repo root onevery sim run; should never be committed).
Test plan
uv run pytest dimos/agents/test_compaction_middleware.py -v— 15/15pass.
uv run mypy dimos/agents/compaction_middleware.py dimos/agents/test_compaction_middleware.py dimos/agents/mcp/mcp_client.py— clean.dimos --simulation run unitree-go2-agenticwithAGENT_COMPACTION_THRESHOLD=2000, drive the agent until the thresholdis crossed, confirm a
Compaction fired (summarize)log line appearsand the next prompt sent to the LLM contains the summary
SystemMessageinstead of the older turns.Known limitations
Documented in the module docstring as "Known limitations":
Progressive disclosure with a content store is the right long-term answer.
very long session could exceed the summarizer model's own context window.
Mitigation deferred to a follow-up (chunked summarization).
@retry(on_exception=Exception)is intentionally broad because thesummarizer is duck-typed; permanent errors cost up to 3 attempts + 1s of
sleeps before propagating.