from __future__ import annotations import io import json import tarfile import unittest from agent_builder.config import Settings 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, ToolContext, build_tools, ) 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.PLATFORM", files["agent.py"]) self.assertIn("WorkspaceAccess.dynamic", files["agent.py"]) self.assertIn("RUNTIME_SKILLS_DIR", 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 test_init_agent_template_writes_through_workspace_store(self) -> None: workspace = _FakeWorkspace() tools = build_tools( ToolContext( bucket="bucket", settings=_settings(), workspace=workspace, ) ) init = _tool_by_name(tools, "init_agent_template") result = json.loads( init.invoke({ "name": "website-scraper", "description": "Browser scraper", "frontend": "none", }) ) self.assertTrue(result["ok"]) self.assertIn("agents/website-scraper/agent.py", workspace.objects) self.assertIn("agents/website-scraper/a2a.yaml", workspace.objects) self.assertIn("agents/website-scraper/requirements.txt", workspace.objects) 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) ] } ] class _FakeWorkspace: def __init__(self) -> None: self.objects: dict[str, bytes] = {} def iter_paths(self) -> list[str]: return sorted(self.objects) def read_bytes(self, key: str) -> bytes: return self.objects[key] def write_bytes(self, key: str, body: bytes) -> None: self.objects[key] = bytes(body) def delete_path(self, key: str) -> None: self.objects.pop(key, None) def _tool_by_name(tools: list[object], name: str) -> object: for tool in tools: if getattr(tool, "name", None) == name: return tool raise AssertionError(f"tool not found: {name}") def _settings() -> Settings: return Settings( sandbox_url="http://sandbox.test", sandbox_timeout_s=1, sandbox_token=None, deploy_wait_timeout_s=1, litellm_url="http://litellm.test", litellm_key="", litellm_model="gpt-test", cp_url="http://cp.test", minio_endpoint="http://minio.test", minio_access_key="key", minio_secret_key="secret", image="python:3.11-slim", )