Python: fix(devui): discover entities exported with an annotated assignment

python/packages/devui/agent_framework_devui/_discovery.py

@@ -20,6 +20,19 @@
 logger = logging.getLogger(__name__)
 
 
+def _global_names(scope: ast.FunctionDef | ast.AsyncFunctionDef | ast.ClassDef) -> set[str]:
+    """Return the names a function or class body declares ``global``, not counting nested scopes."""
+    names: set[str] = set()
+    pending: list[ast.AST] = list(scope.body)
+    while pending:
+        node = pending.pop()
+        if isinstance(node, ast.Global):
+            names.update(node.names)
+        elif not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)):
+            pending.extend(ast.iter_child_nodes(node))
+    return names
+
+
 class EntityDiscovery:
     """Discovery for Agent Framework entities - agents and workflows."""
 
@@ -555,12 +568,28 @@ def _has_entity_exports(self, file_path: Path) -> bool:
             source = file_path.read_text(encoding="utf-8")
             tree = ast.parse(source, filename=str(file_path))
 
-            # Look for module-level assignments of 'agent' or 'workflow'
-            for node in ast.walk(tree):
+            # Look for assignments of 'agent' or 'workflow' that create module attributes. In a
+            # function or class body that is only the case for names the scope declares global;
+            # otherwise they are locals or class attributes.
+            exports = {"agent", "workflow"}
+            pending: list[tuple[ast.AST, set[str]]] = [(node, exports) for node in tree.body]
+            while pending:
+                node, names = pending.pop()
+                if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
+                    scope_names = exports & _global_names(node)
+                    pending.extend((child, scope_names) for child in node.body)
+                    continue
+                pending.extend((child, names) for child in ast.iter_child_nodes(node))
                 if isinstance(node, ast.Assign):
-                    for target in node.targets:
-                        if isinstance(target, ast.Name) and target.id in ("agent", "workflow"):
-                            return True
+                    targets = node.targets
+                elif isinstance(node, ast.AnnAssign) and node.value is not None:
+                    # Annotated assignment, e.g. `agent: Agent = Agent(...)`
+                    targets = [node.target]
+                else:
+                    continue
+                for target in targets:
+                    if isinstance(target, ast.Name) and target.id in names:
+                        return True
         except Exception as e:
             logger.debug(f"Could not parse {file_path} for entity exports: {e}")
             return False

python/packages/devui/tests/devui/test_discovery.py

@@ -56,6 +56,27 @@ async def test_empty_directory():
         assert len(entities) == 0
 
 
+async def test_discover_file_entity_with_annotated_export():
+    """A file exporting `agent: Agent = ...` is discovered like `agent = ...`."""
+    with tempfile.TemporaryDirectory() as temp_dir:
+        (Path(temp_dir) / "typed_agent.py").write_text("agent: object = object()\n")
+        (Path(temp_dir) / "declared_only.py").write_text("agent: object\n")
+        (Path(temp_dir) / "conditional_agent.py").write_text("if True:\n    agent: object = object()\n")
+        (Path(temp_dir) / "local_only.py").write_text(
+            "def build():\n    agent: object = object()\n    return agent\n\n\n"
+            "class Holder:\n    workflow = object()\n\n\n"
+            "def other():\n    global name\n    agent = object()\n"
+        )
+        (Path(temp_dir) / "global_factory.py").write_text(
+            "def build():\n    global agent\n    if True:\n        agent = object()\n\n\nbuild()\n"
+        )
+
+        discovery = EntityDiscovery(temp_dir)
+        entities = await discovery.discover_entities()
+
+        assert sorted(entity.id for entity in entities) == ["conditional_agent", "global_factory", "typed_agent"]
+
+
 async def test_discovery_accepts_agents_with_only_run():
     """Test that discovery accepts agents with only run() method.