fix: Fixes a bug that causes intermittent pydantic validation errors when uploading files

The root cause is an unsafe in-memory mutation. The `SaveFilesAsArtifactsPlugin` was saving a direct reference to the message part and then modifying the message list in-place. This created a race condition where downstream code could alter the original part *after* it had been saved as an artifact, leading to a corrupted state.
This CL saves a `copy.copy()` of the artifact, which create a snapshot of the data.
Also Changes the plugin to return a new `types.Content` object instead of modifying the original message in-place

PiperOrigin-RevId: 814308070
This commit is contained in:
George Weale
2025-10-02 13:43:34 -07:00
committed by Copybara-Service
parent f667c7445e
commit e68006386f
5 changed files with 113 additions and 33 deletions
+8 -2
View File
@@ -16,6 +16,7 @@
from __future__ import annotations from __future__ import annotations
import contextlib import contextlib
import copy
from functools import cached_property from functools import cached_property
import logging import logging
import os import os
@@ -300,8 +301,13 @@ class Gemini(BaseLlm):
if not content.parts: if not content.parts:
continue continue
for part in content.parts: for part in content.parts:
_remove_display_name_if_present(part.inline_data) # Create copies to avoid mutating the original objects
_remove_display_name_if_present(part.file_data) if part.inline_data:
part.inline_data = copy.copy(part.inline_data)
_remove_display_name_if_present(part.inline_data)
if part.file_data:
part.file_data = copy.copy(part.file_data)
_remove_display_name_if_present(part.file_data)
# Initialize config if needed # Initialize config if needed
if llm_request.config and llm_request.config.tools: if llm_request.config and llm_request.config.tools:
@@ -14,6 +14,7 @@
from __future__ import annotations from __future__ import annotations
import copy
import logging import logging
from typing import Optional from typing import Optional
@@ -29,14 +30,15 @@ class SaveFilesAsArtifactsPlugin(BasePlugin):
"""A plugin that saves files embedded in user messages as artifacts. """A plugin that saves files embedded in user messages as artifacts.
This is useful to allow users to upload files in the chat experience and have This is useful to allow users to upload files in the chat experience and have
those files available to the agent. those files available to the agent within the current session.
We use Blob.display_name to determine We use Blob.display_name to determine the file name. By default, artifacts are
the file name. Artifacts with the same name will be overwritten. A placeholder session-scoped. For cross-session persistence, prefix the filename with
with the artifact name will be put in place of the embedded file in the user "user:".
message so the model knows where to find the file. You may want to add Artifacts with the same name will be overwritten. A placeholder with the
load_artifacts tool to the agent, or load the artifacts in your own tool to artifact name will be put in place of the embedded file in the user message
use the files. so the model knows where to find the file. You may want to add load_artifacts
tool to the agent, or load the artifacts in your own tool to use the files.
""" """
def __init__(self, name: str = 'save_files_as_artifacts_plugin'): def __init__(self, name: str = 'save_files_as_artifacts_plugin'):
@@ -62,10 +64,14 @@ class SaveFilesAsArtifactsPlugin(BasePlugin):
return user_message return user_message
if not user_message.parts: if not user_message.parts:
return user_message return None
new_parts = []
modified = False
for i, part in enumerate(user_message.parts): for i, part in enumerate(user_message.parts):
if part.inline_data is None: if part.inline_data is None:
new_parts.append(part)
continue continue
try: try:
@@ -77,23 +83,32 @@ class SaveFilesAsArtifactsPlugin(BasePlugin):
f'No display_name found, using generated filename: {file_name}' f'No display_name found, using generated filename: {file_name}'
) )
# Store original filename for display to user/ placeholder
display_name = file_name
# Create a copy to stop mutation of the saved artifact if the original part is modified
await invocation_context.artifact_service.save_artifact( await invocation_context.artifact_service.save_artifact(
app_name=invocation_context.app_name, app_name=invocation_context.app_name,
user_id=invocation_context.user_id, user_id=invocation_context.user_id,
session_id=invocation_context.session.id, session_id=invocation_context.session.id,
filename=file_name, filename=file_name,
artifact=part, artifact=copy.copy(part),
) )
# Replace the inline data with a placeholder text # Replace the inline data with a placeholder text (using the clean name)
user_message.parts[i] = types.Part( new_parts.append(
text=f'[Uploaded Artifact: "{file_name}"]' types.Part(text=f'[Uploaded Artifact: "{display_name}"]')
) )
modified = True
logger.info(f'Successfully saved artifact: {file_name}') logger.info(f'Successfully saved artifact: {file_name}')
except Exception as e: except Exception as e:
logger.error(f'Failed to save artifact for part {i}: {e}') logger.error(f'Failed to save artifact for part {i}: {e}')
# Keep the original part if saving fails # Keep the original part if saving fails
new_parts.append(part)
continue continue
return user_message if modified:
return types.Content(role=user_message.role, parts=new_parts)
else:
return None
@@ -15,6 +15,7 @@
from __future__ import annotations from __future__ import annotations
import json import json
import logging
from typing import Any from typing import Any
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
@@ -27,6 +28,8 @@ if TYPE_CHECKING:
from ..models.llm_request import LlmRequest from ..models.llm_request import LlmRequest
from .tool_context import ToolContext from .tool_context import ToolContext
logger = logging.getLogger('google_adk.' + __name__)
class LoadArtifactsTool(BaseTool): class LoadArtifactsTool(BaseTool):
"""A tool that loads the artifacts and adds them to the session.""" """A tool that loads the artifacts and adds them to the session."""
@@ -96,7 +99,18 @@ class LoadArtifactsTool(BaseTool):
if function_response and function_response.name == 'load_artifacts': if function_response and function_response.name == 'load_artifacts':
artifact_names = function_response.response['artifact_names'] artifact_names = function_response.response['artifact_names']
for artifact_name in artifact_names: for artifact_name in artifact_names:
# Try session-scoped first (default behavior)
artifact = await tool_context.load_artifact(artifact_name) artifact = await tool_context.load_artifact(artifact_name)
# If not found and name doesn't already have user: prefix,
# try cross-session artifacts with user: prefix
if artifact is None and not artifact_name.startswith('user:'):
prefixed_name = f'user:{artifact_name}'
artifact = await tool_context.load_artifact(prefixed_name)
if artifact is None:
logger.warning('Artifact "%s" not found, skipping', artifact_name)
continue
llm_request.contents.append( llm_request.contents.append(
types.Content( types.Content(
role='user', role='user',
@@ -277,3 +277,48 @@ async def test_list_versions(service_type):
) )
assert response_versions == list(range(4)) assert response_versions == list(range(4))
@pytest.mark.asyncio
async def test_list_keys_preserves_user_prefix():
"""Tests that list_artifact_keys preserves 'user:' prefix in returned names."""
artifact_service = InMemoryArtifactService()
artifact = types.Part.from_bytes(data=b"test_data", mime_type="text/plain")
app_name = "app0"
user_id = "user0"
session_id = "123"
# Save artifacts with "user:" prefix (cross-session artifacts)
await artifact_service.save_artifact(
app_name=app_name,
user_id=user_id,
session_id=session_id,
filename="user:document.pdf",
artifact=artifact,
)
await artifact_service.save_artifact(
app_name=app_name,
user_id=user_id,
session_id=session_id,
filename="user:image.png",
artifact=artifact,
)
# Save session-scoped artifact without prefix
await artifact_service.save_artifact(
app_name=app_name,
user_id=user_id,
session_id=session_id,
filename="session_file.txt",
artifact=artifact,
)
# List artifacts should return names with "user:" prefix for user-scoped artifacts
artifact_keys = await artifact_service.list_artifact_keys(
app_name=app_name, user_id=user_id, session_id=session_id
)
# Should contain prefixed names and session file
expected_keys = ["user:document.pdf", "user:image.png", "session_file.txt"]
assert sorted(artifact_keys) == sorted(expected_keys)
@@ -57,7 +57,7 @@ class TestSaveFilesAsArtifactsPlugin:
invocation_context=self.mock_context, user_message=user_message invocation_context=self.mock_context, user_message=user_message
) )
# Verify artifact was saved with correct filename # Verify artifact was saved with correct filename (session-scoped by default)
self.mock_context.artifact_service.save_artifact.assert_called_once_with( self.mock_context.artifact_service.save_artifact.assert_called_once_with(
app_name="test_app", app_name="test_app",
user_id="test_user", user_id="test_user",
@@ -66,7 +66,7 @@ class TestSaveFilesAsArtifactsPlugin:
artifact=original_part, artifact=original_part,
) )
# Verify message was modified with placeholder # Verify message was modified with placeholder (clean name)
assert result.parts[0].text == '[Uploaded Artifact: "test_document.pdf"]' assert result.parts[0].text == '[Uploaded Artifact: "test_document.pdf"]'
@pytest.mark.asyncio @pytest.mark.asyncio
@@ -85,7 +85,7 @@ class TestSaveFilesAsArtifactsPlugin:
invocation_context=self.mock_context, user_message=user_message invocation_context=self.mock_context, user_message=user_message
) )
# Verify artifact was saved with generated filename # Verify artifact was saved with generated filename (session-scoped by default)
expected_filename = "artifact_test_invocation_123_0" expected_filename = "artifact_test_invocation_123_0"
self.mock_context.artifact_service.save_artifact.assert_called_once_with( self.mock_context.artifact_service.save_artifact.assert_called_once_with(
app_name="test_app", app_name="test_app",
@@ -95,8 +95,12 @@ class TestSaveFilesAsArtifactsPlugin:
artifact=original_part, artifact=original_part,
) )
# Verify message was modified with generated filename # Verify message was modified with generated filename (clean name)
assert result.parts[0].text == f'[Uploaded Artifact: "{expected_filename}"]' generated_display_name = "artifact_test_invocation_123_0"
assert (
result.parts[0].text
== f'[Uploaded Artifact: "{generated_display_name}"]'
)
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_multiple_files_in_message(self): async def test_multiple_files_in_message(self):
@@ -138,7 +142,7 @@ class TestSaveFilesAsArtifactsPlugin:
) )
assert second_call[1]["filename"] == "file2.jpg" assert second_call[1]["filename"] == "file2.jpg"
# Verify message parts were modified correctly # Verify message parts were modified correctly (clean names)
assert result.parts[0].text == '[Uploaded Artifact: "file1.txt"]' assert result.parts[0].text == '[Uploaded Artifact: "file1.txt"]'
assert result.parts[1].text == "Some text between files" # Unchanged assert result.parts[1].text == "Some text between files" # Unchanged
assert result.parts[2].text == '[Uploaded Artifact: "file2.jpg"]' assert result.parts[2].text == '[Uploaded Artifact: "file2.jpg"]'
@@ -174,9 +178,8 @@ class TestSaveFilesAsArtifactsPlugin:
invocation_context=self.mock_context, user_message=user_message invocation_context=self.mock_context, user_message=user_message
) )
# Should return original message unchanged # Should return None to proceed with original message
assert result == user_message assert result is None
assert result.parts == []
# Should not try to save any artifacts # Should not try to save any artifacts
self.mock_context.artifact_service.save_artifact.assert_not_called() self.mock_context.artifact_service.save_artifact.assert_not_called()
@@ -193,10 +196,8 @@ class TestSaveFilesAsArtifactsPlugin:
invocation_context=self.mock_context, user_message=user_message invocation_context=self.mock_context, user_message=user_message
) )
# Should return original message unchanged # Should return None to proceed with original message
assert result == user_message assert result is None
assert result.parts[0].text == "Hello world"
assert result.parts[1].text == "No files here"
# Should not try to save any artifacts # Should not try to save any artifacts
self.mock_context.artifact_service.save_artifact.assert_not_called() self.mock_context.artifact_service.save_artifact.assert_not_called()
@@ -221,9 +222,8 @@ class TestSaveFilesAsArtifactsPlugin:
invocation_context=self.mock_context, user_message=user_message invocation_context=self.mock_context, user_message=user_message
) )
# Should preserve original part when saving fails # Should return None when saving fails (no modifications made)
assert result.parts[0] == original_part assert result is None
assert result.parts[0].inline_data == inline_data
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_mixed_success_and_failure(self): async def test_mixed_success_and_failure(self):
@@ -264,7 +264,7 @@ class TestSaveFilesAsArtifactsPlugin:
invocation_context=self.mock_context, user_message=user_message invocation_context=self.mock_context, user_message=user_message
) )
# First file should be replaced with placeholder # First file should be replaced with placeholder (clean name)
assert result.parts[0].text == '[Uploaded Artifact: "success.pdf"]' assert result.parts[0].text == '[Uploaded Artifact: "success.pdf"]'
# Second file should remain unchanged due to failure # Second file should remain unchanged due to failure
@@ -287,7 +287,7 @@ class TestSaveFilesAsArtifactsPlugin:
invocation_context=self.mock_context, user_message=user_message invocation_context=self.mock_context, user_message=user_message
) )
# Verify exact format of placeholder text # Verify exact format of placeholder text (clean name)
expected_text = '[Uploaded Artifact: "test file with spaces.docx"]' expected_text = '[Uploaded Artifact: "test file with spaces.docx"]'
assert result.parts[0].text == expected_text assert result.parts[0].text == expected_text