mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
fix: Harden YAML builder tmp save/cleanup
- Add path-safe helpers so all builder filesystem operations stay under <agents_dir>/<app_name> and reject traversal/invalid upload paths.
- Rework /builder/save to support tmp=true writes under <app>/tmp/<app>, promote tmp → app root on final save (preserving tools.py/tools/), then clean up tmp on success.
- Simplify /builder/app/{app_name}/cancel to best-effort delete tmp; update GET /builder/app/{app_name}?tmp=true to auto-recreate tmp from the app root and safely serve requested files.
Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 852366567
This commit is contained in:
George Weale
committed by
Copybara-Service
Copybara-Service
parent
3ec7ae3b8d
commit
6f259f08b3
@@ -17,6 +17,7 @@ import json
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
import signal
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
@@ -31,6 +32,7 @@ from google.adk.agents.base_agent import BaseAgent
|
||||
from google.adk.agents.run_config import RunConfig
|
||||
from google.adk.apps.app import App
|
||||
from google.adk.artifacts.base_artifact_service import ArtifactVersion
|
||||
from google.adk.cli import fast_api as fast_api_module
|
||||
from google.adk.cli.fast_api import get_fast_api_app
|
||||
from google.adk.errors.input_validation_error import InputValidationError
|
||||
from google.adk.evaluation.eval_case import EvalCase
|
||||
@@ -414,29 +416,41 @@ def test_app(
|
||||
|
||||
# Patch multiple services and signal handlers
|
||||
with (
|
||||
patch("signal.signal", return_value=None),
|
||||
patch(
|
||||
"google.adk.cli.fast_api.create_session_service_from_options",
|
||||
patch.object(signal, "signal", autospec=True, return_value=None),
|
||||
patch.object(
|
||||
fast_api_module,
|
||||
"create_session_service_from_options",
|
||||
autospec=True,
|
||||
return_value=mock_session_service,
|
||||
),
|
||||
patch(
|
||||
"google.adk.cli.fast_api.create_artifact_service_from_options",
|
||||
patch.object(
|
||||
fast_api_module,
|
||||
"create_artifact_service_from_options",
|
||||
autospec=True,
|
||||
return_value=mock_artifact_service,
|
||||
),
|
||||
patch(
|
||||
"google.adk.cli.fast_api.create_memory_service_from_options",
|
||||
patch.object(
|
||||
fast_api_module,
|
||||
"create_memory_service_from_options",
|
||||
autospec=True,
|
||||
return_value=mock_memory_service,
|
||||
),
|
||||
patch(
|
||||
"google.adk.cli.fast_api.AgentLoader",
|
||||
patch.object(
|
||||
fast_api_module,
|
||||
"AgentLoader",
|
||||
autospec=True,
|
||||
return_value=mock_agent_loader,
|
||||
),
|
||||
patch(
|
||||
"google.adk.cli.fast_api.LocalEvalSetsManager",
|
||||
patch.object(
|
||||
fast_api_module,
|
||||
"LocalEvalSetsManager",
|
||||
autospec=True,
|
||||
return_value=mock_eval_sets_manager,
|
||||
),
|
||||
patch(
|
||||
"google.adk.cli.fast_api.LocalEvalSetResultsManager",
|
||||
patch.object(
|
||||
fast_api_module,
|
||||
"LocalEvalSetResultsManager",
|
||||
autospec=True,
|
||||
return_value=mock_eval_set_results_manager,
|
||||
),
|
||||
):
|
||||
@@ -459,6 +473,70 @@ def test_app(
|
||||
return client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def builder_test_client(
|
||||
tmp_path,
|
||||
mock_session_service,
|
||||
mock_artifact_service,
|
||||
mock_memory_service,
|
||||
mock_agent_loader,
|
||||
mock_eval_sets_manager,
|
||||
mock_eval_set_results_manager,
|
||||
):
|
||||
"""Return a TestClient rooted in a temporary agents directory."""
|
||||
with (
|
||||
patch.object(signal, "signal", autospec=True, return_value=None),
|
||||
patch.object(
|
||||
fast_api_module,
|
||||
"create_session_service_from_options",
|
||||
autospec=True,
|
||||
return_value=mock_session_service,
|
||||
),
|
||||
patch.object(
|
||||
fast_api_module,
|
||||
"create_artifact_service_from_options",
|
||||
autospec=True,
|
||||
return_value=mock_artifact_service,
|
||||
),
|
||||
patch.object(
|
||||
fast_api_module,
|
||||
"create_memory_service_from_options",
|
||||
autospec=True,
|
||||
return_value=mock_memory_service,
|
||||
),
|
||||
patch.object(
|
||||
fast_api_module,
|
||||
"AgentLoader",
|
||||
autospec=True,
|
||||
return_value=mock_agent_loader,
|
||||
),
|
||||
patch.object(
|
||||
fast_api_module,
|
||||
"LocalEvalSetsManager",
|
||||
autospec=True,
|
||||
return_value=mock_eval_sets_manager,
|
||||
),
|
||||
patch.object(
|
||||
fast_api_module,
|
||||
"LocalEvalSetResultsManager",
|
||||
autospec=True,
|
||||
return_value=mock_eval_set_results_manager,
|
||||
),
|
||||
):
|
||||
app = get_fast_api_app(
|
||||
agents_dir=str(tmp_path),
|
||||
web=True,
|
||||
session_service_uri="",
|
||||
artifact_service_uri="",
|
||||
memory_service_uri="",
|
||||
allow_origins=["*"],
|
||||
a2a=False,
|
||||
host="127.0.0.1",
|
||||
port=8000,
|
||||
)
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def create_test_session(
|
||||
test_app, test_session_info, mock_session_service
|
||||
@@ -1175,5 +1253,103 @@ def test_patch_memory(test_app, create_test_session, mock_memory_service):
|
||||
logger.info("Add session to memory test completed successfully")
|
||||
|
||||
|
||||
def test_builder_final_save_preserves_tools_and_cleans_tmp(
|
||||
builder_test_client, tmp_path
|
||||
):
|
||||
files = [
|
||||
("files", ("app/__init__.py", b"from . import agent\n", "text/plain")),
|
||||
("files", ("app/tools.py", b"def tool():\n return 1\n", "text/plain")),
|
||||
(
|
||||
"files",
|
||||
("app/root_agent.yaml", b"name: app\n", "application/x-yaml"),
|
||||
),
|
||||
]
|
||||
response = builder_test_client.post("/builder/save?tmp=true", files=files)
|
||||
assert response.status_code == 200
|
||||
assert response.json() is True
|
||||
|
||||
response = builder_test_client.post(
|
||||
"/builder/save",
|
||||
files=[(
|
||||
"files",
|
||||
(
|
||||
"app/root_agent.yaml",
|
||||
b"name: app_updated\n",
|
||||
"application/x-yaml",
|
||||
),
|
||||
)],
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert response.json() is True
|
||||
|
||||
assert (tmp_path / "app" / "tools.py").is_file()
|
||||
assert not (tmp_path / "app" / "tmp" / "app").exists()
|
||||
tmp_dir = tmp_path / "app" / "tmp"
|
||||
assert not tmp_dir.exists() or not any(tmp_dir.iterdir())
|
||||
|
||||
|
||||
def test_builder_cancel_deletes_tmp_idempotent(builder_test_client, tmp_path):
|
||||
tmp_agent_root = tmp_path / "app" / "tmp" / "app"
|
||||
tmp_agent_root.mkdir(parents=True, exist_ok=True)
|
||||
(tmp_agent_root / "root_agent.yaml").write_text("name: app\n")
|
||||
|
||||
response = builder_test_client.post("/builder/app/app/cancel")
|
||||
assert response.status_code == 200
|
||||
assert response.json() is True
|
||||
assert not (tmp_path / "app" / "tmp").exists()
|
||||
|
||||
response = builder_test_client.post("/builder/app/app/cancel")
|
||||
assert response.status_code == 200
|
||||
assert response.json() is True
|
||||
assert not (tmp_path / "app" / "tmp").exists()
|
||||
|
||||
|
||||
def test_builder_get_tmp_true_recreates_tmp(builder_test_client, tmp_path):
|
||||
app_root = tmp_path / "app"
|
||||
app_root.mkdir(parents=True, exist_ok=True)
|
||||
(app_root / "root_agent.yaml").write_text("name: app\n")
|
||||
nested_dir = app_root / "nested"
|
||||
nested_dir.mkdir(parents=True, exist_ok=True)
|
||||
(nested_dir / "nested.yaml").write_text("nested: true\n")
|
||||
|
||||
assert not (app_root / "tmp").exists()
|
||||
response = builder_test_client.get("/builder/app/app?tmp=true")
|
||||
assert response.status_code == 200
|
||||
assert response.text == "name: app\n"
|
||||
|
||||
tmp_agent_root = app_root / "tmp" / "app"
|
||||
assert (tmp_agent_root / "root_agent.yaml").is_file()
|
||||
assert (tmp_agent_root / "nested" / "nested.yaml").is_file()
|
||||
|
||||
response = builder_test_client.get(
|
||||
"/builder/app/app?tmp=true&file_path=nested/nested.yaml"
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert response.text == "nested: true\n"
|
||||
|
||||
|
||||
def test_builder_get_tmp_true_missing_app_returns_empty(
|
||||
builder_test_client, tmp_path
|
||||
):
|
||||
response = builder_test_client.get("/builder/app/missing?tmp=true")
|
||||
assert response.status_code == 200
|
||||
assert response.text == ""
|
||||
assert not (tmp_path / "missing").exists()
|
||||
|
||||
|
||||
def test_builder_save_rejects_traversal(builder_test_client, tmp_path):
|
||||
response = builder_test_client.post(
|
||||
"/builder/save?tmp=true",
|
||||
files=[(
|
||||
"files",
|
||||
("app/../escape.yaml", b"nope\n", "application/x-yaml"),
|
||||
)],
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert response.json() is False
|
||||
assert not (tmp_path / "escape.yaml").exists()
|
||||
assert not (tmp_path / "app" / "tmp" / "escape.yaml").exists()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main(["-xvs", __file__])
|
||||
|
||||
Reference in New Issue
Block a user