561 lines
21 KiB
Python
561 lines
21 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}$")
|
|
|
|
|
|
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 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 ():
|
|
out.append({
|
|
"path": obj["Key"][len(prefix):],
|
|
"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)})
|
|
key = prefix + path.lstrip("/")
|
|
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)})
|
|
key = prefix + path.lstrip("/")
|
|
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) -> 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``.
|
|
|
|
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",
|
|
})
|
|
|
|
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.",
|
|
}
|
|
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")
|
|
|
|
# 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_,
|
|
"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()
|
|
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
|
|
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))
|
|
return buf.getvalue() if buf.getbuffer().nbytes > 0 else b""
|
|
|
|
|
|
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
|