From e88039648813dadc42adf884e87d62cd877ff4fb Mon Sep 17 00:00:00 2001
From: tt-a1i <53142663+tt-a1i@users.noreply.github.com>
Date: Fri, 19 Jun 2026 16:18:47 +0800
Subject: [PATCH] fix(gateway): key native image handoff by session

---
 gateway/run.py                                  | 11 +++++++----
 .../test_native_image_buffer_isolation.py       | 17 +++++++++++++++++
 2 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/gateway/run.py b/gateway/run.py
index d3992a760..17ccd4c1c 100644
--- a/gateway/run.py
+++ b/gateway/run.py
@@ -9962,6 +9962,7 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
         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 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
         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 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
             event=event,
             source=source,
             history=history,
+            session_key=session_key,
         )
         if message_text is None:
             return
@@ -19073,6 +19075,7 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
                         event=pending_event,
                         source=next_source,
                         history=updated_history,
+                        session_key=session_key,
                     )
                     if next_message is None:
                         return result
diff --git a/tests/gateway/test_native_image_buffer_isolation.py b/tests/gateway/test_native_image_buffer_isolation.py
index f8fb2e65a..55f96730a 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"]