diff --git a/src/google/adk/cli/fast_api.py b/src/google/adk/cli/fast_api.py index 131213ec..b5abaa50 100644 --- a/src/google/adk/cli/fast_api.py +++ b/src/google/adk/cli/fast_api.py @@ -211,75 +211,214 @@ def get_fast_api_app( **extra_fast_api_args, ) + agents_base_path = (Path.cwd() / agents_dir).resolve() + + def _get_app_root(app_name: str) -> Path: + if app_name in ("", ".", ".."): + raise ValueError(f"Invalid app name: {app_name!r}") + if Path(app_name).name != app_name or "\\" in app_name: + raise ValueError(f"Invalid app name: {app_name!r}") + app_root = (agents_base_path / app_name).resolve() + if not app_root.is_relative_to(agents_base_path): + raise ValueError(f"Invalid app name: {app_name!r}") + return app_root + + def _normalize_relative_path(path: str) -> str: + return path.replace("\\", "/").lstrip("/") + + def _has_parent_reference(path: str) -> bool: + return any(part == ".." for part in path.split("/")) + + def _parse_upload_filename(filename: Optional[str]) -> tuple[str, str]: + if not filename: + raise ValueError("Upload filename is missing.") + filename = _normalize_relative_path(filename) + if "/" not in filename: + raise ValueError(f"Invalid upload filename: {filename!r}") + app_name, rel_path = filename.split("/", 1) + if not app_name or not rel_path: + raise ValueError(f"Invalid upload filename: {filename!r}") + if rel_path.startswith("/"): + raise ValueError(f"Absolute upload path rejected: {filename!r}") + if _has_parent_reference(rel_path): + raise ValueError(f"Path traversal rejected: {filename!r}") + return app_name, rel_path + + def _parse_file_path(file_path: str) -> str: + file_path = _normalize_relative_path(file_path) + if not file_path: + raise ValueError("file_path is missing.") + if file_path.startswith("/"): + raise ValueError(f"Absolute file_path rejected: {file_path!r}") + if _has_parent_reference(file_path): + raise ValueError(f"Path traversal rejected: {file_path!r}") + return file_path + + def _resolve_under_dir(root_dir: Path, rel_path: str) -> Path: + file_path = root_dir / rel_path + resolved_root_dir = root_dir.resolve() + resolved_file_path = file_path.resolve() + if not resolved_file_path.is_relative_to(resolved_root_dir): + raise ValueError(f"Path escapes root_dir: {rel_path!r}") + return file_path + + def _get_tmp_agent_root(app_root: Path, app_name: str) -> Path: + tmp_agent_root = app_root / "tmp" / app_name + resolved_tmp_agent_root = tmp_agent_root.resolve() + if not resolved_tmp_agent_root.is_relative_to(app_root): + raise ValueError(f"Invalid tmp path for app: {app_name!r}") + return tmp_agent_root + + def copy_dir_contents(source_dir: Path, dest_dir: Path) -> None: + dest_dir.mkdir(parents=True, exist_ok=True) + for source_path in source_dir.iterdir(): + if source_path.name == "tmp": + continue + + dest_path = dest_dir / source_path.name + if source_path.is_dir(): + if dest_path.exists() and dest_path.is_file(): + dest_path.unlink() + shutil.copytree(source_path, dest_path, dirs_exist_ok=True) + elif source_path.is_file(): + if dest_path.exists() and dest_path.is_dir(): + shutil.rmtree(dest_path) + shutil.copy2(source_path, dest_path) + + def cleanup_tmp(app_name: str) -> bool: + try: + app_root = _get_app_root(app_name) + except ValueError as exc: + logger.exception("Error in cleanup_tmp: %s", exc) + return False + + try: + tmp_agent_root = _get_tmp_agent_root(app_root, app_name) + except ValueError as exc: + logger.exception("Error in cleanup_tmp: %s", exc) + return False + + try: + shutil.rmtree(tmp_agent_root) + except FileNotFoundError: + pass + except OSError as exc: + logger.exception("Error deleting tmp agent root: %s", exc) + return False + + tmp_dir = app_root / "tmp" + resolved_tmp_dir = tmp_dir.resolve() + if not resolved_tmp_dir.is_relative_to(app_root): + logger.error( + "Refusing to delete tmp outside app_root: %s", resolved_tmp_dir + ) + return False + + try: + tmp_dir.rmdir() + except OSError: + pass + + return True + + def ensure_tmp_exists(app_name: str) -> bool: + try: + app_root = _get_app_root(app_name) + except ValueError as exc: + logger.exception("Error in ensure_tmp_exists: %s", exc) + return False + + if not app_root.is_dir(): + return False + + try: + tmp_agent_root = _get_tmp_agent_root(app_root, app_name) + except ValueError as exc: + logger.exception("Error in ensure_tmp_exists: %s", exc) + return False + + if tmp_agent_root.exists(): + return True + + try: + tmp_agent_root.mkdir(parents=True, exist_ok=True) + copy_dir_contents(app_root, tmp_agent_root) + except OSError as exc: + logger.exception("Error in ensure_tmp_exists: %s", exc) + return False + + return True + @app.post("/builder/save", response_model_exclude_none=True) async def builder_build( files: list[UploadFile], tmp: Optional[bool] = False ) -> bool: - base_path = Path.cwd() / agents_dir - for file in files: - if not file.filename: - logger.exception("Agent name is missing in the input files") - return False - agent_name, filename = file.filename.split("/") - agent_dir = os.path.join(base_path, agent_name) - try: - # File name format: {app_name}/{agent_name}.yaml - if tmp: - agent_dir = os.path.join(agent_dir, "tmp/" + agent_name) - os.makedirs(agent_dir, exist_ok=True) - file_path = os.path.join(agent_dir, filename) - with open(file_path, "wb") as buffer: + try: + if tmp: + app_names = set() + uploads = [] + for file in files: + app_name, rel_path = _parse_upload_filename(file.filename) + app_names.add(app_name) + uploads.append((rel_path, file)) + + if len(app_names) != 1: + logger.error( + "Exactly one app name is required, found: %s", sorted(app_names) + ) + return False + + app_name = next(iter(app_names)) + app_root = _get_app_root(app_name) + tmp_agent_root = _get_tmp_agent_root(app_root, app_name) + tmp_agent_root.mkdir(parents=True, exist_ok=True) + + for rel_path, file in uploads: + destination_path = _resolve_under_dir(tmp_agent_root, rel_path) + destination_path.parent.mkdir(parents=True, exist_ok=True) + with destination_path.open("wb") as buffer: shutil.copyfileobj(file.file, buffer) - else: - source_dir = os.path.join(agent_dir, "tmp/" + agent_name) - destination_dir = agent_dir - for item in os.listdir(source_dir): - source_item = os.path.join(source_dir, item) - destination_item = os.path.join(destination_dir, item) - if os.path.isdir(source_item): - shutil.copytree(source_item, destination_item, dirs_exist_ok=True) - # Check if the item is a file - elif os.path.isfile(source_item): - shutil.copy2(source_item, destination_item) - except Exception as e: - logger.exception("Error in builder_build: %s", e) + return True + + app_names = set() + uploads = [] + for file in files: + app_name, rel_path = _parse_upload_filename(file.filename) + app_names.add(app_name) + uploads.append((rel_path, file)) + + if len(app_names) != 1: + logger.error( + "Exactly one app name is required, found: %s", sorted(app_names) + ) return False - return True + app_name = next(iter(app_names)) + app_root = _get_app_root(app_name) + app_root.mkdir(parents=True, exist_ok=True) + + tmp_agent_root = _get_tmp_agent_root(app_root, app_name) + if tmp_agent_root.is_dir(): + copy_dir_contents(tmp_agent_root, app_root) + + for rel_path, file in uploads: + destination_path = _resolve_under_dir(app_root, rel_path) + destination_path.parent.mkdir(parents=True, exist_ok=True) + with destination_path.open("wb") as buffer: + shutil.copyfileobj(file.file, buffer) + + return cleanup_tmp(app_name) + except ValueError as exc: + logger.exception("Error in builder_build: %s", exc) + return False + except OSError as exc: + logger.exception("Error in builder_build: %s", exc) + return False @app.post("/builder/app/{app_name}/cancel", response_model_exclude_none=True) async def builder_cancel(app_name: str) -> bool: - base_path = Path.cwd() / agents_dir - agent_dir = os.path.join(base_path, app_name) - destination_dir = os.path.join(agent_dir, "tmp/" + app_name) - source_dir = agent_dir - source_items = set(os.listdir(source_dir)) - try: - for item in os.listdir(destination_dir): - if item in source_items: - continue - # If it doesn't exist in the source, delete it from the destination - item_path = os.path.join(destination_dir, item) - if os.path.isdir(item_path): - shutil.rmtree(item_path) - elif os.path.isfile(item_path): - os.remove(item_path) - - for item in os.listdir(source_dir): - source_item = os.path.join(source_dir, item) - destination_item = os.path.join(destination_dir, item) - if item == "tmp" and os.path.isdir(source_item): - continue - if os.path.isdir(source_item): - shutil.copytree(source_item, destination_item, dirs_exist_ok=True) - # Check if the item is a file - elif os.path.isfile(source_item): - shutil.copy2(source_item, destination_item) - except Exception as e: - logger.exception("Error in builder_build: %s", e) - return False - return True + return cleanup_tmp(app_name) @app.get( "/builder/app/{app_name}", @@ -291,34 +430,42 @@ def get_fast_api_app( file_path: Optional[str] = None, tmp: Optional[bool] = False, ): - base_path = Path.cwd() / agents_dir - agent_dir = base_path / app_name + try: + app_root = _get_app_root(app_name) + except ValueError as exc: + logger.exception("Error in get_agent_builder: %s", exc) + return "" + + agent_dir = app_root if tmp: - agent_dir = agent_dir / "tmp" - agent_dir = agent_dir / app_name + if not ensure_tmp_exists(app_name): + return "" + agent_dir = app_root / "tmp" / app_name + if not file_path: - file_name = "root_agent.yaml" - root_file_path = agent_dir / file_name - if not root_file_path.is_file(): - return "" - else: - return FileResponse( - path=root_file_path, - media_type="application/x-yaml", - filename="${app_name}.yaml", - headers={"Cache-Control": "no-store"}, - ) + rel_path = "root_agent.yaml" else: - agent_file_path = agent_dir / file_path - if not agent_file_path.is_file(): + try: + rel_path = _parse_file_path(file_path) + except ValueError as exc: + logger.exception("Error in get_agent_builder: %s", exc) return "" - else: - return FileResponse( - path=agent_file_path, - media_type="application/x-yaml", - filename=file_path, - headers={"Cache-Control": "no-store"}, - ) + + try: + agent_file_path = _resolve_under_dir(agent_dir, rel_path) + except ValueError as exc: + logger.exception("Error in get_agent_builder: %s", exc) + return "" + + if not agent_file_path.is_file(): + return "" + + return FileResponse( + path=agent_file_path, + media_type="application/x-yaml", + filename=file_path or f"{app_name}.yaml", + headers={"Cache-Control": "no-store"}, + ) if a2a: from a2a.server.apps import A2AStarletteApplication diff --git a/tests/unittests/cli/test_fast_api.py b/tests/unittests/cli/test_fast_api.py index 263e4704..7a5627ff 100755 --- a/tests/unittests/cli/test_fast_api.py +++ b/tests/unittests/cli/test_fast_api.py @@ -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__])