From 0b422559f3b21b6dc2c94039b70df8a5b8a103cf Mon Sep 17 00:00:00 2001 From: kshitijk4poor <82637225+kshitijk4poor@users.noreply.github.com> Date: Mon, 13 Jul 2026 23:52:32 +0530 Subject: [PATCH] fix(session): preserve clean multimodal persistence override --- agent/conversation_loop.py | 4 ++-- agent/turn_context.py | 4 ++-- gateway/run.py | 4 ++-- run_agent.py | 15 +++++++++------ tests/run_agent/test_run_agent.py | 25 +++++++++++++++++++++++++ 5 files changed, 40 insertions(+), 12 deletions(-) diff --git a/agent/conversation_loop.py b/agent/conversation_loop.py index 75c07540d885..16083c7cab6d 100644 --- a/agent/conversation_loop.py +++ b/agent/conversation_loop.py @@ -522,12 +522,12 @@ def _sync_failover_system_message(agent, api_messages, active_system_prompt): def run_conversation( agent, - user_message: str, + user_message: Any, system_message: str = None, conversation_history: List[Dict[str, Any]] = None, task_id: str = None, stream_callback: Optional[callable] = None, - persist_user_message: Optional[str] = None, + persist_user_message: Optional[Any] = None, persist_user_timestamp: Optional[float] = None, moa_config: Optional[dict[str, Any]] = None, ) -> Dict[str, Any]: diff --git a/agent/turn_context.py b/agent/turn_context.py index cf1cdff47517..ea150ff30a7c 100644 --- a/agent/turn_context.py +++ b/agent/turn_context.py @@ -118,12 +118,12 @@ class TurnContext: def build_turn_context( agent, - user_message: str, + user_message: Any, system_message: Optional[str], conversation_history: Optional[List[Dict[str, Any]]], task_id: Optional[str], stream_callback, - persist_user_message: Optional[str], + persist_user_message: Optional[Any], persist_user_timestamp: Optional[float] = None, *, restore_or_build_system_prompt, diff --git a/gateway/run.py b/gateway/run.py index d5b3fbff4743..9b19190338d4 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -17123,7 +17123,7 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew event_message_id: Optional[str] = None, channel_prompt: Optional[str] = None, moa_config: Optional[dict] = None, - persist_user_message: Optional[str] = None, + persist_user_message: Optional[Any] = None, persist_user_timestamp: Optional[float] = None, ) -> Dict[str, Any]: """Profile-scoping wrapper around the agent run. @@ -17184,7 +17184,7 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew event_message_id: Optional[str] = None, channel_prompt: Optional[str] = None, moa_config: Optional[dict] = None, - persist_user_message: Optional[str] = None, + persist_user_message: Optional[Any] = None, persist_user_timestamp: Optional[float] = None, ) -> Dict[str, Any]: """ diff --git a/run_agent.py b/run_agent.py index 2d710d77e2d9..82e93a5dc04a 100644 --- a/run_agent.py +++ b/run_agent.py @@ -1871,11 +1871,14 @@ class AIAgent: content = msg.get("content") _row_timestamp = msg.get("timestamp") # Apply the persist override to THIS row's written values only - # (never to the live dict). Match the original guard: text-only - # content is replaced; multimodal (list) content is left intact - # so image/audio blocks aren't clobbered by the text override. + # (never to the live dict). A multimodal override is a complete + # clean replacement for an API-local noted payload. Preserve the + # historical text-only guard for a list payload, though: a plain + # text override must not erase its image/audio transcript summary. if _ov_idx == _msg_idx and msg.get("role") == "user": - if _ov_content is not None and not isinstance(content, list): + if _ov_content is not None and ( + not isinstance(content, list) or isinstance(_ov_content, list) + ): content = _ov_content if _ov_timestamp is not None: _row_timestamp = _ov_timestamp @@ -5798,12 +5801,12 @@ class AIAgent: def run_conversation( self, - user_message: str, + user_message: Any, system_message: str = None, conversation_history: List[Dict[str, Any]] = None, task_id: str = None, stream_callback: Optional[callable] = None, - persist_user_message: Optional[str] = None, + persist_user_message: Optional[Any] = None, persist_user_timestamp: Optional[float] = None, moa_config: Optional[dict[str, Any]] = None, ) -> Dict[str, Any]: diff --git a/tests/run_agent/test_run_agent.py b/tests/run_agent/test_run_agent.py index 396657ea4e85..07e54a50ded4 100644 --- a/tests/run_agent/test_run_agent.py +++ b/tests/run_agent/test_run_agent.py @@ -143,6 +143,31 @@ def test_persist_user_message_override_preserves_multimodal_turns(agent): assert messages == [{"role": "user", "content": multimodal_content}] +def test_flush_persist_override_replaces_api_local_multimodal_note(agent): + """A note-added multimodal API payload stores the original clean content.""" + clean_content = [ + {"type": "text", "text": "Describe this screenshot"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}}, + ] + api_content = [ + {"type": "text", "text": "[MODEL SWITCH NOTE]\n\nDescribe this screenshot"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}}, + ] + agent._session_db = MagicMock() + agent._session_db_created = True + agent.session_id = "session-123" + agent._last_flushed_db_idx = 0 + agent._persist_user_message_idx = 0 + agent._persist_user_message_override = clean_content + agent._persist_user_message_timestamp = None + + agent._flush_messages_to_session_db([{"role": "user", "content": api_content}], []) + + db_write = agent._session_db.append_message.call_args.kwargs + assert db_write["content"] == "Describe this screenshot\n[screenshot]" + assert api_content[0]["text"] == "[MODEL SWITCH NOTE]\n\nDescribe this screenshot" + + @pytest.fixture() def agent_with_memory_tool(): """Agent whose valid_tool_names includes 'memory'."""