Compaction: clear server-side tool uses before client-side removal shifts indices#4529
Open
ebarkhordar wants to merge 1 commit into
Open
Conversation
…ifts indices (UKGovernmentBEIS#4528) CompactionEdit captures client-side and server-side tool use indices in one enumeration of the message list. With keep_tool_inputs=False the client-side pass deletes tool messages, shifting every later message down, but the server-side indices from that same enumeration were consumed afterwards with nothing rebasing them. Server-side clearing replaces messages without changing the list length, so applying it before the client-side pass keeps both sets of indices valid.
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
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.
This PR contains:
What is the current behavior? (You can also link to an open issue here)
Fixes #4528.
CompactionEditcollects client-side and server-side tool uses in a single enumeration of the message list (Phase 2 in_compaction/edit.py), recording each one's positional message index. Those indices share one implicit invariant: they stay valid only while the list's length does not change.When
keep_tool_inputs=False, the client-side pass breaks that. It removes tool messages entirely (del result[tool_idx]), shifting every later message down. The client-side pass defends its own indices by iterating in reverse (the existingprocess in reverse to preserve indicescomment). The server-side indices come from the same pre-deletion enumeration but were consumed after those deletions, and nothing rebased them.So with
keep_tool_inputs=Falseplus any server-side tool use (for exampleweb_search) positioned after a cleared client-side tool call, the server-side clearing addressed the wrong message. Two outcomes, depending on what shifted into the stale index:IndexError: list assignment index out of range.The second case is the one that worried me more: the history sent to the model then contains a tool use the assistant never made, while compaction silently fails to compact the entry it chose.
The gap dates to #3120, which added server-side tool compaction. The existing tests cover each half but never cross them:
test_mixed_client_and_server_tool_clearinguses the defaultkeep_tool_inputs=True(nothing is removed, so nothing shifts), and thekeep_tool_inputs=Falsetests contain no server-side tool use.What is the new behavior?
Server-side clearing runs before the client-side pass.
_apply_server_tool_clearingonly replaces messages and never changes the list length, so doing it first leaves the client-side indices valid (that pass still guards itself by iterating in reverse). This is a reordering of the existing call, not a change to either clearing routine.Since the invariant is about every writer of the message list rather than the one path a test exercises, I enumerated them by parsing
edit.pywithastrather than by reading. Insidecompact()the writers are:result = list(messages)(creation);result[i] = _clear_reasoning(msg);result = _apply_server_tool_clearing(...);result[tool_idx] = ...model_copy(...);del result[tool_idx];result[assistant_idx] = _replace_tool_call_with_text(...);result = clear_memory_content(result);result = strip_citations(result); and the trailingresult.pop(). Of these, onlydel result[tool_idx]and the trailingpop()change the length. Thepop()runs after all captured indices are consumed, which leavesdel result[tool_idx]as the only one that can invalidate a captured index, and after this change no captured index is consumed after it._apply_server_tool_clearingis length-preserving by the same enumeration: its only writers areresult = list(messages)andresult[msg_idx] = ....Two regression tests, one per case above:
test_server_tool_clearing_after_client_message_removalfails on main withIndexError.test_server_tool_clearing_does_not_write_into_other_messagesfails on main withassert ['A'] == ['A', 'B'], i.e. the overwritten content. It also asserts the final state: exactly one tool use survives anywhere, it is the original, and it is cleared.Does this PR introduce a breaking change? (What changes might users need to make in their application due to this PR?)
No. It only affects
CompactionEditwithkeep_tool_inputs=Falseand server-side tool uses present, which currently either raises or corrupts the history.Other information:
Verified in a clean
python:3.12-slimDocker container against8117bf2, with the loadededit.pysha256-matched to the checkout on both sides of the differential (maincd21b345..., branch002060f8...), so the bytes under test are the ones in this diff.tests/model/test_compaction_edit.py+tests/model/test_compaction.py: 79 passed.tests/model/: 1700 passed, 0 failed.ruff checkandruff format --checkwith the pinnedruff==0.9.6fromrequirements-dev.txt: clean.pre-commit run --fileson the changed files: all hooks pass (controlled against pristine HEAD first).mypyonsrc/inspect_ai/model/_compaction/and the test file: clean.What I did not verify: the whole-repo
mypy --exclude tests/test_package src testsand the full pytest suite did not complete on my machine (3GB RAM, the mypy run was OOM-killed), so I ran the scoped equivalents above and am relying on CI for the rest. I also did not exercise this against a live provider; the repro drivesCompactionEdit.compactdirectly withmockllm/model, as the existing tests in this file do.This contribution was made with AI assistance. I reproduced the bug, ran everything above myself, and am accountable for the change.