fix: use self._session_db directly + add regression test

- Replace getattr(self.session_store, '_db', None) with self._session_db
  (the GatewayRunner's own SessionDB, consistent with existing usage in
  slash_commands.py L240/L499).
- Remove verbose comment referencing a branch name as an issue number.
- Update stale comment in run.py that said 'today it has no session_db'.
- Add regression test verifying session_db is passed and rotated session
  is persisted (adapted from #51624 by @LeonSGP43).
- Add _session_db=None to _make_runner fixtures in test_compress_command,
  test_compress_focus, and test_compress_plugin_engine.
This commit is contained in:
kshitijk4poor 2026-06-26 00:03:52 +05:30 committed by kshitij
parent 1a38a8ff7d
commit 73c8d5a1e7
5 changed files with 65 additions and 12 deletions

View file

@ -9542,7 +9542,6 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
]
if len(_hyg_msgs) >= 4:
_hyg_session_db = getattr(self.session_store, "_db", None)
_hyg_agent = AIAgent(
**_hyg_runtime,
model=_hyg_model,
@ -9551,15 +9550,15 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
skip_memory=True,
enabled_toolsets=["memory"],
session_id=session_entry.session_id,
session_db=_hyg_session_db,
session_db=self._session_db,
)
try:
# The hygiene agent rotates the session
# forward to a continuation id that becomes
# the gateway session's live row. It must
# never finalize on close() (today it has no
# session_db so close() no-ops, but this
# guards a future where one is wired in).
# never finalize on close() — close() would
# end the newly rotated session the gateway
# entry now points at.
_hyg_agent._end_session_on_close = False
_hyg_agent._print_fn = lambda *a, **kw: None

View file

@ -2810,12 +2810,6 @@ class GatewaySlashCommandsMixin:
partial = False
head = msgs
# Pass session_db so the temp agent can persist its session
# rotation to the database. Without it, compress_context computes
# the compressed messages in memory but never writes them to the
# session store — the next user message reloads the full original
# transcript and the apparent compression is lost (#fix/compress-gateway-persistence).
_tmp_session_db = getattr(self.session_store, "_db", None)
tmp_agent = AIAgent(
**runtime_kwargs,
model=model,
@ -2824,7 +2818,7 @@ class GatewaySlashCommandsMixin:
skip_memory=True,
enabled_toolsets=["memory"],
session_id=session_entry.session_id,
session_db=_tmp_session_db,
session_db=self._session_db,
)
try:
tmp_agent._print_fn = lambda *a, **kw: None

View file

@ -54,6 +54,7 @@ def _make_runner(history: list[dict[str, str]]):
runner.session_store.rewrite_transcript = MagicMock()
runner.session_store.update_session = MagicMock()
runner.session_store._save = MagicMock()
runner._session_db = None
return runner
@ -247,3 +248,60 @@ async def test_compress_command_surfaces_aux_model_failure_even_when_recovered()
assert "intact" in result
agent_instance.shutdown_memory_provider.assert_called_once()
agent_instance.close.assert_called_once()
@pytest.mark.asyncio
async def test_compress_command_passes_session_db_and_persists_rotated_session():
"""session_db must be wired into the /compress temp agent so that
_compress_context can actually rotate the session and persist the
compressed transcript without it compression is a silent no-op."""
history = _make_history()
compressed = [
history[0],
{"role": "assistant", "content": "compressed summary"},
history[-1],
]
runner = _make_runner(history)
runner._session_db = object()
agent_instance = MagicMock()
agent_instance.shutdown_memory_provider = MagicMock()
agent_instance.close = MagicMock()
agent_instance._cached_system_prompt = ""
agent_instance.tools = None
agent_instance.context_compressor.has_content_to_compress.return_value = True
agent_instance.compression_in_place = False
agent_instance.session_id = "sess-1"
def _compress(messages, *_args, **_kwargs):
agent_instance.session_id = "sess-2"
return compressed, ""
agent_instance._compress_context.side_effect = _compress
def _estimate(messages, **_kwargs):
if messages == history:
return 100
if messages == compressed:
return 60
raise AssertionError(f"unexpected transcript: {messages!r}")
with (
patch("gateway.run._resolve_runtime_agent_kwargs", return_value={"api_key": "***"}),
patch("gateway.run._resolve_gateway_model", return_value="test-model"),
patch("run_agent.AIAgent", return_value=agent_instance) as mock_agent_cls,
patch("agent.model_metadata.estimate_request_tokens_rough", side_effect=_estimate),
):
result = await runner._handle_compress_command(_make_event())
assert "Compressed:" in result
mock_agent_cls.assert_called_once()
assert mock_agent_cls.call_args.kwargs["session_db"] is runner._session_db
runner.session_store._save.assert_called_once()
runner.session_store.rewrite_transcript.assert_called_once_with(
"sess-2", compressed
)
runner.session_store.update_session.assert_called_once_with(
build_session_key(_make_source()), last_prompt_tokens=0
)
agent_instance.shutdown_memory_provider.assert_called_once()
agent_instance.close.assert_called_once()

View file

@ -54,6 +54,7 @@ def _make_runner(history: list[dict[str, str]]):
runner.session_store.rewrite_transcript = MagicMock()
runner.session_store.update_session = MagicMock()
runner.session_store._save = MagicMock()
runner._session_db = None
return runner

View file

@ -101,6 +101,7 @@ def _make_runner(history: list[dict[str, str]]):
runner.session_store.rewrite_transcript = MagicMock()
runner.session_store.update_session = MagicMock()
runner.session_store._save = MagicMock()
runner._session_db = None
return runner