From 5a3e4e11bdbc84900b231602adb3fa31b2d7f2f7 Mon Sep 17 00:00:00 2001
From: Tongzhou Jiang <tzjiang@google.com>
Date: Fri, 20 Feb 2026 14:35:42 -0800
Subject: [PATCH] fix: save artifact in streaming agent run with events when
 multiturn

PiperOrigin-RevId: 873078156
---
 vertexai/agent_engines/templates/adk.py       | 21 ++++++++++++--
 .../reasoning_engines/templates/adk.py        | 28 +++++++++++++++++--
 2 files changed, 44 insertions(+), 5 deletions(-)

diff --git a/vertexai/agent_engines/templates/adk.py b/vertexai/agent_engines/templates/adk.py
index 3cb2e0bceb..f52b46f1ea 100644
--- a/vertexai/agent_engines/templates/adk.py
+++ b/vertexai/agent_engines/templates/adk.py
@@ -683,6 +683,18 @@ async def _init_session(
         if request.events:
             for event in request.events:
                 await session_service.append_event(session, Event(**event))
+        if request.artifacts:
+            await self._save_artifacts(session.id, artifact_service, request)
+        return session
+
+    async def _save_artifacts(
+        self,
+        session_id: str,
+        artifact_service: "BaseArtifactService",
+        request: _StreamRunRequest,
+    ):
+        """Saves the artifacts."""
+        app = self._tmpl_attrs.get("app")
         if request.artifacts:
             for artifact in request.artifacts:
                 artifact = _Artifact(**artifact)
@@ -693,7 +705,7 @@ async def _init_session(
                     saved_version = await artifact_service.save_artifact(
                         app_name=app.name if app else self._tmpl_attrs.get("app_name"),
                         user_id=request.user_id,
-                        session_id=session.id,
+                        session_id=session_id,
                         filename=artifact.file_name,
                         artifact=version_data.data,
                     )
@@ -707,7 +719,6 @@ async def _init_session(
                             saved_version,
                             version_data.version,
                         )
-        return session
 
     async def _convert_response_events(
         self,
@@ -1209,6 +1220,12 @@ async def streaming_agent_run_with_events(self, request_json: str):
                     user_id=request.user_id,
                     session_id=request.session_id,
                 )
+                if session:
+                    await self._save_artifacts(
+                        session_id=request.session_id,
+                        artifact_service=artifact_service,
+                        request=request,
+                    )
             except ClientError:
                 pass
             if not session:
diff --git a/vertexai/preview/reasoning_engines/templates/adk.py b/vertexai/preview/reasoning_engines/templates/adk.py
index 1e55bf40f7..8d970d66f5 100644
--- a/vertexai/preview/reasoning_engines/templates/adk.py
+++ b/vertexai/preview/reasoning_engines/templates/adk.py
@@ -616,6 +616,23 @@ async def _init_session(
         if request.events:
             for event in request.events:
                 await session_service.append_event(session, Event(**event))
+        if request.artifacts:
+            await self._save_artifacts(
+                session_id=session.id,
+                artifact_service=artifact_service,
+                request=request,
+            )
+
+        return session
+
+    async def _save_artifacts(
+        self,
+        session_id: str,
+        artifact_service: "BaseArtifactService",
+        request: _StreamRunRequest,
+    ):
+        """Saves the artifacts."""
+        app = self._tmpl_attrs.get("app")
         if request.artifacts:
             for artifact in request.artifacts:
                 artifact = _Artifact(**artifact)
@@ -624,9 +641,9 @@ async def _init_session(
                 ):
                     version_data = _ArtifactVersion(**version_data)
                     saved_version = await artifact_service.save_artifact(
-                        app_name=self._tmpl_attrs.get("app_name"),
+                        app_name=app.name if app else self._tmpl_attrs.get("app_name"),
                         user_id=request.user_id,
-                        session_id=session.id,
+                        session_id=session_id,
                         filename=artifact.file_name,
                         artifact=version_data.data,
                     )
@@ -640,7 +657,6 @@ async def _init_session(
                             saved_version,
                             version_data.version,
                         )
-        return session
 
     async def _convert_response_events(
         self,
@@ -1043,6 +1059,12 @@ async def _invoke_agent_async():
                         user_id=request.user_id,
                         session_id=request.session_id,
                     )
+                    if session:
+                        await self._save_artifacts(
+                            session_id=request.session_id,
+                            artifact_service=artifact_service,
+                            request=request,
+                        )
                 except ClientError:
                     pass
                 if not session: