fix: harden generated agent platform access

agent_builder/tools.py

@@ -1,6 +1,6 @@
 """LangChain tools the inner deepagents loop uses.
 
-  - workspace file CRUD (boto3 → MinIO, scoped to the user's bucket)
+  - workspace file CRUD through the caller's scoped workspace grant
   - sandbox python (for `a2a card` / `a2a validate` round-trips)
   - cp_deploy_tarball (call /v1/agents/from-tarball on the user's behalf)
 """
@@ -28,7 +28,7 @@ if TYPE_CHECKING:
     from .config import Settings
 
 
-A2A_PACK_MIN_VERSION = "0.1.25"
+A2A_PACK_MIN_VERSION = "0.1.26"
 
 
 @dataclass(frozen=True)
@@ -36,6 +36,8 @@ class ToolContext:
     bucket: str
     settings: "Settings"
     cp_jwt: str | None = None
+    workspace: Any | None = None
+    grant_token: str | None = None
 
 
 async def _wait_for_live_card(
@@ -95,6 +97,99 @@ def _ensure_bucket(s3: Any, bucket: str) -> None:
             raise
 
 
+class _ObjectStore:
+    def iter_keys(self, prefix: str) -> list[str]:
+        raise NotImplementedError
+
+    def get(self, key: str) -> bytes:
+        raise NotImplementedError
+
+    def put(self, key: str, body: bytes, content_type: str) -> None:
+        raise NotImplementedError
+
+    def delete(self, key: str) -> None:
+        raise NotImplementedError
+
+
+class _S3ObjectStore(_ObjectStore):
+    def __init__(self, s3: Any, bucket: str) -> None:
+        self._s3 = s3
+        self._bucket = bucket
+
+    def iter_keys(self, prefix: str) -> list[str]:
+        keys: list[str] = []
+        paginator = self._s3.get_paginator("list_objects_v2")
+        for page in paginator.paginate(Bucket=self._bucket, Prefix=prefix):
+            for obj in page.get("Contents") or ():
+                key = obj.get("Key")
+                if isinstance(key, str):
+                    keys.append(key)
+        return sorted(keys)
+
+    def get(self, key: str) -> bytes:
+        obj = self._s3.get_object(Bucket=self._bucket, Key=key)
+        return obj["Body"].read()
+
+    def put(self, key: str, body: bytes, content_type: str) -> None:
+        self._s3.put_object(
+            Bucket=self._bucket,
+            Key=key,
+            Body=body,
+            ContentType=content_type,
+        )
+
+    def delete(self, key: str) -> None:
+        self._s3.delete_object(Bucket=self._bucket, Key=key)
+
+
+class _WorkspaceObjectStore(_ObjectStore):
+    def __init__(self, workspace: Any) -> None:
+        self._workspace = workspace
+
+    def iter_keys(self, prefix: str) -> list[str]:
+        return sorted(
+            key
+            for key in self._workspace.iter_paths()
+            if isinstance(key, str) and key.startswith(prefix)
+        )
+
+    def get(self, key: str) -> bytes:
+        return self._workspace.read_bytes(key)
+
+    def put(self, key: str, body: bytes, content_type: str) -> None:
+        del content_type
+        self._workspace.write_bytes(key, body)
+
+    def delete(self, key: str) -> None:
+        self._workspace.delete_path(key)
+
+
+def _store_from_args(
+    store_or_s3: Any,
+    bucket_or_prefix: str,
+    prefix: str | None,
+) -> tuple[_ObjectStore, str]:
+    """Support new helper calls with ObjectStore and old tests with fake S3."""
+    if isinstance(store_or_s3, _ObjectStore):
+        return store_or_s3, bucket_or_prefix
+    if hasattr(store_or_s3, "iter_keys") and hasattr(store_or_s3, "get"):
+        return store_or_s3, bucket_or_prefix
+    if prefix is None:
+        raise TypeError("prefix is required when passing an S3 client")
+    return _S3ObjectStore(store_or_s3, bucket_or_prefix), prefix
+
+
+def _sandbox_headers(
+    bearer_token: str | None,
+    grant_token: str | None,
+) -> dict[str, str] | None:
+    if bearer_token:
+        return {"authorization": f"Bearer {bearer_token}"}
+    if grant_token:
+        return {"X-A2A-Grant": grant_token}
+    return None
+
+
 _AGENT_NAME_RE = re.compile(r"^[a-z][a-z0-9-]{1,62}$")
 _SKILL_NAME_RE = re.compile(r"^[a-z0-9](?:[a-z0-9-]{0,62}[a-z0-9])?$")
 _BUILDER_STATE_FILE = ".a2a-builder-state.json"
@@ -128,8 +223,12 @@ def _builder_state_key(name: str) -> str:
 def build_tools(ctx: ToolContext) -> list[Any]:
     bucket = ctx.bucket
     settings = ctx.settings
-    s3 = _s3(ctx)
-    _ensure_bucket(s3, bucket)
+    if ctx.workspace is not None:
+        store: _ObjectStore = _WorkspaceObjectStore(ctx.workspace)
+    else:
+        s3 = _s3(ctx)
+        _ensure_bucket(s3, bucket)
+        store = _S3ObjectStore(s3, bucket)
 
     @tool
     def init_agent_template(
@@ -190,16 +289,11 @@ def build_tools(ctx: ToolContext) -> list[Any]:
         except ValueError as exc:
             return json.dumps({"error": str(exc)})
         out: list[dict[str, Any]] = []
-        paginator = s3.get_paginator("list_objects_v2")
-        for page in paginator.paginate(Bucket=bucket, Prefix=prefix):
-            for obj in page.get("Contents") or ():
-                rel = obj["Key"][len(prefix):]
-                if rel == _BUILDER_STATE_FILE or rel.startswith(_BUILDER_INTERNAL_PREFIX):
-                    continue
-                out.append({
-                    "path": rel,
-                    "size": int(obj.get("Size") or 0),
-                })
+        for key in store.iter_keys(prefix):
+            rel = key[len(prefix):]
+            if rel == _BUILDER_STATE_FILE or rel.startswith(_BUILDER_INTERNAL_PREFIX):
+                continue
+            out.append({"path": rel, "size": len(store.get(key))})
         return json.dumps({"agent": name, "files": out})
 
     @tool
@@ -217,10 +311,7 @@ def build_tools(ctx: ToolContext) -> list[Any]:
         if rel == _BUILDER_STATE_FILE or rel.startswith(_BUILDER_INTERNAL_PREFIX):
             return json.dumps({"error": f"{rel} is managed by agent-builder"})
         key = prefix + rel
-        s3.put_object(
-            Bucket=bucket, Key=key, Body=content.encode("utf-8"),
-            ContentType="text/plain; charset=utf-8",
-        )
+        store.put(key, content.encode("utf-8"), "text/plain; charset=utf-8")
         return json.dumps({"ok": True, "path": key, "size": len(content)})
 
     @tool
@@ -235,10 +326,10 @@ def build_tools(ctx: ToolContext) -> list[Any]:
             return json.dumps({"error": f"{rel} is managed by agent-builder"})
         key = prefix + rel
         try:
-            obj = s3.get_object(Bucket=bucket, Key=key)
-        except ClientError as exc:
+            data = store.get(key)
+        except Exception as exc:  # noqa: BLE001
             return json.dumps({"error": f"read failed: {exc}"})
-        body = obj["Body"].read().decode("utf-8", errors="replace")
+        body = data.decode("utf-8", errors="replace")
         return json.dumps({"path": key, "content": body[:200_000]})
 
     @tool
@@ -275,12 +366,7 @@ def build_tools(ctx: ToolContext) -> list[Any]:
         for rel, content in payloads.items():
             key = base + rel
             encoded = content.encode("utf-8")
-            s3.put_object(
-                Bucket=bucket,
-                Key=key,
-                Body=encoded,
-                ContentType="text/plain; charset=utf-8",
-            )
+            store.put(key, encoded, "text/plain; charset=utf-8")
             written.append({"path": key, "size": len(encoded)})
         return json.dumps({
             "ok": True,
@@ -308,7 +394,7 @@ def build_tools(ctx: ToolContext) -> list[Any]:
 
         # Bundle the workspace files into the sandbox via base64 tar so
         # the sandbox doesn't need MinIO read access for this skill.
-        bundle_bytes = _tarball_workspace_dir(s3, bucket, prefix)
+        bundle_bytes = _tarball_workspace_dir(store, prefix)
         if not bundle_bytes:
             return json.dumps({
                 "error": "no files in agents/" + name + "/ — write some first",
@@ -331,10 +417,7 @@ def build_tools(ctx: ToolContext) -> list[Any]:
             "  a2a frontend info --project . || true\n"
             "fi\n"
         )
-        headers = (
-            {"authorization": f"Bearer {settings.sandbox_token}"}
-            if settings.sandbox_token else None
-        )
+        headers = _sandbox_headers(settings.sandbox_token, ctx.grant_token)
         async with httpx.AsyncClient(timeout=settings.sandbox_timeout_s + 30) as c:
             r = await c.post(
                 f"{settings.sandbox_url}/v1/run_shell",
@@ -399,7 +482,7 @@ def build_tools(ctx: ToolContext) -> list[Any]:
             latest_deploy = await _latest_cp_deployment(settings.cp_url, ctx.cp_jwt, name)
         except RuntimeError as exc:
             return json.dumps({"error": str(exc)})
-        builder_state = _read_builder_state(s3, bucket, name)
+        builder_state = _read_builder_state(store, name)
         drift = _deployment_drift_error(
             name,
             latest_deploy,
@@ -412,7 +495,7 @@ def build_tools(ctx: ToolContext) -> list[Any]:
         if force:
             base_head = None
 
-        bundle_bytes = _tarball_workspace_dir(s3, bucket, prefix)
+        bundle_bytes = _tarball_workspace_dir(store, prefix)
         if not bundle_bytes:
             return json.dumps({"error": f"no files at agents/{name}/"})
         # Read the entrypoint from a2a.yaml so we can pass it verbatim.
@@ -446,8 +529,7 @@ def build_tools(ctx: ToolContext) -> list[Any]:
         head_sha = body.get("head_sha")
         if isinstance(head_sha, str) and head_sha:
             _write_builder_state(
-                s3,
-                bucket,
+                store,
                 name,
                 {
                     "agent": name,
@@ -579,10 +661,9 @@ def build_tools(ctx: ToolContext) -> list[Any]:
             return json.dumps({"error": f"invalid source tarball: {exc}"})
         if not files:
             return json.dumps({"error": "managed repo source export was empty"})
-        written = _replace_workspace_files(s3, bucket, prefix, files)
+        written = _replace_workspace_files(store, prefix, files)
         _write_builder_state(
-            s3,
-            bucket,
+            store,
             name,
             {
                 "agent": name,
@@ -691,58 +772,46 @@ def build_tools(ctx: ToolContext) -> list[Any]:
 # --- helpers (not tools) ---
 
 
-def _tarball_workspace_dir(s3: Any, bucket: str, prefix: str) -> bytes:
-    """Pull every object under ``prefix`` from MinIO and produce a
-    gzipped tarball whose entries are relative to ``prefix``."""
-    paginator = s3.get_paginator("list_objects_v2")
+def _tarball_workspace_dir(
+    store_or_s3: Any,
+    bucket_or_prefix: str,
+    prefix: str | None = None,
+) -> bytes:
+    """Pull every object under ``prefix`` and produce a relative tarball."""
+    store, prefix = _store_from_args(store_or_s3, bucket_or_prefix, prefix)
     buf = io.BytesIO()
     added = False
     with tarfile.open(fileobj=buf, mode="w:gz") as tf:
-        for page in paginator.paginate(Bucket=bucket, Prefix=prefix):
-            for obj in page.get("Contents") or ():
-                key = obj["Key"]
-                rel = key[len(prefix):]
-                if not rel:
-                    continue
-                if rel == _BUILDER_STATE_FILE or rel.startswith(_BUILDER_INTERNAL_PREFIX):
-                    continue
-                body = s3.get_object(Bucket=bucket, Key=key)["Body"].read()
-                info = tarfile.TarInfo(name=rel)
-                info.size = len(body)
-                info.mode = 0o644
-                tf.addfile(info, io.BytesIO(body))
-                added = True
+        for key in store.iter_keys(prefix):
+            rel = key[len(prefix):]
+            if not rel:
+                continue
+            if rel == _BUILDER_STATE_FILE or rel.startswith(_BUILDER_INTERNAL_PREFIX):
+                continue
+            body = store.get(key)
+            info = tarfile.TarInfo(name=rel)
+            info.size = len(body)
+            info.mode = 0o644
+            tf.addfile(info, io.BytesIO(body))
+            added = True
     return buf.getvalue() if added else b""
 
 
-def _read_builder_state(s3: Any, bucket: str, name: str) -> dict[str, Any]:
+def _read_builder_state(store: _ObjectStore, name: str) -> dict[str, Any]:
     try:
-        obj = s3.get_object(Bucket=bucket, Key=_builder_state_key(name))
-    except ClientError as exc:
-        code = exc.response.get("Error", {}).get("Code")
-        if code in {"404", "NoSuchKey", "NotFound"}:
-            return {}
-        return {}
-    try:
-        data = json.loads(obj["Body"].read().decode("utf-8"))
-    except (TypeError, ValueError, UnicodeDecodeError):
+        data = json.loads(store.get(_builder_state_key(name)).decode("utf-8"))
+    except Exception:  # noqa: BLE001
         return {}
     return data if isinstance(data, dict) else {}
 
 
 def _write_builder_state(
-    s3: Any,
-    bucket: str,
+    store: _ObjectStore,
     name: str,
     state: dict[str, Any],
 ) -> None:
     body = json.dumps(state, sort_keys=True, indent=2).encode("utf-8")
-    s3.put_object(
-        Bucket=bucket,
-        Key=_builder_state_key(name),
-        Body=body,
-        ContentType="application/json",
-    )
+    store.put(_builder_state_key(name), body, "application/json")
 
 
 def _render_skill_md(skill_name: str, description: str, instructions: str) -> str:
@@ -798,16 +867,9 @@ def _safe_supporting_skill_path(path: str) -> str:
     return rel
 
 
-def _delete_workspace_objects(s3: Any, bucket: str, prefix: str) -> None:
-    # MinIO rejects boto3's batch DeleteObjects request in this environment
-    # when Content-MD5 is not present. Per-object deletes avoid that fragile
-    # API path and are fine for agent source workspaces.
-    paginator = s3.get_paginator("list_objects_v2")
-    for page in paginator.paginate(Bucket=bucket, Prefix=prefix):
-        for obj in page.get("Contents") or ():
-            key = obj.get("Key")
-            if key:
-                s3.delete_object(Bucket=bucket, Key=key)
+def _delete_workspace_objects(store: _ObjectStore, prefix: str) -> None:
+    for key in store.iter_keys(prefix):
+        store.delete(key)
 
 
 def _files_from_tarball(bundle: bytes) -> dict[str, bytes]:
@@ -834,21 +896,25 @@ def _files_from_tarball(bundle: bytes) -> dict[str, bytes]:
 
 
 def _replace_workspace_files(
-    s3: Any,
-    bucket: str,
-    prefix: str,
-    files: dict[str, bytes],
+    store_or_s3: Any,
+    bucket_or_prefix: str,
+    prefix_or_files: str | dict[str, bytes],
+    files: dict[str, bytes] | None = None,
 ) -> list[dict[str, Any]]:
-    _delete_workspace_objects(s3, bucket, prefix)
+    if files is None:
+        if not isinstance(prefix_or_files, dict):
+            raise TypeError("files are required")
+        store, prefix = _store_from_args(store_or_s3, bucket_or_prefix, None)
+        files = prefix_or_files
+    else:
+        if not isinstance(prefix_or_files, str):
+            raise TypeError("prefix must be a string")
+        store, prefix = _store_from_args(store_or_s3, bucket_or_prefix, prefix_or_files)
     written: list[dict[str, Any]] = []
+    _delete_workspace_objects(store, prefix)
     for rel, body in sorted(files.items()):
         key = prefix + rel
-        s3.put_object(
-            Bucket=bucket,
-            Key=key,
-            Body=body,
-            ContentType="application/octet-stream",
-        )
+        store.put(key, body, "application/octet-stream")
         written.append({"path": rel, "size": len(body)})
     return written