diff --git a/gateway/run.py b/gateway/run.py index 94e8f9e43a1..3927b7f371a 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -326,8 +326,7 @@ def _non_conversational_metadata( def _seed_hygiene_system_prompt( agent: Any, - session_db: Any, - session_id: str, + session_row: Optional[Dict[str, Any]], ) -> bool: """Keep gateway hygiene from rebuilding a live session's system prompt. @@ -340,22 +339,10 @@ def _seed_hygiene_system_prompt( 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, - ) + 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 agent._cached_system_prompt = stored_prompt return bool(stored_prompt) @@ -13671,6 +13658,21 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew ] if len(_hyg_msgs) >= 4: + try: + _hyg_session_row = await self._session_db.get_session( + session_entry.session_id + ) + except Exception as exc: + _hyg_session_row = None + 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_entry.session_id, + exc, + exc_info=True, + ) _hyg_session_db = getattr(self._session_db, "_db", self._session_db) _hyg_agent = AIAgent( **_hyg_runtime, @@ -13684,8 +13686,7 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew ) _seed_hygiene_system_prompt( _hyg_agent, - _hyg_session_db, - session_entry.session_id, + _hyg_session_row, ) # If compression must rebuild instead of retaining # the cached prompt, make the persisted result diff --git a/tests/gateway/test_session_hygiene.py b/tests/gateway/test_session_hygiene.py index a586a26263f..22847aa4f28 100644 --- a/tests/gateway/test_session_hygiene.py +++ b/tests/gateway/test_session_hygiene.py @@ -1101,9 +1101,14 @@ async def test_session_hygiene_forces_in_place_compaction_with_bound_session_db( "" ) fake_db = MagicMock() - fake_db.get_session.return_value = { - "system_prompt": stored_system_prompt, - } + async_session_db = SimpleNamespace( + _db=fake_db, + get_session=AsyncMock( + return_value={ + "system_prompt": stored_system_prompt, + } + ), + ) class FakeInPlaceCompressAgent: last_instance = None @@ -1165,7 +1170,7 @@ async def test_session_hygiene_forces_in_place_compaction_with_bound_session_db( runner._running_agents = {} runner._pending_messages = {} runner._pending_approvals = {} - runner._session_db = SimpleNamespace(_db=fake_db) + runner._session_db = async_session_db runner._is_user_authorized = lambda _source: True runner._set_session_env = lambda _context: None runner._run_agent = AsyncMock( @@ -1203,7 +1208,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") + async_session_db.get_session.assert_awaited_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