From c24676356d1f56e3d392b7b54d8203df537a62a2 Mon Sep 17 00:00:00 2001
From: robert <robert@microcash.local>
Date: Sun, 7 Jun 2026 09:06:13 -0300
Subject: [PATCH] Fix builder graph state capture

---
 agent.py                     |  1 +
 tests/test_builder_prompt.py | 72 +++++++++++++++++++++++++++++++++++-
 2 files changed, 72 insertions(+), 1 deletion(-)

diff --git a/agent.py b/agent.py
index 87c1243..a3907de 100644
--- a/agent.py
+++ b/agent.py
@@ -167,6 +167,7 @@ class AgentBuilder(A2AAgent[BuilderConfig, NoAuth]):
                 await ctx.emit_progress(f"agent-builder still working on {name!r}")
 
         async def _run_build_graph() -> None:
+            nonlocal deployed_url, final_state
             async for event in graph.astream_events(
                 {"messages": [{"role": "user", "content": user_msg}]},
                 version="v2",
diff --git a/tests/test_builder_prompt.py b/tests/test_builder_prompt.py
index 62b56e4..c42a41d 100644
--- a/tests/test_builder_prompt.py
+++ b/tests/test_builder_prompt.py
@@ -1,7 +1,9 @@
 from __future__ import annotations
 
 from pathlib import Path
+from types import SimpleNamespace
 import unittest
+from unittest.mock import patch
 
 from agent_builder.builder import (
     BUILDER_SKILL_SOURCE,
@@ -15,7 +17,75 @@ from agent_builder.builder import (
 from agent_builder.config import Settings
 from agent_builder.tools import A2A_PACK_MIN_VERSION, _strip_unsupported_pricing_fields
 from deepagents.backends import StateBackend
-from agent import _build_failure_error
+from agent import AgentBuilder, _build_failure_error
+
+
+class FakeWorkspace:
+    bucket = "user-1-files"
+    _grant_token = "grant-token"
+
+
+class FakeRunContext:
+    workspace = FakeWorkspace()
+    cp_jwt = "cp-jwt"
+    sandbox = None
+    llm = SimpleNamespace(
+        model="gpt-5.5",
+        source="test",
+        base_url="http://litellm/v1",
+        api_key="sk-test",
+        temperature_mode="default",
+        temperature=None,
+        extra_body={},
+    )
+
+    def __init__(self) -> None:
+        self.progress: list[str] = []
+
+    async def emit_progress(self, message: str) -> None:
+        self.progress.append(message)
+
+
+class FakeGraph:
+    async def astream_events(self, *_args, **_kwargs):
+        yield {
+            "event": "on_tool_end",
+            "name": "cp_deploy_tarball",
+            "data": {"output": {"url": "https://generated-agent.a2acloud.io"}},
+        }
+        yield {
+            "event": "on_chain_end",
+            "parent_ids": [],
+            "data": {
+                "output": {
+                    "messages": [
+                        {
+                            "content": (
+                                "Deployed at https://generated-agent.a2acloud.io"
+                            )
+                        }
+                    ]
+                }
+            },
+        }
+
+
+class AgentBuilderBuildTests(unittest.IsolatedAsyncioTestCase):
+    async def test_build_captures_nested_graph_state_and_deployed_url(self) -> None:
+        ctx = FakeRunContext()
+
+        with patch("agent.build_agent_builder", return_value=FakeGraph()):
+            result = await AgentBuilder().build(
+                ctx,
+                name="generated-agent",
+                prompt="make a test agent",
+                public=False,
+            )
+
+        self.assertEqual(result["ok"], True)
+        self.assertEqual(result["url"], "https://generated-agent.a2acloud.io")
+        self.assertIn("Deployed at https://generated-agent.a2acloud.io", result["reply"])
+        self.assertIn("deployed: https://generated-agent.a2acloud.io", ctx.progress)
 
 
 class BuilderPromptTests(unittest.TestCase):