diff --git a/gateway/run.py b/gateway/run.py index d3992a760e82..e187ff8a3ec9 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -9962,6 +9962,7 @@ async def _prepare_inbound_message_text( event: MessageEvent, source: SessionSource, history: List[Dict[str, Any]], + session_key: Optional[str] = None, ) -> Optional[str]: """Prepare inbound event text for the agent. @@ -9980,10 +9981,10 @@ async def _prepare_inbound_message_text( message_text = event.text or "" _group_sessions_per_user = getattr(self.config, "group_sessions_per_user", True) _thread_sessions_per_user = getattr(self.config, "thread_sessions_per_user", False) - # Use the same helper every other call site uses so the write key here - # matches the consume key at the run_conversation site — even if the - # session store overrides build_session_key's default behavior. - session_key = self._session_key_for_source(source) + # Prefer the already resolved session key from the caller so this write + # key matches the consume key at the run_conversation site. Fall back + # to deriving it here for tests and legacy standalone callers. + session_key = session_key or self._session_key_for_source(source) # Reset only this session's per-call buffer; other sessions may be # concurrently preparing multimodal turns on the same runner. self._consume_pending_native_image_paths(session_key) @@ -10997,6 +10998,7 @@ async def _handle_message_with_agent(self, event, source, _quick_key: str, run_g event=event, source=source, history=history, + session_key=session_key, ) if message_text is None: return @@ -19061,6 +19063,7 @@ def _stream_confirmed_final_delivery( next_message = pending next_message_id = None next_channel_prompt = None + next_session_key = session_key if pending_event is not None: next_source = getattr(pending_event, "source", None) or source if self._is_goal_continuation_event(pending_event) and not self._goal_still_active_for_session(session_id): @@ -19069,10 +19072,24 @@ def _stream_confirmed_final_delivery( session_key or "?", ) return result + # Resolve the follow-up's session key BEFORE preparing the + # inbound text: _prepare_inbound_message_text buffers native + # image paths under the key it is given, and the recursive + # _run_agent below consumes them under next_session_key. + # The write and consume keys must match or the images drop. + try: + next_session_key = self._session_key_for_source(next_source) + except Exception: + logger.debug( + "Queued follow-up session-key resolution failed; reusing %s", + session_key or "?", + exc_info=True, + ) next_message = await self._prepare_inbound_message_text( event=pending_event, source=next_source, history=updated_history, + session_key=next_session_key, ) if next_message is None: return result @@ -19114,7 +19131,7 @@ def _stream_confirmed_final_delivery( history=updated_history, source=next_source, session_id=session_id, - session_key=session_key, + session_key=next_session_key, run_generation=run_generation, _interrupt_depth=_interrupt_depth + 1, event_message_id=next_message_id, diff --git a/tests/gateway/test_native_image_buffer_isolation.py b/tests/gateway/test_native_image_buffer_isolation.py index f8fb2e65a71a..55f96730a598 100644 --- a/tests/gateway/test_native_image_buffer_isolation.py +++ b/tests/gateway/test_native_image_buffer_isolation.py @@ -77,3 +77,20 @@ async def test_native_image_buffer_not_cleared_by_other_sessions_without_images( assert runner._consume_pending_native_image_paths(build_session_key(source_a)) == ["/tmp/a.png"] assert runner._consume_pending_native_image_paths(build_session_key(source_b)) == [] + + +@pytest.mark.asyncio +async def test_native_image_buffer_uses_resolved_session_key_when_provided(): + runner = _make_runner() + source = _source("chat-a") + runner._session_key_for_source = lambda _source: "source-derived-key" + + await runner._prepare_inbound_message_text( + event=_image_event(source, "/tmp/a.png"), + source=source, + history=[], + session_key="canonical-session-key", + ) + + assert runner._consume_pending_native_image_paths("source-derived-key") == [] + assert runner._consume_pending_native_image_paths("canonical-session-key") == ["/tmp/a.png"] diff --git a/tests/gateway/test_queued_native_image_session_key.py b/tests/gateway/test_queued_native_image_session_key.py new file mode 100644 index 000000000000..86ba1620ade9 --- /dev/null +++ b/tests/gateway/test_queued_native_image_session_key.py @@ -0,0 +1,151 @@ +import base64 +import importlib +import sys +import types +from types import SimpleNamespace + +import pytest + +from gateway.config import Platform, PlatformConfig +from gateway.platforms.base import BasePlatformAdapter, MessageEvent, MessageType, SendResult +from gateway.session import SessionSource + + +_ONE_BY_ONE_PNG = base64.b64decode( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO6L2ioAAAAASUVORK5CYII=" +) + + +class CaptureAdapter(BasePlatformAdapter): + def __init__(self): + super().__init__(PlatformConfig(enabled=True, token="***"), Platform.TELEGRAM) + self.sent = [] + self.typing = [] + + async def connect(self) -> bool: + return True + + async def disconnect(self) -> None: + return None + + async def send(self, chat_id, content, reply_to=None, metadata=None) -> SendResult: + self.sent.append( + { + "chat_id": chat_id, + "content": content, + "reply_to": reply_to, + "metadata": metadata, + } + ) + return SendResult(success=True, message_id="sent-1") + + async def send_typing(self, chat_id, metadata=None) -> None: + self.typing.append({"chat_id": chat_id, "metadata": metadata}) + + async def stop_typing(self, chat_id) -> None: + return None + + async def get_chat_info(self, chat_id: str): + return {"id": chat_id} + + +class CaptureQueuedNativeImageAgent: + calls = [] + + def __init__(self, **kwargs): + self.tools = [] + self.tool_progress_callback = kwargs.get("tool_progress_callback") + + def run_conversation(self, message, conversation_history=None, task_id=None): + type(self).calls.append(message) + return { + "final_response": f"done-{len(type(self).calls)}", + "messages": [], + "api_calls": 1, + } + + +def _make_runner(adapter): + gateway_run = importlib.import_module("gateway.run") + runner = object.__new__(gateway_run.GatewayRunner) + runner.adapters = {adapter.platform: adapter} + runner._voice_mode = {} + runner._prefill_messages = [] + runner._ephemeral_system_prompt = "" + runner._reasoning_config = None + runner._provider_routing = {} + runner._fallback_model = None + runner._session_db = None + runner._running_agents = {} + runner._session_run_generation = {} + runner.hooks = SimpleNamespace(loaded_hooks=False) + runner.config = SimpleNamespace( + thread_sessions_per_user=False, + group_sessions_per_user=False, + stt_enabled=False, + ) + runner._model = "openai/gpt-4.1-mini" + runner._base_url = None + runner._decide_image_input_mode = lambda: "native" + return runner + + +@pytest.mark.asyncio +async def test_queued_followup_uses_pending_event_session_key_for_native_images(monkeypatch, tmp_path): + CaptureQueuedNativeImageAgent.calls = [] + + fake_dotenv = types.ModuleType("dotenv") + fake_dotenv.load_dotenv = lambda *args, **kwargs: None + monkeypatch.setitem(sys.modules, "dotenv", fake_dotenv) + + fake_run_agent = types.ModuleType("run_agent") + fake_run_agent.AIAgent = CaptureQueuedNativeImageAgent + monkeypatch.setitem(sys.modules, "run_agent", fake_run_agent) + + gateway_run = importlib.import_module("gateway.run") + monkeypatch.setattr(gateway_run, "_hermes_home", tmp_path) + monkeypatch.setattr(gateway_run, "_resolve_runtime_agent_kwargs", lambda: {"api_key": "***"}) + + adapter = CaptureAdapter() + runner = _make_runner(adapter) + + image_path = tmp_path / "queued-image.png" + image_path.write_bytes(_ONE_BY_ONE_PNG) + + source = SessionSource( + platform=Platform.TELEGRAM, + chat_id="-1001", + chat_type="group", + ) + pending_source = SessionSource( + platform=Platform.TELEGRAM, + chat_id="-1001", + chat_type="group", + thread_id="17585", + ) + + adapter._pending_messages["agent:main:telegram:group:-1001"] = MessageEvent( + text="describe this", + message_type=MessageType.PHOTO, + source=pending_source, + media_urls=[str(image_path)], + media_types=["image/png"], + message_id="queued-1", + ) + + result = await runner._run_agent( + message="hello", + context_prompt="", + history=[], + source=source, + session_id="sess-native-image-followup", + session_key="agent:main:telegram:group:-1001", + ) + + assert result["final_response"] == "done-2" + assert len(CaptureQueuedNativeImageAgent.calls) == 2 + queued_message = CaptureQueuedNativeImageAgent.calls[1] + assert isinstance(queued_message, list) + assert queued_message[0]["type"] == "text" + assert queued_message[0]["text"].startswith("describe this") + assert any(part.get("type") == "image_url" for part in queued_message)