fix(compaction): chain any input filter already set on the run config

The image scrub for text-only models also sets call_model_input_filter; replacing it dropped the scrub.
This commit is contained in:
Ian 2026-10-06 21:49:28 -04:00
parent 94adabd67e
commit f565cece23
2 changed files with 29 additions and 3 deletions

View file

@ -4,6 +4,7 @@ from __future__ import annotations
import asyncio
import contextlib
import inspect
import logging
import uuid
from collections.abc import Callable
@ -147,22 +148,27 @@ def _with_compaction_check(run_config: RunConfig) -> RunConfig:
model = _run_config_model(run_config)
if model is None:
return run_config
inner = run_config.call_model_input_filter
first_call = True
async def _check(data: CallModelData[Any]) -> ModelInputData:
nonlocal first_call
model_data = data.model_data
if inner is not None:
result = inner(data)
model_data = await result if inspect.isawaitable(result) else result
# Skip the first call: the cycle compacts right before the run, and the run's
# new input is only persisted after this filter, so stopping here would lose it.
if not first_call and await asyncio.to_thread(
needs_compaction,
model,
data.model_data.instructions or "",
model_data.instructions or "",
_agent_tools_text(data.agent),
data.model_data.input,
model_data.input,
):
raise CompactionNeededError(model)
first_call = False
return data.model_data
return model_data
return replace(run_config, call_model_input_filter=_check)

View file

@ -1553,6 +1553,26 @@ async def test_compaction_check_stops_run_after_first_call_when_over_budget(
await guard(data)
@pytest.mark.asyncio
async def test_compaction_check_chains_an_existing_input_filter(
monkeypatch: pytest.MonkeyPatch,
) -> None:
seen: list[Any] = []
monkeypatch.setattr(execution, "needs_compaction", lambda *args: seen.append(args[3]))
filtered = ModelInputData(input=[{"role": "user", "content": "scrubbed"}], instructions="sys")
run_config = RunConfig(model="m", call_model_input_filter=lambda _data: filtered)
check: Any = execution._with_compaction_check(run_config).call_model_input_filter
data = CallModelData(
model_data=ModelInputData(input=[], instructions="sys"),
agent=MagicMock(tools=[]),
context=None,
)
assert await check(data) is filtered
assert await check(data) is filtered
assert seen == [filtered.input]
class _BudgetStream:
def __init__(self, exc: BaseException | None) -> None:
self.run_loop_exception = exc