Skip to content

fix(flush_memories): always deduct headroom + resolve flush aux model + trim defence #15638

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
kshitijk4poor wants to merge 1 commit into from
+356 −52
Closed

fix(flush_memories): always deduct headroom + resolve flush aux model + trim defence #15638

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
129 changes: 107 additions & 22 deletions run_agent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2402,6 +2402,34 @@
provider=getattr(self, "provider", ""),
)

# Also resolve the flush_memories auxiliary model — it may differ
# from the compression model when the user configures separate
# auxiliary.flush_memories.provider/model, or when the fallback
# chain lands on a different provider. flush_memories runs with
# the FULL pre-compression conversation, so its model's context
# must also be respected.
try:
flush_client, flush_model = get_text_auxiliary_client(
"flush_memories",
main_runtime=self._current_main_runtime(),
)
if flush_client and flush_model:
_flush_ctx = get_model_context_length(
flush_model,
base_url=str(getattr(flush_client, "base_url", "") or ""),
api_key=str(getattr(flush_client, "api_key", "") or ""),
provider=getattr(self, "provider", ""),
)
if _flush_ctx and _flush_ctx < aux_context:
logger.info(
"flush_memories model %s context (%d) < compression "
"model %s context (%d) — using the smaller value",
flush_model, _flush_ctx, aux_model, aux_context,
)
aux_context = _flush_ctx
except Exception:
pass # Non-fatal — fall through with compression model's context

# Hard floor: the auxiliary compression model must have at least
# MINIMUM_CONTEXT_LENGTH (64K) tokens of context. The main model
# is already required to meet this floor (checked earlier in
Expand All @@ -2421,29 +2449,25 @@
)

threshold = self.context_compressor.threshold_tokens
if aux_context < threshold:
# Auto-correct: lower the live session threshold so
# compression actually works this session. The hard floor
# above guarantees aux_context >= MINIMUM_CONTEXT_LENGTH,
# so the new threshold is always >= 64K.
#
# Headroom: the threshold budgets RAW MESSAGES only, but the
# actual request auxiliary callers send also includes the
# system prompt and every tool schema. With 50+ tools that
# overhead can be 25-30K tokens; setting new_threshold =
# aux_context directly would let messages grow right to the
# aux limit and the first compression/flush request would
# overflow with HTTP 400. Subtract a dynamic headroom
# estimate so the full request still fits.
from agent.model_metadata import estimate_request_tokens_rough
tool_overhead = estimate_request_tokens_rough([], tools=self.tools)
# System prompt is not yet built at __init__ time; allow a
# conservative 10K budget (SOUL/AGENTS.md + memory snapshot +
# skills guidance) plus 2K for the flush instruction and a
# small safety margin.
headroom = tool_overhead + 12_000

# Headroom: the threshold budgets RAW MESSAGES only, but the
# actual request auxiliary callers (compression summariser and
# flush_memories) send also includes the system prompt and every
# tool schema. We must ensure threshold + headroom <= aux_context
# or the first compression/flush request will overflow.
#
# This applies even when aux_context > threshold (the common
# same-model case after a155b4a1) — e.g. 128K context, 85%
# threshold = 108K, 20K overhead → 108K + 20K = 128K exactly
# at the limit, and any token-estimate variance causes a 400.
from agent.model_metadata import estimate_request_tokens_rough
tool_overhead = estimate_request_tokens_rough([], tools=self.tools)
headroom = tool_overhead + 12_000
effective_limit = max(aux_context - headroom, MINIMUM_CONTEXT_LENGTH)

if effective_limit < threshold:
old_threshold = threshold
new_threshold = max(aux_context - headroom, MINIMUM_CONTEXT_LENGTH)
new_threshold = effective_limit
self.context_compressor.threshold_tokens = new_threshold
# Keep threshold_percent in sync so future main-model
# context_length changes (update_model) re-derive from a
Expand Down Expand Up @@ -7992,6 +8016,67 @@
messages.pop() # remove flush msg
return

# ── Defence-in-depth: trim messages to fit auxiliary context ──
#
# _check_compression_model_feasibility already lowers the
# compression threshold so conversations *triggered by preflight
# compression* should fit. But flush_memories is also called
# from CLI /new and gateway session resets — paths that bypass
# the preflight check entirely. Trim here as a safety net.
try:
from agent.auxiliary_client import get_text_auxiliary_client
from agent.model_metadata import (
get_model_context_length,
estimate_messages_tokens_rough,
)
_fc, _fm = get_text_auxiliary_client(
"flush_memories",
main_runtime=self._current_main_runtime(),
)
_fctx = 0
if _fc and _fm:
_fctx = get_model_context_length(
_fm,
base_url=str(getattr(_fc, "base_url", "") or ""),
api_key=str(getattr(_fc, "api_key", "") or ""),
provider=getattr(self, "provider", ""),
)
if not _fctx:
_fctx = getattr(
getattr(self, "context_compressor", None),
"context_length", 0,
)
if _fctx:
_budget = _fctx - 5120 - 500 # output + tool schema
if _budget > 0:
_est = estimate_messages_tokens_rough(api_messages)
if _est > _budget:
_sys = []
_conv = api_messages
if api_messages and api_messages[0].get("role") == "system":
_sys = [api_messages[0]]
_conv = api_messages[1:]
_rem = _budget - estimate_messages_tokens_rough(_sys)
_kept: list = []
_acc = 0
for _m in reversed(_conv):
_mt = estimate_messages_tokens_rough([_m])
if _acc + _mt > _rem:
break
_kept.append(_m)
_acc += _mt
_kept.reverse()
if len(_kept) < 3 and len(_conv) >= 3:
_kept = _conv[-3:]
api_messages = _sys + _kept
logger.info(
"flush_memories: trimmed %d→%d msgs to fit "
"%d-token aux context",
len(_sys) + len(_conv), len(api_messages), _fctx,
)
except Exception as _te:
logger.debug("flush_memories: context trim failed: %s", _te)

# Use auxiliary client for the flush call when available --
# it's cheaper and avoids Codex Responses API incompatibility.
from agent.auxiliary_client import (
Expand Down
60 changes: 30 additions & 30 deletions tests/run_agent/test_compression_feasibility.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -151,15 +151,14 @@ def test_feasibility_check_passes_live_main_runtime():
agent._emit_status = lambda msg: None
agent._check_compression_model_feasibility()

mock_get_client.assert_called_once_with(
"compression",
main_runtime={
"model": "gpt-5.4",
"provider": "openai-codex",
# Called for both compression + flush_memories; verify compression call present
assert any(
c == (("compression",), {"main_runtime": {
"model": "gpt-5.4", "provider": "openai-codex",
"base_url": "https://chatgpt.com/backend-api/codex",
"api_key": "codex-token",
"api_mode": "codex_responses",
},
"api_key": "codex-token", "api_mode": "codex_responses",
}})
for c in mock_get_client.call_args_list
)


Expand All @@ -179,12 +178,12 @@ def test_feasibility_check_passes_config_context_length(mock_get_client, mock_ct
agent._emit_status = lambda msg: None
agent._check_compression_model_feasibility()

mock_ctx_len.assert_called_once_with(
"custom/big-model",
base_url="http://custom-endpoint:8080/v1",
api_key="sk-custom",
config_context_length=1_000_000,
provider="openrouter",
# First call is the compression model
assert mock_ctx_len.call_args_list[0] == (
("custom/big-model",),
{"base_url": "http://custom-endpoint:8080/v1",
"api_key": "sk-custom", "config_context_length": 1_000_000,
"provider": "openrouter"},
)


Expand All @@ -202,12 +201,11 @@ def test_feasibility_check_ignores_invalid_context_length(mock_get_client, mock_
agent._emit_status = lambda msg: None
agent._check_compression_model_feasibility()

mock_ctx_len.assert_called_once_with(
"custom/model",
base_url="http://custom:8080/v1",
api_key="sk-test",
config_context_length=None,
provider="openrouter",
assert mock_ctx_len.call_args_list[0] == (
("custom/model",),
{"base_url": "http://custom:8080/v1",
"api_key": "sk-test", "config_context_length": None,
"provider": "openrouter"},
)


Expand Down Expand Up @@ -255,13 +253,10 @@ def on_session_start(self, *args, **kwargs):
)

assert agent._aux_compression_context_length_config == 1_000_000
mock_ctx_len.assert_called_once_with(
"custom/big-model",
base_url="http://custom-endpoint:8080/v1",
api_key="sk-custom",
config_context_length=1_000_000,
provider="",
)
c0 = mock_ctx_len.call_args_list[0]
assert c0.args == ("custom/big-model",)
assert c0.kwargs["base_url"] == "http://custom-endpoint:8080/v1"
assert c0.kwargs["config_context_length"] == 1_000_000


@patch("agent.auxiliary_client.get_text_auxiliary_client")
Expand Down Expand Up @@ -311,8 +306,10 @@ def test_exception_does_not_crash(mock_get_client):

@patch("agent.model_metadata.get_model_context_length", return_value=100_000)
@patch("agent.auxiliary_client.get_text_auxiliary_client")
def test_exact_threshold_boundary_no_warning(mock_get_client, mock_ctx_len):
"""No warning when aux context exactly equals the threshold."""
def test_exact_threshold_boundary_triggers_headroom_correction(mock_get_client, mock_ctx_len):
"""When aux context exactly equals the threshold, headroom deduction
still fires — flush_memories adds system prompt + tool schema on top
of the conversation messages, so threshold must be lowered."""
agent = _make_agent(main_context=200_000, threshold_percent=0.50)
mock_client = MagicMock()
mock_client.base_url = "https://openrouter.ai/api/v1"
Expand All @@ -324,7 +321,10 @@ def test_exact_threshold_boundary_no_warning(mock_get_client, mock_ctx_len):

agent._check_compression_model_feasibility()

assert len(messages) == 0
# 100K - headroom < 100K → auto-corrects
assert len(messages) == 1
assert "Auto-lowered" in messages[0]
assert agent.context_compressor.threshold_tokens < 100_000


@patch("agent.model_metadata.get_model_context_length", return_value=99_999)
Expand Down
Loading
Loading