From 73c8d5a1e7ecdf04c3f6a3ad31e4276fc16f4b69 Mon Sep 17 00:00:00 2001 From: kshitijk4poor <82637225+kshitijk4poor@users.noreply.github.com> Date: Fri, 26 Jun 2026 00:03:52 +0530 Subject: [PATCH] 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. --- gateway/run.py | 9 ++- gateway/slash_commands.py | 8 +-- tests/gateway/test_compress_command.py | 58 ++++++++++++++++++++ tests/gateway/test_compress_focus.py | 1 + tests/gateway/test_compress_plugin_engine.py | 1 + 5 files changed, 65 insertions(+), 12 deletions(-) diff --git a/gateway/run.py b/gateway/run.py index 33f60f79c68..b1ee482b794 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -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 diff --git a/gateway/slash_commands.py b/gateway/slash_commands.py index 0b1775d15b8..e0b655c0a7d 100644 --- a/gateway/slash_commands.py +++ b/gateway/slash_commands.py @@ -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 diff --git a/tests/gateway/test_compress_command.py b/tests/gateway/test_compress_command.py index 95211e97722..9c47e76db23 100644 --- a/tests/gateway/test_compress_command.py +++ b/tests/gateway/test_compress_command.py @@ -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() diff --git a/tests/gateway/test_compress_focus.py b/tests/gateway/test_compress_focus.py index 597185e5792..100cba800b8 100644 --- a/tests/gateway/test_compress_focus.py +++ b/tests/gateway/test_compress_focus.py @@ -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 diff --git a/tests/gateway/test_compress_plugin_engine.py b/tests/gateway/test_compress_plugin_engine.py index 4604e772373..79eef5551f0 100644 --- a/tests/gateway/test_compress_plugin_engine.py +++ b/tests/gateway/test_compress_plugin_engine.py @@ -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