This repository has been archived on 2026-06-28. You can view files and clone it. You cannot open issues or pull requests or push a commit.
Files
agent-builder/agent_builder/tools.py
robert a3466c7a9a
All checks were successful
build / build (push) Successful in 27s
Track builder workspace deployment base
2026-05-19 09:10:43 -03:00

723 lines
27 KiB
Python

"""LangChain tools the inner deepagents loop uses.
- workspace file CRUD (boto3 → MinIO, scoped to the user's bucket)
- sandbox python (for `a2a card` / `a2a validate` round-trips)
- cp_deploy_tarball (call /v1/agents/from-tarball on the user's behalf)
"""
from __future__ import annotations
import asyncio
import base64
import io
import json
import re
import tarfile
import time
from dataclasses import dataclass
from importlib import resources
from typing import TYPE_CHECKING, Any
import boto3
import httpx
from botocore.config import Config as _BotoConfig
from botocore.exceptions import ClientError
from langchain_core.tools import tool
if TYPE_CHECKING:
from .config import Settings
@dataclass(frozen=True)
class ToolContext:
bucket: str
settings: "Settings"
cp_jwt: str | None = None
async def _wait_for_live_card(
url: str | None,
expected_version: str | None,
timeout_s: float = 180.0,
) -> tuple[bool, dict[str, Any]]:
"""Poll ``{url}/.well-known/agent-card`` until ``version`` matches
``expected_version`` or ``timeout_s`` elapses. Returns
``(matched, last_seen_card)``. ``last_seen_card`` is ``{}`` if the
endpoint never responded with a parseable card.
Backoff: 2s → 1.5x per try, capped at 10s. Tolerates pod-not-ready
(any HTTPError or 4xx/5xx) — just keeps polling.
"""
if not url or not expected_version:
return False, {}
deadline = time.monotonic() + timeout_s
last_card: dict[str, Any] = {}
interval = 2.0
while time.monotonic() < deadline:
try:
async with httpx.AsyncClient(timeout=5.0) as c:
r = await c.get(f"{url.rstrip('/')}/.well-known/agent-card")
if r.status_code == 200:
card = r.json() or {}
if isinstance(card, dict):
last_card = card
if card.get("version") == expected_version:
return True, card
except (httpx.HTTPError, ValueError):
pass
await asyncio.sleep(interval)
interval = min(interval * 1.5, 10.0)
return False, last_card
def _s3(ctx: ToolContext) -> Any:
return boto3.client(
"s3",
endpoint_url=ctx.settings.minio_endpoint,
aws_access_key_id=ctx.settings.minio_access_key,
aws_secret_access_key=ctx.settings.minio_secret_key,
region_name="us-east-1",
config=_BotoConfig(s3={"addressing_style": "path"}),
)
def _ensure_bucket(s3: Any, bucket: str) -> None:
try:
s3.head_bucket(Bucket=bucket)
except ClientError as exc:
code = exc.response.get("Error", {}).get("Code")
if code in {"404", "NoSuchBucket", "NotFound"}:
s3.create_bucket(Bucket=bucket)
elif code != "403":
raise
_AGENT_NAME_RE = re.compile(r"^[a-z][a-z0-9-]{1,62}$")
_BUILDER_STATE_FILE = ".a2a-builder-state.json"
def _validate_name(name: str) -> None:
if not _AGENT_NAME_RE.match(name):
raise ValueError(
f"invalid agent name {name!r}: must match {_AGENT_NAME_RE.pattern}"
)
def _agent_prefix(name: str) -> str:
_validate_name(name)
return f"agents/{name}/"
def _builder_state_key(name: str) -> str:
return _agent_prefix(name) + _BUILDER_STATE_FILE
def build_tools(ctx: ToolContext) -> list[Any]:
bucket = ctx.bucket
settings = ctx.settings
s3 = _s3(ctx)
_ensure_bucket(s3, bucket)
@tool
def init_agent_template(
name: str,
description: str = "A new A2A agent",
) -> str:
"""Initialize ``agents/<name>/`` from the installed a2a-pack template.
Use this FIRST for a new project. It writes the same baseline files
as ``a2a init`` for the SDK version installed in this runtime:
``agent.py``, ``a2a.yaml``, and ``requirements.txt``. After this,
edit those files with ``read_agent_file`` + ``write_agent_file`` to
implement the user's requested behavior.
"""
try:
files = _render_a2a_init_template(name, description=description)
prefix = _agent_prefix(name)
except (OSError, RuntimeError, ValueError) as exc:
return json.dumps({"error": str(exc)})
written: list[dict[str, Any]] = []
for path, content in files.items():
key = prefix + path
s3.put_object(
Bucket=bucket,
Key=key,
Body=content.encode("utf-8"),
ContentType="text/plain; charset=utf-8",
)
written.append({"path": key, "size": len(content)})
return json.dumps({"ok": True, "agent": name, "files": written})
@tool
def list_agent_files(name: str) -> str:
"""List every file under ``agents/<name>/`` in the user's workspace.
Use this before edits to see what you already wrote / what's left.
"""
try:
prefix = _agent_prefix(name)
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:
continue
out.append({
"path": rel,
"size": int(obj.get("Size") or 0),
})
return json.dumps({"agent": name, "files": out})
@tool
def write_agent_file(name: str, path: str, content: str) -> str:
"""Write a file under ``agents/<name>/<path>`` in the user's workspace.
Use this for ``agent.py``, ``a2a.yaml``, ``requirements.txt``, and
any auxiliary modules the agent needs.
"""
try:
prefix = _agent_prefix(name)
except ValueError as exc:
return json.dumps({"error": str(exc)})
rel = path.lstrip("/")
if rel == _BUILDER_STATE_FILE:
return json.dumps({"error": f"{_BUILDER_STATE_FILE} is managed by deploys"})
key = prefix + rel
s3.put_object(
Bucket=bucket, Key=key, Body=content.encode("utf-8"),
ContentType="text/plain; charset=utf-8",
)
return json.dumps({"ok": True, "path": key, "size": len(content)})
@tool
def read_agent_file(name: str, path: str) -> str:
"""Read a single file from the scaffolded agent project."""
try:
prefix = _agent_prefix(name)
except ValueError as exc:
return json.dumps({"error": str(exc)})
rel = path.lstrip("/")
if rel == _BUILDER_STATE_FILE:
return json.dumps({"error": f"{_BUILDER_STATE_FILE} is managed by deploys"})
key = prefix + rel
try:
obj = s3.get_object(Bucket=bucket, Key=key)
except ClientError as exc:
return json.dumps({"error": f"read failed: {exc}"})
body = obj["Body"].read().decode("utf-8", errors="replace")
return json.dumps({"path": key, "content": body[:200_000]})
@tool
async def test_agent_in_sandbox(name: str) -> str:
"""Spin a microVM, ``pip install a2a-pack`` + the agent's own
requirements.txt, then run ``a2a card`` to verify the scaffold
produces a valid Card. Returns stdout/stderr.
"""
try:
prefix = _agent_prefix(name)
except ValueError as exc:
return json.dumps({"error": str(exc)})
# 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)
if not bundle_bytes:
return json.dumps({
"error": "no files in agents/" + name + "/ — write some first",
})
b64 = base64.b64encode(bundle_bytes).decode("ascii")
script = (
"set -e\n"
"pip install --quiet a2a-pack >/dev/null\n"
"mkdir -p /tmp/agent\n"
f"echo '{b64}' | base64 -d | tar -xzf - -C /tmp/agent\n"
"cd /tmp/agent\n"
"ls -la\n"
"if [ -f requirements.txt ]; then\n"
" pip install --quiet -r requirements.txt >/dev/null || true\n"
"fi\n"
"echo '--- card ---'\n"
"a2a card --project .\n"
)
async with httpx.AsyncClient(timeout=settings.sandbox_timeout_s + 30) as c:
r = await c.post(
f"{settings.sandbox_url}/v1/run_shell",
json={
"bucket": bucket, "script": script,
"image": settings.image, "memory_mib": 512,
"timeout_seconds": settings.sandbox_timeout_s,
},
)
if r.status_code >= 400:
return json.dumps({"error": f"sandbox {r.status_code}", "detail": r.text[:1000]})
out = r.json()
return json.dumps({
"exit_code": out.get("exit_code"),
"stdout": (out.get("stdout") or "")[:4000],
"stderr": (out.get("stderr") or "")[:4000],
})
@tool
async def cp_deploy_tarball(
name: str,
version: str = "0.1.0",
public: bool = True,
force: bool = False,
) -> str:
"""Tar the agent at ``agents/<name>/`` and POST it to the control
plane's ``/v1/agents/from-tarball`` endpoint on the user's behalf,
then BLOCK until the live ``.well-known/agent-card`` reports the
new version (or 180s timeout).
Returns JSON with ``ok``, ``live``, ``url``, ``version``,
``live_version``, and ``live_skills[{name, input_schema}]``. When
``live=true`` the agent is callable; when ``live=false`` the
deploy POST succeeded but the pod never came up with the new
version — caller should treat this as a failure, not success.
Compare ``live_skills[].input_schema`` against the code you wrote
to confirm the advertised schema matches the params the function
actually accepts. Mismatches cause downstream ``agent 400`` from
``call_agent``.
``force`` defaults to false so a stale MinIO workspace cannot
silently overwrite source that was deployed from a user's computer
or another tool. Only set ``force=True`` after the user explicitly
accepts replacing the current managed repo contents.
Requires the orchestrator to have forwarded the user's CP JWT
(the platform does this automatically when this agent's Card
declared ``wants_cp_jwt=True``).
"""
try:
prefix = _agent_prefix(name)
except ValueError as exc:
return json.dumps({"error": str(exc)})
if not ctx.cp_jwt:
return json.dumps({
"error": "no CP JWT forwarded — agent declaration missing wants_cp_jwt=True",
})
try:
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)
drift = _deployment_drift_error(
name,
latest_deploy,
builder_state,
force=force,
)
if drift is not None:
return json.dumps(drift)
base_head = builder_state.get("repo_head_sha")
if force:
base_head = None
bundle_bytes = _tarball_workspace_dir(s3, bucket, 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.
entrypoint = _read_entrypoint(bundle_bytes) or f"agent:{_class_name(name)}"
# NOTE: the CP from-tarball endpoint names the file field
# ``source`` (see apps/control-plane/control_plane/routes/agents.py).
files = {"source": (f"{name}.tar.gz", bundle_bytes, "application/gzip")}
data = {
"name": name, "version": version,
"entrypoint": entrypoint, "public": "true" if public else "false",
"description": "Built by agent-builder.",
}
if isinstance(base_head, str) and base_head:
data["base_head_sha"] = base_head
try:
async with httpx.AsyncClient(timeout=60.0) as c:
r = await c.post(
f"{settings.cp_url}/v1/agents/from-tarball",
headers={"authorization": f"bearer {ctx.cp_jwt}"},
files=files, data=data,
)
except httpx.HTTPError as exc:
return json.dumps({"error": f"cp unreachable: {exc}"})
if r.status_code >= 400:
return json.dumps({"error": f"cp {r.status_code}", "detail": r.text[:1000]})
body = r.json()
name_ = body.get("name")
version_ = body.get("version")
url_ = body.get("url")
head_sha = body.get("head_sha")
if isinstance(head_sha, str) and head_sha:
_write_builder_state(
s3,
bucket,
name,
{
"agent": name,
"repo_head_sha": head_sha,
"deployment_id": body.get("deployment_id"),
"version": version_,
"updated_at": int(time.time()),
"source": "agent-builder",
},
)
# CP returns immediately with status="building". The agent isn't
# actually callable until the pod rolls + serves the new card.
# Poll /.well-known/agent-card until the live card reports the
# version we just deployed (or timeout). Surface the live skills
# so the LLM can verify the input_schema actually matches the
# code it wrote — catches the schema/code skew that causes
# downstream 400s on call_agent.
live, live_card = await _wait_for_live_card(url_, version_)
out: dict[str, Any] = {
"ok": live,
"name": name_,
"version": version_,
"status": body.get("status"),
"url": url_,
"head_sha": head_sha,
"deployment_id": body.get("deployment_id"),
"workspace_base_head_sha": base_head,
"live": live,
}
if live:
out["live_version"] = live_card.get("version")
out["live_skills"] = [
{
"name": s.get("name"),
"input_schema": s.get("input_schema") or {},
}
for s in (live_card.get("skills") or [])
]
else:
out["error"] = (
"deploy did not become live within timeout — pod may still"
" be rolling or build failed. Call cp_deploy_tarball again"
" or inspect via discover_agent before declaring done."
)
if live_card:
out["last_seen_version"] = live_card.get("version")
return json.dumps(out)
@tool
async def cp_refresh_agent(name: str) -> str:
"""Force the control plane to re-fetch the agent's live
``/.well-known/agent-card`` and update its stored copy.
Use this when the pod was redeployed out-of-band (CI bump, manual
rollout) and the CP's cached card is stale — symptoms are stale
``input_schema`` in ``list_my_agents`` / ``discover_agent`` and
4xx from ``call_agent`` on what should be a valid skill call.
``cp_deploy_tarball`` already does this automatically by polling
the live card, so you only need ``cp_refresh_agent`` when you
did NOT deploy through this tool.
Returns JSON: the fresh ``AgentDetailOut`` (with ``card`` +
``skills`` + ``input_schema``), or ``{error, ...}``.
"""
if not ctx.cp_jwt:
return json.dumps({
"error": "no CP JWT forwarded — agent declaration missing wants_cp_jwt=True",
})
try:
async with httpx.AsyncClient(timeout=15.0) as c:
r = await c.get(
f"{settings.cp_url}/v1/agents/{name}",
headers={"authorization": f"bearer {ctx.cp_jwt}"},
)
except httpx.HTTPError as exc:
return json.dumps({"error": f"cp unreachable: {exc}"})
if r.status_code >= 400:
return json.dumps({"error": f"cp {r.status_code}", "detail": r.text[:1000]})
body = r.json() or {}
return json.dumps({
"ok": True,
"name": body.get("name"),
"version": body.get("version"),
"status": body.get("status"),
"url": body.get("url"),
"card": body.get("card") or {},
})
@tool
def list_a2a_pack(subdir: str = "") -> str:
"""List ``.py`` files under the installed ``a2a_pack`` package.
Use this when you need to recall what the SDK actually exposes
— what decorators exist, what ``RunContext`` methods are
available, how ``WorkspaceClient`` is shaped, what the
``AgentRuntime`` fields are, etc. The agent code you scaffold
runs on this exact package, so reading it is the authoritative
reference (the system-prompt examples are intentionally
abridged).
``subdir`` is an optional path segment relative to the package
root (``""``, ``"serve"``, ``"mcp"`` …). Returns JSON:
``{root, files: [{path, size}]}``.
"""
try:
import a2a_pack
import os as _os
base = _os.path.dirname(_os.path.abspath(a2a_pack.__file__))
except Exception as exc: # noqa: BLE001
return json.dumps({"error": f"a2a_pack not importable: {exc}"})
if subdir.startswith("/") or ".." in subdir.split("/"):
return json.dumps({"error": "subdir must be relative; '..' disallowed"})
root = _os.path.join(base, subdir) if subdir else base
if not _os.path.commonpath([base, _os.path.abspath(root)]) == base:
return json.dumps({"error": "subdir escapes package root"})
if not _os.path.isdir(root):
return json.dumps({"error": f"not a directory: {subdir}"})
out: list[dict[str, Any]] = []
for dirpath, _, filenames in _os.walk(root):
for fname in filenames:
if not fname.endswith(".py"):
continue
full = _os.path.join(dirpath, fname)
rel = _os.path.relpath(full, base)
try:
size = _os.path.getsize(full)
except OSError:
continue
out.append({"path": rel, "size": size})
out.sort(key=lambda r: r["path"])
return json.dumps({"root": base, "files": out})
@tool
def read_a2a_pack(path: str) -> str:
"""Read a file from the installed ``a2a_pack`` package.
Path is relative to the package root, e.g. ``agent.py``,
``context.py``, ``runtime.py``, ``workspace.py``, ``card.py``,
``mcp/server.py``. Returns JSON: ``{path, content}``. Capped
at 64 KiB; reach for ``list_a2a_pack`` first if you need to
chase a deeper file.
"""
try:
import a2a_pack
import os as _os
base = _os.path.dirname(_os.path.abspath(a2a_pack.__file__))
except Exception as exc: # noqa: BLE001
return json.dumps({"error": f"a2a_pack not importable: {exc}"})
if path.startswith("/") or ".." in path.split("/"):
return json.dumps({"error": "path must be relative; '..' disallowed"})
full = _os.path.abspath(_os.path.join(base, path))
if _os.path.commonpath([base, full]) != base:
return json.dumps({"error": "path escapes package root"})
if not _os.path.isfile(full):
return json.dumps({"error": f"not a file: {path}"})
try:
with open(full, "r", encoding="utf-8") as fh:
content = fh.read(64 * 1024 + 1)
except OSError as exc:
return json.dumps({"error": str(exc)})
truncated = len(content) > 64 * 1024
if truncated:
content = content[: 64 * 1024]
return json.dumps({
"path": path,
"content": content,
"truncated": truncated,
})
return [
init_agent_template, list_agent_files, write_agent_file, read_agent_file,
test_agent_in_sandbox, cp_deploy_tarball, cp_refresh_agent,
list_a2a_pack, read_a2a_pack,
]
# --- 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")
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:
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
return buf.getvalue() if added else b""
def _read_builder_state(s3: Any, bucket: str, 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):
return {}
return data if isinstance(data, dict) else {}
def _write_builder_state(
s3: Any,
bucket: str,
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",
)
async def _latest_cp_deployment(
cp_url: str,
cp_jwt: str,
name: str,
) -> dict[str, Any] | None:
try:
async with httpx.AsyncClient(timeout=15.0) as c:
r = await c.get(
f"{cp_url}/v1/agents/{name}/deployments",
headers={"authorization": f"bearer {cp_jwt}"},
)
except httpx.HTTPError as exc:
raise RuntimeError(f"cp unreachable: {exc}") from exc
if r.status_code == 404:
return None
if r.status_code >= 400:
raise RuntimeError(f"cp {r.status_code}: {r.text[:1000]}")
body = r.json() or []
if not isinstance(body, list) or not body:
return None
first = body[0]
return first if isinstance(first, dict) else None
def _deployment_drift_error(
name: str,
latest_deployment: dict[str, Any] | None,
builder_state: dict[str, Any],
*,
force: bool = False,
) -> dict[str, Any] | None:
if force or latest_deployment is None:
return None
latest_head = latest_deployment.get("head_sha")
base_head = builder_state.get("repo_head_sha")
if isinstance(latest_head, str) and latest_head:
if isinstance(base_head, str) and base_head == latest_head:
return None
return {
"ok": False,
"error": "workspace_drift",
"agent": name,
"message": (
"This MinIO workspace is not based on the latest managed "
"repo deployment. Refresh it before deploying, or use "
"force=True only if the user wants to replace the repo."
),
"workspace_base_head_sha": base_head,
"current_head_sha": latest_head,
"current_deployment_id": latest_deployment.get("deploy_id"),
}
if not isinstance(base_head, str) or not base_head:
return {
"ok": False,
"error": "workspace_untracked",
"agent": name,
"message": (
"This agent already has managed source, but the MinIO "
"workspace has no deploy base marker. Refresh it before "
"deploying, or use force=True only if the user wants to "
"replace the repo."
),
"current_deployment_id": latest_deployment.get("deploy_id"),
}
return None
def _read_entrypoint(bundle: bytes) -> str | None:
try:
with tarfile.open(fileobj=io.BytesIO(bundle), mode="r:gz") as tf:
for member in tf.getmembers():
if member.name.endswith("a2a.yaml") and tf.extractfile(member):
text = tf.extractfile(member).read().decode("utf-8", "replace") # type: ignore[union-attr]
for line in text.splitlines():
if line.lstrip().startswith("entrypoint:"):
return line.split(":", 1)[1].strip().strip('"\'')
except Exception: # noqa: BLE001
pass
return None
def _read_bundle_text(bundle: bytes, path: str) -> str | None:
try:
with tarfile.open(fileobj=io.BytesIO(bundle), mode="r:gz") as tf:
for member in tf.getmembers():
if member.name == path and member.isfile():
extracted = tf.extractfile(member)
if extracted is None:
return None
return extracted.read().decode("utf-8", "replace")
except Exception: # noqa: BLE001
return None
return None
def _class_name(slug: str) -> str:
return "".join(p.capitalize() for p in re.split(r"[-_]+", slug) if p) or "MyAgent"
def _render_a2a_init_template(
name: str,
*,
description: str = "A new A2A agent",
) -> dict[str, str]:
_validate_name(name)
class_name = _class_name(name)
files = {
"agent.py": _render_sdk_template(
"agent.py.tmpl",
name=name,
class_name=class_name,
description=description,
),
"a2a.yaml": _render_sdk_template(
"a2a.yaml.tmpl",
name=name,
class_name=class_name,
),
"requirements.txt": _render_sdk_template("requirements.txt.tmpl"),
}
return files
def _render_sdk_template(template: str, /, **vars: str) -> str:
text = (
resources.files("a2a_pack.cli.templates")
.joinpath(template)
.read_text(encoding="utf-8")
)
for key, value in vars.items():
text = text.replace("{{ " + key + " }}", value)
return text