Native image session-key handoff issue and patches

Fixes on record

Merged to main. A stamp counts the issues a fix is linked to.

• #57648 fix(gateway): key native image handoff by the same resolved session key (#48912)

  merged to main

  files: gateway/run.py, tests/gateway/test_native_image_buffer_isolation.py, tests/gateway/test_queued_native_image_session_key.py

  Show the diff

  • 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,

  • 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)) == []
    +
    +
    [redacted:email]
    +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"]

  • 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
    +
    +
    [redacted:email]
    +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)