Skip to content

Commit f635214

Browse files
committed
fix(memory): do not leak temp file when RAG upload fails
add_session_to_memory created a NamedTemporaryFile(delete=False) and only removed it on the success path. A missing rag-resources configuration or any rag.upload_file failure left the temp file on disk permanently. Validate rag resources before the temp file is created, and wrap the upload loop in try/finally so the file is always removed.
1 parent e6df097 commit f635214

2 files changed

Lines changed: 84 additions & 15 deletions

File tree

src/google/adk/memory/vertex_ai_rag_memory_service.py

Lines changed: 16 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,9 @@ def __init__(
125125

126126
@override
127127
async def add_session_to_memory(self, session: Session) -> None:
128+
if not self._vertex_rag_store.rag_resources:
129+
raise ValueError("Rag resources must be set.")
130+
128131
with tempfile.NamedTemporaryFile(
129132
mode="w", delete=False, suffix=".txt"
130133
) as temp_file:
@@ -150,23 +153,21 @@ async def add_session_to_memory(self, session: Session) -> None:
150153
temp_file.write(output_string)
151154
temp_file_path = temp_file.name
152155

153-
if not self._vertex_rag_store.rag_resources:
154-
raise ValueError("Rag resources must be set.")
155-
156156
from ..dependencies.vertexai import rag
157157

158-
for rag_resource in self._vertex_rag_store.rag_resources:
159-
rag.upload_file(
160-
corpus_name=rag_resource.rag_corpus,
161-
path=temp_file_path,
162-
# this is the temp workaround as upload file does not support
163-
# adding metadata, thus use display_name to store the session info.
164-
display_name=_build_source_display_name(
165-
session.app_name, session.user_id, session.id
166-
),
167-
)
168-
169-
os.remove(temp_file_path)
158+
try:
159+
for rag_resource in self._vertex_rag_store.rag_resources:
160+
rag.upload_file(
161+
corpus_name=rag_resource.rag_corpus,
162+
path=temp_file_path,
163+
# this is the temp workaround as upload file does not support
164+
# adding metadata, thus use display_name to store the session info.
165+
display_name=_build_source_display_name(
166+
session.app_name, session.user_id, session.id
167+
),
168+
)
169+
finally:
170+
os.remove(temp_file_path)
170171

171172
@override
172173
async def search_memory(

tests/unittests/memory/test_vertex_ai_rag_memory_service.py

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,8 @@
1313
# limitations under the License.
1414

1515
import json
16+
import os
17+
import tempfile
1618
from types import SimpleNamespace
1719

1820
from google.adk.events.event import Event
@@ -31,6 +33,72 @@ def _rag_context(source_display_name: str, text: str) -> SimpleNamespace:
3133
)
3234

3335

36+
def _sample_session() -> Session:
37+
return Session(
38+
app_name="demo",
39+
user_id="alice",
40+
id="session-1",
41+
last_update_time=1,
42+
events=[
43+
Event(
44+
id="event-1",
45+
author="user",
46+
timestamp=1,
47+
content=types.Content(parts=[types.Part(text="hello")]),
48+
)
49+
],
50+
)
51+
52+
53+
def _track_temp_files(mocker) -> list[str]:
54+
"""Records the paths of every NamedTemporaryFile created."""
55+
created_paths: list[str] = []
56+
real_named_temp_file = tempfile.NamedTemporaryFile
57+
58+
def _spy(*args, **kwargs):
59+
temp_file = real_named_temp_file(*args, **kwargs)
60+
created_paths.append(temp_file.name)
61+
return temp_file
62+
63+
mocker.patch(
64+
"google.adk.memory.vertex_ai_rag_memory_service.tempfile.NamedTemporaryFile",
65+
side_effect=_spy,
66+
)
67+
return created_paths
68+
69+
70+
@pytest.mark.asyncio
71+
async def test_add_session_removes_temp_file_when_upload_fails(mocker):
72+
"""The temp file must not leak when rag.upload_file raises."""
73+
memory_service = VertexAiRagMemoryService(rag_corpus="unused")
74+
created_paths = _track_temp_files(mocker)
75+
fake_rag = SimpleNamespace(
76+
upload_file=mocker.Mock(side_effect=RuntimeError("upload boom"))
77+
)
78+
mocker.patch("google.adk.dependencies.vertexai.rag", fake_rag)
79+
80+
with pytest.raises(RuntimeError, match="upload boom"):
81+
await memory_service.add_session_to_memory(_sample_session())
82+
83+
assert created_paths, "expected a temp file to have been created"
84+
assert not os.path.exists(created_paths[0])
85+
86+
87+
@pytest.mark.asyncio
88+
async def test_add_session_does_not_create_temp_file_without_rag_resources(
89+
mocker,
90+
):
91+
"""Validation happens before the temp file is created, so nothing leaks."""
92+
memory_service = VertexAiRagMemoryService(rag_corpus="unused")
93+
memory_service._vertex_rag_store.rag_resources = None
94+
created_paths = _track_temp_files(mocker)
95+
96+
with pytest.raises(ValueError, match="Rag resources must be set."):
97+
await memory_service.add_session_to_memory(_sample_session())
98+
99+
assert created_paths == []
100+
101+
34102
@pytest.mark.asyncio
35103
async def test_search_memory_rejects_ambiguous_legacy_display_names(mocker):
36104
"""Ensures dotted user IDs cannot match another user's legacy memory."""

0 commit comments

Comments
 (0)