fix(gateway): preserve memory prompt during hygiene compression

This commit is contained in:
Gille
2026-07-27 17:40:18 -06:00
committed by kshitij
parent 1dfe781edd
commit 76a17046e2
2 changed files with 63 additions and 2 deletions

View File

@@ -85,6 +85,7 @@ _TELEGRAM_CONNECT_TIMEOUT_SECS_DEFAULT = 180.0
_ADAPTER_DISCONNECT_TIMEOUT_SECS_DEFAULT = 5.0
_GATEWAY_PROXY_SSE_BUFFER_MAX_CHARS = 16 * 1024 * 1024
_TELEGRAM_COMMAND_MENTION_RE = re.compile(r"(?<![\w:/])/([A-Za-z0-9][A-Za-z0-9_-]*)")
_GATEWAY_HYGIENE_PLATFORM = "gateway_hygiene"
_TELEGRAM_NOISY_STATUS_RE = re.compile(
r"(" # transient/auxiliary status that should stay in logs, not gateway chats
@@ -323,6 +324,43 @@ def _non_conversational_metadata(
return merged
def _seed_hygiene_system_prompt(
agent: Any,
session_db: Any,
session_id: str,
) -> bool:
"""Keep gateway hygiene from rebuilding a live session's system prompt.
The hygiene helper intentionally skips memory-provider initialization.
Compression is allowed to persist a system prompt, so letting that helper
rebuild one would strip external provider blocks from the live session.
Seed the exact persisted prompt instead. When no usable prompt can be
restored, seed an empty cache entry. Compression either preserves that
unusable value or rebuilds with the hygiene-only platform marker; the real
turn will rebuild either form with its fully initialized providers.
"""
stored_prompt = ""
if session_db is not None and session_id:
try:
session_row = session_db.get_session(session_id)
if isinstance(session_row, dict):
raw_prompt = session_row.get("system_prompt")
if isinstance(raw_prompt, str) and raw_prompt.strip():
stored_prompt = raw_prompt
except Exception as exc:
logger.warning(
"Session hygiene 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_id,
exc,
exc_info=True,
)
agent._cached_system_prompt = stored_prompt
return bool(stored_prompt)
def _is_transient_network_error(exc: BaseException) -> bool:
"""Return True for transient network errors safe to log + swallow.
@@ -13644,6 +13682,15 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
session_id=session_entry.session_id,
session_db=_hyg_session_db,
)
_seed_hygiene_system_prompt(
_hyg_agent,
_hyg_session_db,
session_entry.session_id,
)
# If compression must rebuild instead of retaining
# the cached prompt, make the persisted result
# deliberately stale for every real gateway surface.
_hyg_agent.platform = _GATEWAY_HYGIENE_PLATFORM
_hyg_cleanup_deferred = False
try:
# Gateway hygiene runs before the user turn

View File

@@ -1094,15 +1094,26 @@ async def test_session_hygiene_forces_in_place_compaction_with_bound_session_db(
fake_dotenv.load_dotenv = lambda *args, **kwargs: None
monkeypatch.setitem(sys.modules, "dotenv", fake_dotenv)
fake_db = object()
stored_system_prompt = (
"You are Hermes.\n\n"
"<memory_provider_context>\n"
"Pinboard provider instructions\n"
"</memory_provider_context>"
)
fake_db = MagicMock()
fake_db.get_session.return_value = {
"system_prompt": stored_system_prompt,
}
class FakeInPlaceCompressAgent:
last_instance = None
def __init__(self, **kwargs):
self.model = kwargs.get("model")
self.platform = kwargs.get("platform")
self.session_id = kwargs.get("session_id", "fake-session")
self._session_db = kwargs.get("session_db")
self._cached_system_prompt = None
self.compression_in_place = False
self._last_compaction_in_place = False
self.context_compressor = SimpleNamespace(
@@ -1118,6 +1129,8 @@ async def test_session_hygiene_forces_in_place_compaction_with_bound_session_db(
def _compress_context(self, messages, *_args, **_kwargs):
assert self.compression_in_place is True
assert self._session_db is fake_db
assert self.platform == "gateway_hygiene"
assert self._cached_system_prompt == stored_system_prompt
self._last_compaction_in_place = True
return ([{"role": "assistant", "content": "compressed in place"}], None)
@@ -1190,6 +1203,7 @@ async def test_session_hygiene_forces_in_place_compaction_with_bound_session_db(
assert result == "ok"
agent = FakeInPlaceCompressAgent.last_instance
assert agent is not None
fake_db.get_session.assert_called_once_with("sess-1")
agent.context_compressor.bind_session_state.assert_called_once_with(fake_db, "sess-1")
# In-place compaction already persisted via archive_and_compact() —
# rewrite_transcript would replace_messages(active_only=False) and DELETE
@@ -1429,7 +1443,7 @@ def _make_progress_runner(monkeypatch, tmp_path, agent_cls, cfg_text):
monkeypatch.setitem(sys.modules, "run_agent", fake_run_agent)
cfg_path = tmp_path / "config.yaml"
cfg_path.write_text(cfg_text)
cfg_path.write_text(cfg_text, encoding="utf-8")
gateway_run = importlib.import_module("gateway.run")
GatewayRunner = gateway_run.GatewayRunner