encounter/adk-python
0
0
Fork 0
mirror of https://github.com/encounter/adk-python.git synced 2026-07-09 18:19:28 -07:00
Code Issues Packages Projects Releases Wiki Activity

fix: Harden YAML builder tmp save/cleanup

Browse Source 
- 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
2026-01-05 10:47:12 -08:00
committed by Copybara-Service
parent 3ec7ae3b8d
commit 6f259f08b3
2 changed files with 418 additions and 95 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/google/adk/cli/fast_api.py
+222 -75
View File
@@ -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:
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,32 +430,40 @@ 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
if tmp:
agent_dir = agent_dir / "tmp"
agent_dir = agent_dir / 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():
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:
if not ensure_tmp_exists(app_name):
return ""
agent_dir = app_root / "tmp" / app_name
if not file_path:
rel_path = "root_agent.yaml"
else:
return FileResponse(
path=root_file_path,
media_type="application/x-yaml",
filename="${app_name}.yaml",
headers={"Cache-Control": "no-store"},
)
else:
agent_file_path = agent_dir / file_path
try:
rel_path = _parse_file_path(file_path)
except ValueError as exc:
logger.exception("Error in get_agent_builder: %s", exc)
return ""
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 ""
else:
return FileResponse(
path=agent_file_path,
media_type="application/x-yaml",
filename=file_path,
filename=file_path or f"{app_name}.yaml",
headers={"Cache-Control": "no-store"},
)
tests/unittests/cli/test_fast_api.py
+189 -13
View File
@@ -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__])