262 lines
9.2 KiB
Python
262 lines
9.2 KiB
Python
from __future__ import annotations
|
|
|
|
import io
|
|
import tarfile
|
|
import unittest
|
|
|
|
from agent_builder.tools import (
|
|
_BUILDER_STATE_FILE,
|
|
_BUILDER_INTERNAL_PREFIX,
|
|
_deployment_drift_error,
|
|
_files_from_tarball,
|
|
_parse_supporting_skill_files,
|
|
_replace_workspace_files,
|
|
_render_a2a_init_template,
|
|
_render_skill_md,
|
|
_tarball_workspace_dir,
|
|
)
|
|
|
|
|
|
class TemplateInitTests(unittest.TestCase):
|
|
def test_render_a2a_init_template_matches_sdk_starter(self) -> None:
|
|
files = _render_a2a_init_template(
|
|
"research-agent",
|
|
description="Research helper",
|
|
)
|
|
|
|
self.assertEqual(set(files), {"agent.py", "a2a.yaml", "requirements.txt"})
|
|
self.assertIn('name = "research-agent"', files["agent.py"])
|
|
self.assertIn('description = "Research helper"', files["agent.py"])
|
|
self.assertIn("LLMProvisioning.CALLER_PROVIDED", files["agent.py"])
|
|
self.assertIn("WorkspaceAccess.dynamic", files["agent.py"])
|
|
self.assertIn("RUNTIME_SKILLS_ROOT", files["agent.py"])
|
|
self.assertIn("DEEPAGENTS_RECURSION_LIMIT = 500", files["agent.py"])
|
|
self.assertIn(
|
|
'config={"recursion_limit": DEEPAGENTS_RECURSION_LIMIT}',
|
|
files["agent.py"],
|
|
)
|
|
self.assertIn("skills=skill_sources or None", files["agent.py"])
|
|
self.assertIn("ctx.llm", files["agent.py"])
|
|
self.assertIn("ctx.workspace_backend()", files["agent.py"])
|
|
self.assertIn("create_deep_agent", files["agent.py"])
|
|
self.assertIn("name: research-agent", files["a2a.yaml"])
|
|
self.assertIn("entrypoint: agent:ResearchAgent", files["a2a.yaml"])
|
|
self.assertNotIn("{{ frontend_block }}", files["a2a.yaml"])
|
|
self.assertNotIn("frontend:", files["a2a.yaml"])
|
|
self.assertIn("deepagents>=0.5.0", files["requirements.txt"])
|
|
|
|
def test_render_a2a_init_template_can_include_react_frontend(self) -> None:
|
|
files = _render_a2a_init_template(
|
|
"chart-agent",
|
|
description="Chart helper",
|
|
frontend="react",
|
|
)
|
|
|
|
self.assertTrue(
|
|
{
|
|
"agent.py",
|
|
"a2a.yaml",
|
|
"requirements.txt",
|
|
"frontend/package.json",
|
|
"frontend/index.html",
|
|
"frontend/vite.config.js",
|
|
"frontend/src/main.jsx",
|
|
"frontend/src/App.jsx",
|
|
"frontend/src/a2a.js",
|
|
"frontend/src/style.css",
|
|
}.issubset(files)
|
|
)
|
|
self.assertIn("frontend:", files["a2a.yaml"])
|
|
self.assertIn("build: npm run build", files["a2a.yaml"])
|
|
self.assertIn("mount: /app", files["a2a.yaml"])
|
|
self.assertIn("auth: inherit", files["a2a.yaml"])
|
|
self.assertIn("Skill runner", files["frontend/src/App.jsx"])
|
|
self.assertIn("callSkill", files["frontend/src/App.jsx"])
|
|
self.assertIn("CONFIG_URL", files["frontend/src/a2a.js"])
|
|
self.assertIn('"/invoke": agent', files["frontend/vite.config.js"])
|
|
|
|
def test_render_a2a_init_template_can_include_static_frontend(self) -> None:
|
|
files = _render_a2a_init_template("status-agent", frontend="static")
|
|
|
|
self.assertIn("frontend/dist/index.html", files)
|
|
self.assertIn("frontend:", files["a2a.yaml"])
|
|
self.assertIn("dist: dist", files["a2a.yaml"])
|
|
self.assertIn("mount: /app", files["a2a.yaml"])
|
|
self.assertNotIn("build: npm run build", files["a2a.yaml"])
|
|
self.assertIn("./a2a-client.js", files["frontend/dist/index.html"])
|
|
|
|
def test_render_a2a_init_template_rejects_invalid_names(self) -> None:
|
|
with self.assertRaises(ValueError):
|
|
_render_a2a_init_template("Bad Name")
|
|
|
|
def test_render_a2a_init_template_rejects_invalid_frontend(self) -> None:
|
|
with self.assertRaises(ValueError):
|
|
_render_a2a_init_template("demo-agent", frontend="vue")
|
|
|
|
def test_agent_projects_can_diverge_from_template(self) -> None:
|
|
files = _render_a2a_init_template("regex-explainer")
|
|
files["agent.py"] = "from a2a_pack import A2AAgent\n"
|
|
|
|
self.assertIn("A2AAgent", files["agent.py"])
|
|
|
|
def test_deployment_drift_error_blocks_untracked_existing_workspace(self) -> None:
|
|
latest = {"deploy_id": "dpl_1", "head_sha": "abc1234"}
|
|
|
|
err = _deployment_drift_error("demo-agent", latest, {}, force=False)
|
|
|
|
self.assertIsNotNone(err)
|
|
assert err is not None
|
|
self.assertEqual(err["error"], "workspace_untracked")
|
|
self.assertEqual(err["current_head_sha"], "abc1234")
|
|
|
|
def test_deployment_drift_error_allows_tracked_or_forced_workspace(self) -> None:
|
|
latest = {"deploy_id": "dpl_1", "head_sha": "abc1234"}
|
|
|
|
self.assertIsNone(
|
|
_deployment_drift_error(
|
|
"demo-agent",
|
|
latest,
|
|
{"repo_head_sha": "abc1234"},
|
|
force=False,
|
|
)
|
|
)
|
|
self.assertIsNone(
|
|
_deployment_drift_error(
|
|
"demo-agent",
|
|
latest,
|
|
{"repo_head_sha": "repohead999"},
|
|
force=False,
|
|
)
|
|
)
|
|
self.assertIsNone(
|
|
_deployment_drift_error(
|
|
"demo-agent",
|
|
latest,
|
|
{"repo_head_sha": "old9999"},
|
|
force=True,
|
|
)
|
|
)
|
|
|
|
def test_tarball_workspace_excludes_builder_state(self) -> None:
|
|
prefix = "agents/demo-agent/"
|
|
s3 = _FakeS3({
|
|
prefix + "agent.py": b"print('ok')\n",
|
|
prefix + _BUILDER_STATE_FILE: b'{"repo_head_sha":"abc1234"}',
|
|
prefix + _BUILDER_INTERNAL_PREFIX + "skills/x/SKILL.md": b"hidden",
|
|
})
|
|
|
|
bundle = _tarball_workspace_dir(s3, "bucket", prefix)
|
|
|
|
with tarfile.open(fileobj=io.BytesIO(bundle), mode="r:gz") as tf:
|
|
self.assertEqual(tf.getnames(), ["agent.py"])
|
|
|
|
def test_render_skill_md_creates_valid_frontmatter(self) -> None:
|
|
text = _render_skill_md(
|
|
"market-research",
|
|
"Plan market research and synthesize findings.",
|
|
"Use subagents for separate market segments.",
|
|
)
|
|
|
|
self.assertIn("name: market-research", text)
|
|
self.assertIn('description: "Plan market research', text)
|
|
self.assertIn("# market-research", text)
|
|
|
|
def test_parse_supporting_skill_files_rejects_unsafe_paths(self) -> None:
|
|
with self.assertRaises(ValueError):
|
|
_parse_supporting_skill_files('{"../secrets.txt": "bad"}')
|
|
|
|
files = _parse_supporting_skill_files(
|
|
'{"references/schema.md": "schema", "scripts/check.py": "print(1)"}'
|
|
)
|
|
self.assertEqual(
|
|
sorted(files),
|
|
["references/schema.md", "scripts/check.py"],
|
|
)
|
|
|
|
def test_files_from_tarball_rejects_unsafe_paths(self) -> None:
|
|
bundle = _tarball({"../agent.py": "bad"})
|
|
|
|
with self.assertRaises(ValueError):
|
|
_files_from_tarball(bundle)
|
|
|
|
def test_replace_workspace_files_deletes_stale_objects(self) -> None:
|
|
prefix = "agents/demo-agent/"
|
|
s3 = _FakeS3({
|
|
prefix + "old.py": b"old",
|
|
prefix + _BUILDER_STATE_FILE: b"{}",
|
|
})
|
|
|
|
written = _replace_workspace_files(
|
|
s3,
|
|
"bucket",
|
|
prefix,
|
|
{"agent.py": b"print('ok')\n", "a2a.yaml": b"name: demo-agent\n"},
|
|
)
|
|
|
|
self.assertEqual(
|
|
sorted(s3.objects),
|
|
[prefix + "a2a.yaml", prefix + "agent.py"],
|
|
)
|
|
self.assertEqual(
|
|
written,
|
|
[
|
|
{"path": "a2a.yaml", "size": len(b"name: demo-agent\n")},
|
|
{"path": "agent.py", "size": len(b"print('ok')\n")},
|
|
],
|
|
)
|
|
|
|
|
|
def _tarball(files: dict[str, str]) -> bytes:
|
|
buf = io.BytesIO()
|
|
with tarfile.open(fileobj=buf, mode="w:gz") as tf:
|
|
for path, text in files.items():
|
|
body = text.encode("utf-8")
|
|
info = tarfile.TarInfo(path)
|
|
info.size = len(body)
|
|
tf.addfile(info, io.BytesIO(body))
|
|
return buf.getvalue()
|
|
|
|
|
|
class _FakeS3:
|
|
def __init__(self, objects: dict[str, bytes]) -> None:
|
|
self.objects = objects
|
|
|
|
def get_paginator(self, name: str) -> "_FakePaginator":
|
|
assert name == "list_objects_v2"
|
|
return _FakePaginator(self.objects)
|
|
|
|
def get_object(self, *, Bucket: str, Key: str) -> dict[str, io.BytesIO]:
|
|
return {"Body": io.BytesIO(self.objects[Key])}
|
|
|
|
def put_object(
|
|
self,
|
|
*,
|
|
Bucket: str,
|
|
Key: str,
|
|
Body: bytes,
|
|
ContentType: str,
|
|
) -> None:
|
|
self.objects[Key] = bytes(Body)
|
|
|
|
def delete_object(self, *, Bucket: str, Key: str) -> None:
|
|
self.objects.pop(Key, None)
|
|
|
|
def delete_objects(self, *, Bucket: str, Delete: dict[str, object]) -> None:
|
|
raise AssertionError("batch DeleteObjects should not be used")
|
|
|
|
|
|
class _FakePaginator:
|
|
def __init__(self, objects: dict[str, bytes]) -> None:
|
|
self.objects = objects
|
|
|
|
def paginate(self, *, Bucket: str, Prefix: str) -> list[dict[str, object]]:
|
|
return [
|
|
{
|
|
"Contents": [
|
|
{"Key": key, "Size": len(value)}
|
|
for key, value in sorted(self.objects.items())
|
|
if key.startswith(Prefix)
|
|
]
|
|
}
|
|
]
|