diff --git a/agent/context_compressor.py b/agent/context_compressor.py index 4d95eda14fb0f..be66f8d2f87a9 100644 --- a/agent/context_compressor.py +++ b/agent/context_compressor.py @@ -4706,8 +4706,17 @@ This compaction should PRIORITISE preserving all information related to the focu # only genuinely orphaned calls in the *discarded* region are # stripped. trailing_inflight: Optional[Dict[str, Any]] = None - if messages and messages[-1].get("role") == "assistant": - trailing_inflight = messages[-1] + # Walk back over any trailing tool results first: with a + # multi-call batch the executor appends results one at a time, so + # a snapshot taken between appends looks like + # ``[..., assistant(c1,c2,c3), tool(c1)]`` — the chain is still + # in flight even though the last message is a tool result. The + # last NON-tool message is the live request in both shapes. + idx = len(messages) - 1 + while idx >= 0 and messages[idx].get("role") == "tool": + idx -= 1 + if idx >= 0 and messages[idx].get("role") == "assistant": + trailing_inflight = messages[idx] # ----------------------------------------------------------------- for msg in messages: if msg.get("role") != "assistant": diff --git a/tests/agent/test_context_compressor.py b/tests/agent/test_context_compressor.py index 4ef54085f00e3..eaa730c14d01f 100644 --- a/tests/agent/test_context_compressor.py +++ b/tests/agent/test_context_compressor.py @@ -2150,8 +2150,16 @@ class TestSanitizerStripsOrphanedToolCalls: assert asst.get("content") == "(tool call removed)" def test_sanitizer_strips_orphaned_keeps_valid(self, compressor): - """When an assistant has both valid and orphaned tool_calls, only - the orphans are stripped. #51218""" + """When a MID-LIST assistant has both valid and orphaned tool_calls, + only the orphans are stripped. #51218 + + The shape must sit mid-list: the same shape at the TAIL is + indistinguishable from a partial multi-call batch whose remaining + results are still in flight, and the sanitizer now presumes in-flight + there (#79278) — preserving is safe because the pre-API chokepoint + injects stub results for genuinely unanswered calls, while stripping + a live call silently loses its late result. + """ msgs = [ { "role": "assistant", @@ -2162,6 +2170,8 @@ class TestSanitizerStripsOrphanedToolCalls: ], }, {"role": "tool", "tool_call_id": "tc_valid", "content": "file content"}, + # Later turn: the chain above is settled history, not in flight. + {"role": "assistant", "content": "done"}, ] sanitized = compressor._sanitize_tool_pairs(msgs) @@ -2396,6 +2406,34 @@ class TestSanitizerPreservesInFlightToolChain: assert len(results) == 1 assert "42" in results[0]["content"] + def test_partial_batch_inflight_calls_preserved(self, compressor): + """Multi-call batch snapshotted BETWEEN result appends: the executor + has appended tool(c1) but not yet tool(c2)/tool(c3), so the last + message is a tool result while c2/c3 are still pending. The walk-back + must find the assistant behind the trailing results and preserve the + whole batch — stripping c2/c3 there loses their late results exactly + like the tail-is-assistant shape. #79278 follow-up.""" + msgs = [ + {"role": "user", "content": "run the batch"}, + {"role": "assistant", "content": "", "tool_calls": [ + {"id": "c1", "function": {"name": "a", "arguments": "{}"}}, + {"id": "c2", "function": {"name": "b", "arguments": "{}"}}, + {"id": "c3", "function": {"name": "c", "arguments": "{}"}}, + ]}, + {"role": "tool", "tool_call_id": "c1", "content": "done 1"}, + # snapshot taken here: c2/c3 results not yet appended + ] + + sanitized = compressor._sanitize_tool_pairs(msgs) + + batch = [m for m in sanitized if m.get("role") == "assistant"][-1] + assert [tc["id"] for tc in batch["tool_calls"]] == ["c1", "c2", "c3"] + # And the already-arrived result survives too. + assert any( + m.get("role") == "tool" and m.get("tool_call_id") == "c1" + for m in sanitized + ) + class TestCooldownReentryAbort: """Regression: a second compress() call during the failure cooldown must