| """Sync session workspace files to/from AgentCore invoke payloads.""" |
|
|
| from __future__ import annotations |
|
|
| import base64 |
| import os |
| from pathlib import Path |
|
|
| _DEFAULT_MAX_BYTES = 8 * 1024 * 1024 |
|
|
|
|
| def max_workspace_sync_bytes() -> int: |
| raw = (os.environ.get("AGENTCORE_MAX_UPLOAD_BYTES") or "").strip() |
| if raw.isdigit(): |
| return int(raw) |
| return _DEFAULT_MAX_BYTES |
|
|
|
|
| def _ensure_session_workspace(session_hash: str | None) -> Path: |
| try: |
| from bundle_support.session_workspace import ensure_session_workspace |
| except ImportError: |
| from pi.session_workspace import ( |
| ensure_session_workspace, |
| ) |
|
|
| return ensure_session_workspace(session_hash or "") |
|
|
|
|
| def apply_workspace_files(session_hash: str | None, files: list[dict]) -> list[str]: |
| """Write base64-encoded files into the session workspace. Returns relative paths written.""" |
| if not files: |
| return [] |
| root = _ensure_session_workspace(session_hash).resolve() |
| written: list[str] = [] |
| for item in files: |
| if not isinstance(item, dict): |
| continue |
| relative = str(item.get("relative_path") or item.get("name") or "").strip() |
| encoded = str(item.get("content_base64") or "").strip() |
| if not relative or not encoded: |
| continue |
| dest = (root / relative).resolve() |
| try: |
| dest.relative_to(root) |
| except ValueError: |
| continue |
| payload = base64.b64decode(encoded, validate=True) |
| if len(payload) > max_workspace_sync_bytes(): |
| continue |
| dest.parent.mkdir(parents=True, exist_ok=True) |
| dest.write_bytes(payload) |
| written.append(str(dest.relative_to(root)).replace("\\", "/")) |
| return written |
|
|
|
|
| def collect_workspace_files_for_sync( |
| session_hash: str | None, |
| *, |
| prefix: str = "redact/", |
| ) -> list[dict[str, str]]: |
| """Collect workspace files under *prefix* for download to the Gradio client.""" |
| root = _ensure_session_workspace(session_hash).resolve() |
| if not root.is_dir(): |
| return [] |
| limit = max_workspace_sync_bytes() |
| out: list[dict[str, str]] = [] |
| for path in sorted(root.rglob("*")): |
| if not path.is_file(): |
| continue |
| rel = str(path.relative_to(root)).replace("\\", "/") |
| if prefix and not rel.startswith(prefix.lstrip("/")): |
| continue |
| size = path.stat().st_size |
| if size > limit: |
| continue |
| out.append( |
| { |
| "relative_path": rel, |
| "content_base64": base64.b64encode(path.read_bytes()).decode("ascii"), |
| } |
| ) |
| return out |
|
|
