fix(gateway): preserve memory prompt during manual compression

This commit is contained in:
Gille 2026-07-28 11:12:16 -06:00 committed by kshitij
parent 28d11ab38c
commit 678916b427
2 changed files with 92 additions and 2 deletions

View File

@ -3914,7 +3914,11 @@ class GatewaySlashCommandsMixin:
# unbound/default "cli" host source — see #50422. _platform_config_key
# maps LOCAL->"cli" exactly like the live turn, avoiding a new
# "local" vs "cli" mismatch.
from gateway.run import _platform_config_key
from gateway.run import (
_GATEWAY_HYGIENE_PLATFORM,
_platform_config_key,
_seed_hygiene_system_prompt,
)
platform_key = (
_platform_config_key(source.platform) if source.platform else None
)
@ -3963,6 +3967,25 @@ class GatewaySlashCommandsMixin:
if platform_key is not None:
runtime_kwargs["platform"] = platform_key
runtime_kwargs["gateway_session_key"] = session_key
# The manual compression helper skips memory-provider initialization,
# but _compress_context may persist its cached system prompt. Restore
# the exact live-session prompt so provider blocks are retained.
session_row = None
get_session = getattr(self._session_db, "get_session", None)
if callable(get_session):
try:
session_row = await get_session(session_entry.session_id)
except Exception as exc:
logger.warning(
"Manual compression could not restore the system prompt "
"for session %s: %s. Preserving an empty prompt so the "
"live turn rebuilds it with its configured providers.",
session_entry.session_id,
exc,
exc_info=True,
)
tmp_agent = AIAgent(
**runtime_kwargs,
model=model,
@ -3973,6 +3996,12 @@ class GatewaySlashCommandsMixin:
session_id=session_entry.session_id,
session_db=getattr(self._session_db, "_db", self._session_db),
)
_seed_hygiene_system_prompt(tmp_agent, session_row)
# Keep the real source platform during construction so external
# context engines bind correctly. If compression has to rebuild the
# prompt, stamp that provider-less fallback as stale for the next
# real gateway turn.
tmp_agent.platform = _GATEWAY_HYGIENE_PLATFORM
try:
tmp_agent._print_fn = lambda *a, **kw: None
# Prevent close() from ending the newly rotated session —

View File

@ -1,7 +1,7 @@
"""Tests for gateway /compress user-facing messaging."""
from datetime import datetime
from unittest.mock import MagicMock, patch
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -511,6 +511,67 @@ async def test_compress_command_preserves_platform_and_gateway_session_key():
assert kwargs["gateway_session_key"]
@pytest.mark.asyncio
async def test_compress_command_preserves_persisted_provider_prompt():
"""Manual /compress must not replace a provider-aware session prompt.
Its temporary agent intentionally skips memory-provider initialization, so
it must reuse the exact persisted prompt. If compression rebuilds instead,
the hygiene-only marker makes that fallback stale for the next live turn.
"""
from gateway.run import _GATEWAY_HYGIENE_PLATFORM
history = _make_history()
stored_prompt = (
"base prompt\n\n"
"<hindsight_memories>\n"
"## Personal Memory\n"
"- pinned: exact provider content\n"
"</hindsight_memories>\n"
)
runner = _make_runner(history)
runner._session_db = MagicMock()
runner._session_db.get_session = AsyncMock(
return_value={"system_prompt": stored_prompt}
)
agent_instance = MagicMock()
agent_instance.shutdown_memory_provider = MagicMock()
agent_instance.close = MagicMock()
agent_instance._cached_system_prompt = "provider-less prompt"
agent_instance.platform = "telegram"
agent_instance.tools = None
agent_instance.context_compressor.has_content_to_compress.return_value = True
agent_instance.session_id = "sess-1"
agent_instance._compression_skipped_due_to_lock = False
def _compress(messages, *_args, **_kwargs):
assert messages == history
assert agent_instance._cached_system_prompt == stored_prompt
assert agent_instance.platform == _GATEWAY_HYGIENE_PLATFORM
return list(history), ""
agent_instance._compress_context.side_effect = _compress
def _estimate(messages, **kwargs):
assert messages == history
assert kwargs["system_prompt"] == stored_prompt
return 100
with (
patch("gateway.run._resolve_runtime_agent_kwargs", return_value={"api_key": "test-key"}),
patch("gateway.run._resolve_gateway_model", return_value="test-model"),
patch("run_agent.AIAgent", return_value=agent_instance) as mock_agent,
patch("agent.model_metadata.estimate_request_tokens_rough", side_effect=_estimate),
):
await runner._handle_compress_command(_make_event())
runner._session_db.get_session.assert_awaited_once_with("sess-1")
assert mock_agent.call_args.kwargs["platform"] == "telegram"
assert agent_instance._cached_system_prompt == stored_prompt
assert agent_instance.platform == _GATEWAY_HYGIENE_PLATFORM
@pytest.mark.asyncio
async def test_compress_command_overrides_stale_resolver_identity():
"""If the resolver already supplies platform/gateway_session_key, the