Make site count telemetry authoritative and tested

agent/pipeline_engine.py

@@ -29,6 +29,8 @@ from .netbox_client import NetBoxClient
 from agent.network_ops import create_gns3_project, build_network_from_description, get_lab_topology
 
 logger = logging.getLogger(__name__)
+SITE_COUNT_MIN = 1
+SITE_COUNT_MAX = 10
 
 
 @dataclass
@@ -249,29 +251,29 @@ class OvergrowthPipeline:
         text_lower = text.lower()
         for word, val in words.items():
             if f"{word}-site" in text_lower or f"{word} site" in text_lower:
-                return max(1, min(10, val))
+                return max(SITE_COUNT_MIN, min(SITE_COUNT_MAX, val))
         m = re.search(r"(\d+)[-\s]?site", text_lower)
         if m:
             try:
-                return max(1, min(10, int(m.group(1))))
+                return max(SITE_COUNT_MIN, min(SITE_COUNT_MAX, int(m.group(1))))
             except Exception:
                 pass
         m = re.search(r"(\d+)\s*sites?", text_lower)
         if m:
             try:
-                return max(1, min(10, int(m.group(1))))
+                return max(SITE_COUNT_MIN, min(SITE_COUNT_MAX, int(m.group(1))))
             except Exception:
                 pass
         m = re.search(r"(\d+)\s*(locations?|branches?)", text_lower)
         if m:
             try:
-                return max(1, min(10, int(m.group(1))))
+                return max(SITE_COUNT_MIN, min(SITE_COUNT_MAX, int(m.group(1))))
             except Exception:
                 pass
         m = re.search(r"(\d+)\s*(stores?|shops?)", text_lower)
         if m:
             try:
-                return max(1, min(10, int(m.group(1))))
+                return max(SITE_COUNT_MIN, min(SITE_COUNT_MAX, int(m.group(1))))
             except Exception:
                 pass
         return 3
@@ -280,6 +282,51 @@ class OvergrowthPipeline:
         # Match the naming used by the GNS3 builder (store clouds)
         return [f"☕-Store-{idx}" for idx in range(1, site_cnt + 1)]
 
+    def _count_branch_sites(
+        self,
+        nodes: List[Dict[str, Any]],
+        expected_site_count: Optional[int] = None
+    ) -> Dict[str, Any]:
+        """
+        Count branch/store sites using strict naming from the retail builder.
+        Only counts clouds named ☕-Store-N and switches named SW-Store-N.
+        """
+        import re
+
+        cloud_pat = re.compile(r"^☕-Store-(\d+)$")
+        sw_pat = re.compile(r"^SW-Store-(\d+)$")
+        cloud_sites: List[str] = []
+        switch_sites: List[str] = []
+        site_indices: set[int] = set()
+
+        for node in nodes or []:
+            name = str(node.get("name", "")).strip()
+            node_type = (node.get("node_type") or "").lower()
+            m_cloud = cloud_pat.match(name)
+            m_sw = sw_pat.match(name)
+            if m_cloud and node_type == "cloud":
+                idx = int(m_cloud.group(1))
+                site_indices.add(idx)
+                cloud_sites.append(name)
+            elif m_sw:
+                idx = int(m_sw.group(1))
+                site_indices.add(idx)
+                switch_sites.append(name)
+
+        site_count = len(site_indices)
+        missing: List[str] = []
+        if expected_site_count:
+            expected_names = [f"☕-Store-{i}" for i in range(1, expected_site_count + 1)]
+            missing = [name for name in expected_names if name not in cloud_sites]
+
+        return {
+            "count": site_count,
+            "indices": sorted(site_indices),
+            "cloud_sites": sorted(cloud_sites),
+            "switch_sites": sorted(switch_sites),
+            "missing": missing,
+        }
+
     def _ensure_wireguard(self):
         """
         If WG_CONFIG_B64 is provided, bring up wg0 using that config.
@@ -343,6 +390,40 @@ class OvergrowthPipeline:
             logger.error(f"WireGuard setup error: {e}")
             self.wg_status = f"error: {e}"
 
+    def _extract_builder_summary(self, build_resp: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Extract builder summary (site counts, nodes) from the MCP response.
+        Handles both parsed dicts and JSON-encoded 'text' fields.
+        """
+        candidates: List[Dict[str, Any]] = []
+        if isinstance(build_resp, dict):
+            candidates.append(build_resp)
+            txt = build_resp.get("text")
+            if isinstance(txt, str):
+                try:
+                    candidates.append(json.loads(txt))
+                except Exception:
+                    pass
+            raw = build_resp.get("raw", {})
+            if isinstance(raw, dict):
+                content = raw.get("content") or []
+                if content and isinstance(content, list) and isinstance(content[0], dict):
+                    txt2 = content[0].get("text")
+                    if isinstance(txt2, str):
+                        try:
+                            candidates.append(json.loads(txt2))
+                        except Exception:
+                            pass
+
+        for cand in candidates:
+            if isinstance(cand, dict) and (
+                "site_count_built" in cand
+                or "site_count_requested" in cand
+                or "nodes" in cand
+            ):
+                return cand
+        return {}
+
     def _prepare_lab_info(self, intent: NetworkIntent, create_project: bool = False) -> Dict[str, Any]:
         """
         Return lab link info and optionally create a project via MCP to get the project_id.
@@ -1782,27 +1863,8 @@ Be specific and practical. Use RFC1918 addressing. Consider scalability and secu
                     return []
             return []
 
-        def _count_branch_sites(nodes: List[Dict[str, Any]]) -> Dict[str, Any]:
-            import re
-            pattern = re.compile(r"^(?:☕-)?store[-\s]", re.IGNORECASE)
-            cloud_sites = {
-                str(n.get("name", ""))
-                for n in nodes
-                if pattern.match(str(n.get("name", "")))
-                and (n.get("node_type") or "").lower() == "cloud"
-            }
-            switch_sites = {
-                str(n.get("name", ""))
-                for n in nodes
-                if str(n.get("name", "")).startswith("SW-Store-")
-            }
-            return {
-                "count": max(len(cloud_sites), len(switch_sites)),
-                "cloud_sites": sorted(cloud_sites),
-                "switch_sites": sorted(switch_sites),
-            }
-
         if self.gns3_server and self.enable_gns3_build and lab_info:
+            builder_site_count: Optional[int] = None
             try:
                 base_desc = (
                     intent.description
@@ -1819,63 +1881,54 @@ Be specific and practical. Use RFC1918 addressing. Consider scalability and secu
                 )
                 lab_info["build_result"] = build_resp
                 nodes = _nodes_from_resp(build_resp)
-                site_summary_resp = _count_branch_sites(nodes)
-                node_sites = {
-                    s.lower() for s in (site_summary_resp.get("cloud_sites") or site_summary_resp.get("switch_sites") or [])
-                }
-                missing_sites = [s for s in site_names if s.lower() not in node_sites]
-                lab_info["missing_sites"] = missing_sites
-                if missing_sites:
-                    desc_retry = base_desc + " Second attempt: ensure these sites exist: " + ", ".join(missing_sites)
-                    retry_resp = build_network_from_description(
-                        description=desc_retry,
-                        project_name=lab_info.get("project_name") or "overgrowth",
-                        auto_configure=True,
-                        site_count=site_count,
-                    )
-                    lab_info["build_retry"] = retry_resp
+                builder_summary = self._extract_builder_summary(build_resp)
+                builder_site_count = builder_summary.get("site_count_built")
+                lab_info["builder_site_count"] = builder_site_count
+                results["builder_site_count"] = builder_site_count
+                lab_info["builder_summary"] = builder_summary
+                site_summary_resp = self._count_branch_sites(nodes, expected_site_count=site_count)
+                lab_info["initial_site_count"] = site_summary_resp.get("count")
+                lab_info["initial_missing_sites"] = site_summary_resp.get("missing", [])
             except Exception as e:
                 lab_info["build_error"] = str(e)
 
             try:
                 topology = get_lab_topology(lab_info.get("project_name") or "overgrowth")
-                site_summary = _count_branch_sites(topology.get("nodes", []))
-                actual_count = site_summary["count"]
-                lab_info["actual_site_count"] = actual_count
+                site_summary = self._count_branch_sites(topology.get("nodes", []), expected_site_count=site_count)
+                topology_count = site_summary["count"]
+                missing_sites = site_summary.get("missing") or []
+                lab_info["topology_site_count"] = topology_count
+                results["topology_site_count"] = topology_count
+                lab_info["missing_sites"] = missing_sites
                 lab_info["site_nodes"] = site_summary["cloud_sites"] or site_summary["switch_sites"]
                 lab_info["topology"] = {
                     "nodes": topology.get("nodes", []),
                     "links": topology.get("links", []),
                 }
+                actual_count = builder_site_count if builder_site_count is not None else topology_count
+                lab_info["actual_site_count"] = actual_count
                 results["actual_site_count"] = actual_count
-                if actual_count < site_count:
-                    mismatch = {"requested": site_count, "actual": actual_count}
+
+                mismatch = None
+                if (
+                    builder_site_count is not None
+                    and builder_site_count != site_count
+                ) or (topology_count is not None and topology_count != site_count):
+                    mismatch = {
+                        "requested": site_count,
+                        "builder": builder_site_count,
+                        "topology": topology_count,
+                    }
+                elif builder_site_count is not None and topology_count is not None and builder_site_count != topology_count:
+                    mismatch = {
+                        "requested": site_count,
+                        "builder": builder_site_count,
+                        "topology": topology_count,
+                    }
+
+                if mismatch:
                     lab_info["site_count_mismatch"] = mismatch
                     results["site_count_mismatch"] = mismatch
-                    if not lab_info.get("build_retry"):
-                        retry_resp = build_network_from_description(
-                            description=base_desc + " Retry to satisfy branch count.",
-                            project_name=lab_info.get("project_name") or "overgrowth",
-                            auto_configure=True,
-                            site_count=site_count,
-                        )
-                        lab_info["build_retry"] = retry_resp
-                        retry_topology = get_lab_topology(lab_info.get("project_name") or "overgrowth")
-                        retry_summary = _count_branch_sites(retry_topology.get("nodes", []))
-                        lab_info["actual_site_count"] = retry_summary["count"]
-                        lab_info["site_nodes"] = retry_summary["cloud_sites"] or retry_summary["switch_sites"]
-                        lab_info["topology"] = {
-                            "nodes": retry_topology.get("nodes", []),
-                            "links": retry_topology.get("links", []),
-                        }
-                        results["actual_site_count"] = retry_summary["count"]
-                        if retry_summary["count"] >= site_count:
-                            lab_info.pop("site_count_mismatch", None)
-                            results.pop("site_count_mismatch", None)
-                        else:
-                            mismatch = {"requested": site_count, "actual": retry_summary["count"]}
-                            lab_info["site_count_mismatch"] = mismatch
-                            results["site_count_mismatch"] = mismatch
             except Exception as e:
                 lab_info["topology_error"] = str(e)
             results["lab"] = lab_info

app.py

@@ -607,16 +607,23 @@ def build_ui():
                     if isinstance(lab_info, dict):
                         missing_sites = lab_info.get("missing_sites") or []
                         requested_sites = lab_info.get("requested_site_count") or results.get("requested_site_count")
+                        builder_sites = lab_info.get("builder_site_count") or results.get("builder_site_count")
+                        topology_sites = lab_info.get("topology_site_count") or results.get("topology_site_count")
                         actual_sites = lab_info.get("actual_site_count") or results.get("actual_site_count")
                         mismatch = lab_info.get("site_count_mismatch") or results.get("site_count_mismatch")
-                        if requested_sites:
-                            status += f"- Requested site count: {requested_sites}\n"
+                        if requested_sites is not None:
+                            status += f"- Requested site count (parsed): {requested_sites}\n"
+                        if builder_sites is not None:
+                            status += f"- Builder site count (GNS3): {builder_sites}\n"
+                        if topology_sites is not None:
+                            status += f"- Topology site count (telemetry): {topology_sites}\n"
                         if actual_sites is not None:
                             status += f"- Actual site count: {actual_sites}\n"
                         if mismatch:
                             status += (
                                 f"- Site count mismatch: requested {mismatch.get('requested')}, "
-                                f"actual {mismatch.get('actual')}\n"
+                                f"builder {mismatch.get('builder')}, "
+                                f"topology {mismatch.get('topology')}\n"
                             )
                         build_err = lab_info.get("build_error") or (lab_info.get("build_result") or {}).get("error")
                         if build_err:

tests/test_site_counts.py

@@ -0,0 +1,61 @@
+#!/usr/bin/env python3
+"""
+Unit tests for site count parsing and topology counting helpers.
+"""
+
+import unittest
+from agent.pipeline_engine import OvergrowthPipeline
+
+
+class SiteCountTests(unittest.TestCase):
+    def setUp(self):
+        # Keep NetBox disabled for fast, offline unit tests.
+        self.pipeline = OvergrowthPipeline(use_netbox=False)
+
+    def test_parse_site_count_variants(self):
+        self.assertEqual(
+            self.pipeline._parse_site_count("Design a 4-site retail network with WiFi"),
+            4,
+        )
+        self.assertEqual(
+            self.pipeline._parse_site_count("We need a seven site deployment"), 7
+        )
+        self.assertEqual(
+            self.pipeline._parse_site_count("Build 9 stores across the region"), 9
+        )
+
+    def test_site_names_alignment(self):
+        self.assertEqual(
+            self.pipeline._site_names(3),
+            ["☕-Store-1", "☕-Store-2", "☕-Store-3"],
+        )
+
+    def test_count_branch_sites_exact(self):
+        nodes = [
+            {"name": "☕-Store-1", "node_type": "cloud"},
+            {"name": "☕-Store-2", "node_type": "cloud"},
+            {"name": "SW-Store-1", "node_type": "qemu"},
+            {"name": "SW-Store-2", "node_type": "qemu"},
+            {"name": "☁️-Internet-MPLS", "node_type": "cloud"},
+            {"name": "-HQ-Building", "node_type": "cloud"},
+        ]
+        summary = self.pipeline._count_branch_sites(nodes, expected_site_count=2)
+        self.assertEqual(summary["count"], 2)
+        self.assertEqual(summary["missing"], [])
+
+    def test_count_branch_sites_missing(self):
+        nodes = [
+            {"name": "☕-Store-1", "node_type": "cloud"},
+            {"name": "☕-Store-2", "node_type": "cloud"},
+            {"name": "☕-Store-3", "node_type": "cloud"},
+            {"name": "SW-Store-1", "node_type": "qemu"},
+            {"name": "SW-Store-2", "node_type": "qemu"},
+            {"name": "SW-Store-3", "node_type": "qemu"},
+        ]
+        summary = self.pipeline._count_branch_sites(nodes, expected_site_count=4)
+        self.assertEqual(summary["count"], 3)
+        self.assertEqual(summary["missing"], ["☕-Store-4"])
+
+
+if __name__ == "__main__":
+    unittest.main()