fix: save artifact in streaming agent run with events when multiturn

Tongzhou-Jiang · copybara-github · commit e9f1c88d8a74 · 2026-03-03T11:45:34.000-08:00
PiperOrigin-RevId: 878043339
diff --git 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
@@ -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:
