Skip to content

Compaction: clear server-side tool uses before client-side removal shifts indices#4529

Open
ebarkhordar wants to merge 1 commit into
UKGovernmentBEIS:mainfrom
ebarkhordar:fix/compaction-server-tool-stale-index
Open

Compaction: clear server-side tool uses before client-side removal shifts indices#4529
ebarkhordar wants to merge 1 commit into
UKGovernmentBEIS:mainfrom
ebarkhordar:fix/compaction-server-tool-stale-index

Conversation

@ebarkhordar

Copy link
Copy Markdown

This PR contains:

  • New features
  • Changes to dev-tools e.g. CI config / github tooling
  • Docs
  • Bug fixes
  • Code refactor

What is the current behavior? (You can also link to an open issue here)

Fixes #4528.

CompactionEdit collects 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 existing process in reverse to preserve indices comment). 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=False plus any server-side tool use (for example web_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:

  1. A message with string content: IndexError: list assignment index out of range.
  2. An assistant message with list content: nothing raises. The tool use is written into that unrelated message, overwriting the content at that position, and the tool use that was selected for clearing keeps its result.

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_clearing uses the default keep_tool_inputs=True (nothing is removed, so nothing shifts), and the keep_tool_inputs=False tests contain no server-side tool use.

What is the new behavior?

Server-side clearing runs before the client-side pass. _apply_server_tool_clearing only 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.py with ast rather than by reading. Inside compact() 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 trailing result.pop(). Of these, only del result[tool_idx] and the trailing pop() change the length. The pop() runs after all captured indices are consumed, which leaves del 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_clearing is length-preserving by the same enumeration: its only writers are result = list(messages) and result[msg_idx] = ....

Two regression tests, one per case above:

  • test_server_tool_clearing_after_client_message_removal fails on main with IndexError.
  • test_server_tool_clearing_does_not_write_into_other_messages fails on main with assert ['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 CompactionEdit with keep_tool_inputs=False and server-side tool uses present, which currently either raises or corrupts the history.

Other information:

Verified in a clean python:3.12-slim Docker container against 8117bf2, with the loaded edit.py sha256-matched to the checkout on both sides of the differential (main cd21b345..., branch 002060f8...), so the bytes under test are the ones in this diff.

  • Both new tests: fail on main, pass on this branch.
  • tests/model/test_compaction_edit.py + tests/model/test_compaction.py: 79 passed.
  • tests/model/: 1700 passed, 0 failed.
  • ruff check and ruff format --check with the pinned ruff==0.9.6 from requirements-dev.txt: clean.
  • pre-commit run --files on the changed files: all hooks pass (controlled against pristine HEAD first).
  • mypy on src/inspect_ai/model/_compaction/ and the test file: clean.

What I did not verify: the whole-repo mypy --exclude tests/test_package src tests and 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 drives CompactionEdit.compact directly with mockllm/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.

…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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

CompactionEdit: server-side tool uses are cleared at stale message indices when keep_tool_inputs=False removes messages

1 participant