diff --git a/libs/langchain_v1/langchain/agents/middleware/context_editing.py b/libs/langchain_v1/langchain/agents/middleware/context_editing.py index 6da6714bd9..08f6980bac 100644 --- a/libs/langchain_v1/langchain/agents/middleware/context_editing.py +++ b/libs/langchain_v1/langchain/agents/middleware/context_editing.py @@ -317,7 +317,18 @@ class ContextEditingMiddleware(AgentMiddleware[AgentState[ResponseT], ContextT, different counter could disagree with `wrap_model_call`'s decision and either skip persisting a clear that was applied, or persist one that wasn't. + + `ContextEdit.apply` is only required to mutate `messages` in place — + custom strategies are free to insert or remove entries, not just + replace them one-for-one. A same-length assumption would either raise + (pairing up mismatched indices) or silently mismatch messages, so this + only persists when the edit preserved length; otherwise the edited + request still reaches the model exactly as before, just without the + persistence benefit for that edit. """ + if len(original_messages) != len(edited_messages): + return response + changed = [ new_msg for old_msg, new_msg in zip(original_messages, edited_messages, strict=True) diff --git a/libs/langchain_v1/tests/unit_tests/agents/middleware/implementations/test_context_editing.py b/libs/langchain_v1/tests/unit_tests/agents/middleware/implementations/test_context_editing.py index 72b436e2ea..4357e407f0 100644 --- a/libs/langchain_v1/tests/unit_tests/agents/middleware/implementations/test_context_editing.py +++ b/libs/langchain_v1/tests/unit_tests/agents/middleware/implementations/test_context_editing.py @@ -355,6 +355,48 @@ def test_wrap_model_call_returns_plain_response_when_no_edit_applied() -> None: assert isinstance(result, ModelResponse) +def test_wrap_model_call_skips_persistence_for_length_changing_edit() -> None: + """A length-changing custom `ContextEdit` must not crash persistence. + + Regression test for a review comment on #37815: `ContextEdit.apply` is + only required to mutate the message list in place — it isn't required to + preserve length (e.g. a custom strategy could `messages.pop(0)`). Pairing + up messages by index to detect changes would previously raise + `ValueError` from a length-mismatched `zip(..., strict=True)`. The edited, + shortened request must still reach the model, just without persistence + for that edit. + """ + + class DropFirstMessageEdit: + """Custom edit that removes the oldest message instead of replacing it.""" + + def apply(self, messages: list[AnyMessage], *, count_tokens: Any) -> None: + del count_tokens + if messages: + messages.pop(0) + + ai_message = AIMessage(content="", tool_calls=[]) + tool_message = ToolMessage(content="x" * 200, tool_call_id="call-1") + + _state, request = _make_state_and_request([ai_message, tool_message]) + middleware = ContextEditingMiddleware(edits=[DropFirstMessageEdit()]) + + captured_request: ModelRequest | None = None + + def mock_handler(req: ModelRequest) -> ModelResponse: + nonlocal captured_request + captured_request = req + return ModelResponse(result=[AIMessage(content="mock response")]) + + result = middleware.wrap_model_call(request, mock_handler) + + assert captured_request is not None + # The model still receives the shortened request. + assert captured_request.messages == [tool_message] + # No persistence is attempted for a length-changing edit. + assert isinstance(result, ModelResponse) + + def test_persisted_edit_is_idempotent_across_turns() -> None: """Once persisted, a later turn must not re-clear an already-cleared message.