Agent service filtering + ReDoS fix in transform handlers

- Add stage_services field to AgentContext, populated by _agent_build_context()
  and passed through _apply_governance_check() to reduce false positive
  anti-pattern warnings for irrelevant service namespaces
- Extract _find_azapi_blocks() shared brace-counting helper and rewrite
  _add_response_export_values, _add_resource_group_parent_id, and
  _remove_private_endpoint_resources to eliminate nested-quantifier regex
- 5 new tests (2 service filtering, 3 brace counting safety)

Author: a11smiles

HISTORY.rst

@@ -41,6 +41,22 @@ Generation prompt improvements
   now appears before the architecture context in the generation prompt,
   reducing unused ``terraform_remote_state`` data sources.
 
+Agent-level service filtering
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+* **Agent governance checks now filter by service namespace** — added
+  ``stage_services`` field to ``AgentContext``, populated by
+  ``_agent_build_context()``.  ``_apply_governance_check()`` now passes
+  stage services to ``validate_response()``, reducing false positive
+  anti-pattern warnings for irrelevant service namespaces.
+
+ReDoS fix in transform handlers
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+* **Replaced nested-quantifier regex with brace counting** — extracted
+  shared ``_find_azapi_blocks()`` helper and rewrote
+  ``_add_response_export_values``, ``_add_resource_group_parent_id``,
+  and ``_remove_private_endpoint_resources`` to use it.  Eliminates
+  potential exponential backtracking on pathological input.
+
 Test suite consolidation
 ~~~~~~~~~~~~~~~~~~~~~~~~~~
 * **Consolidated and enhanced unit test coverage** — migrated flat test

azext_prototype/agents/base.py

@@ -81,6 +81,7 @@ class AgentContext:
     artifacts: dict[str, Any] = field(default_factory=dict)
     shared_state: dict[str, Any] = field(default_factory=dict)
     mcp_manager: Any = None  # MCPManager | None — typed as Any to avoid circular import
+    stage_services: list[str] | None = None  # ARM namespaces for service filtering
 
     def add_artifact(self, key: str, value: Any):
         """Store an artifact for other agents to reference."""
@@ -299,7 +300,7 @@ def _apply_governance_check(self, response: AIResponse, context: AgentContext) -
         avoid duplicating the governance warning block.
         """
         iac_tool = context.project_config.get("project", {}).get("iac_tool") if context.project_config else None
-        warnings = self.validate_response(response.content, iac_tool=iac_tool, services=None)
+        warnings = self.validate_response(response.content, iac_tool=iac_tool, services=context.stage_services)
         if warnings:
             for w in warnings:
                 logger.warning("Governance: %s", w)

azext_prototype/governance/transforms/__init__.py

@@ -268,6 +268,32 @@ def _remove_unused_remote_state(content: str, stage_content: str | None = None)
     return result
 
 
+def _find_azapi_blocks(content: str) -> list[tuple[int, int, str, str]]:
+    """Find all ``azapi_resource`` blocks using brace counting.
+
+    Returns a list of ``(start, end, resource_name, block_text)`` tuples
+    where *start*/*end* are character offsets into *content*.
+    """
+    pattern = re.compile(r'resource\s+"azapi_resource"\s+"(\w+)"\s*\{')
+    blocks: list[tuple[int, int, str, str]] = []
+    for match in pattern.finditer(content):
+        name = match.group(1)
+        start = match.start()
+        brace_start = match.end() - 1
+        depth = 1
+        pos = brace_start + 1
+        while pos < len(content) and depth > 0:
+            if content[pos] == "{":
+                depth += 1
+            elif content[pos] == "}":
+                depth -= 1
+            pos += 1
+        if depth != 0:
+            continue  # malformed block
+        blocks.append((start, pos, name, content[start:pos]))
+    return blocks
+
+
 def _remove_private_endpoint_resources(content: str) -> str:
     """Remove private endpoint and DNS zone resources from non-networking stages.
 
@@ -287,42 +313,20 @@ def _remove_private_endpoint_resources(content: str) -> str:
         "virtualnetworklinks",
     )
 
-    # Find resource block starts and use brace counting to find the end
-    block_start_pattern = re.compile(
-        r'resource\s+"azapi_resource"\s+"(\w+)"\s*\{',
-    )
-
     removed_names: list[str] = []
     result = content
 
-    for match in reversed(list(block_start_pattern.finditer(result))):
-        resource_name = match.group(1)
-        # Find the matching closing brace using brace counting
-        start = match.start()
-        brace_start = match.end() - 1  # position of opening {
-        depth = 1
-        pos = brace_start + 1
-        while pos < len(result) and depth > 0:
-            if result[pos] == "{":
-                depth += 1
-            elif result[pos] == "}":
-                depth -= 1
-            pos += 1
-        if depth != 0:
-            continue  # malformed block, skip
-
-        block_text = result[start:pos]
-        # Check if this block's type is a PE/DNS type
+    for start, end, resource_name, block_text in reversed(_find_azapi_blocks(result)):
         type_match = re.search(r'type\s*=\s*"([^"]+)"', block_text)
         if not type_match:
             continue
         resource_type = type_match.group(1).lower()
         if any(pt in resource_type for pt in pe_types):
             # Remove the block plus any trailing whitespace/newlines
-            end = pos
-            while end < len(result) and result[end] in ("\n", "\r", " "):
-                end += 1
-            result = result[:start] + result[end:]
+            trim_end = end
+            while trim_end < len(result) and result[trim_end] in ("\n", "\r", " "):
+                trim_end += 1
+            result = result[:start] + result[trim_end:]
             removed_names.append(resource_name)
             logger.debug("Removed PE/DNS resource: azapi_resource.%s", resource_name)
 
@@ -351,29 +355,26 @@ def _remove_private_endpoint_resources(content: str) -> str:
 def _add_response_export_values(content: str) -> str:
     """Add ``response_export_values = ["*"]`` to azapi_resource blocks missing it.
 
-    Finds each ``resource "azapi_resource" "name" { ... }`` block and checks
-    if ``response_export_values`` appears inside it.  If missing, inserts it
-    after the ``parent_id`` line (or after ``type`` if no ``parent_id``).
+    Uses brace-counting via :func:`_find_azapi_blocks` to avoid nested-quantifier
+    regex (ReDoS risk).  Inserts after ``parent_id``, ``location``, or ``type``.
     """
-    # Match azapi_resource blocks
-    block_pattern = re.compile(
-        r'(resource\s+"azapi_resource"\s+"\w+"\s*\{)(.*?\n)((?:.*?\n)*?)(})',
-        re.DOTALL,
-    )
-
-    def _inject(match: re.Match) -> str:  # type: ignore[type-arg]
-        full = match.group(0)
-        if "response_export_values" in full:
-            return full  # already has it
+    result = content
+    for start, end, _name, block_text in reversed(_find_azapi_blocks(result)):
+        if "response_export_values" in block_text:
+            continue
 
-        header = match.group(1)
-        first_line = match.group(2)
-        body = match.group(3)
-        closing = match.group(4)
+        # Split block body (after the opening { line) into lines
+        header_end = block_text.index("{") + 1
+        header = block_text[:header_end]
+        body_plus_close = block_text[header_end:]
+        # Remove the final closing brace
+        body = body_plus_close.rstrip()
+        if body.endswith("}"):
+            body = body[:-1]
+        closing = "}"
 
-        # Find insertion point: after parent_id, or after location, or after type
-        lines = (first_line + body).splitlines(keepends=True)
-        insert_idx = len(lines)  # fallback: before closing brace
+        lines = body.splitlines(keepends=True)
+        insert_idx = len(lines)
         for i, line in enumerate(lines):
             stripped = line.strip()
             if stripped.startswith("parent_id"):
@@ -384,69 +385,62 @@ def _inject(match: re.Match) -> str:  # type: ignore[type-arg]
             elif stripped.startswith("type") and insert_idx == len(lines):
                 insert_idx = i + 1
 
-        # Detect indentation from the type/parent_id line
         indent = "  "
-        if insert_idx > 0 and insert_idx <= len(lines):
+        if 0 < insert_idx <= len(lines):
             prev_line = lines[insert_idx - 1]
             leading = len(prev_line) - len(prev_line.lstrip())
             indent = " " * leading
 
         lines.insert(insert_idx, f'\n{indent}response_export_values = ["*"]\n')
-        return header + "".join(lines) + closing
+        new_block = header + "".join(lines) + closing
+        result = result[:start] + new_block + result[end:]
 
-    new_content = block_pattern.sub(_inject, content)
-    if new_content != content:
+    if result != content:
         logger.debug("Added response_export_values to azapi_resource blocks")
-    return new_content
+    return result
 
 
 def _add_resource_group_parent_id(content: str) -> str:
     """Add ``parent_id`` to resource group azapi_resource blocks missing it.
 
-    Finds ``azapi_resource`` blocks whose type contains
-    ``Microsoft.Resources/resourceGroups`` and injects
-    ``parent_id = "/subscriptions/${var.subscription_id}"``
-    after the ``name`` line.
+    Uses brace-counting via :func:`_find_azapi_blocks` to avoid nested-quantifier
+    regex (ReDoS risk).  Injects after the ``name`` line.
     """
-    # Match azapi_resource blocks with resourceGroups type
-    block_pattern = re.compile(
-        r'(resource\s+"azapi_resource"\s+"\w+"\s*\{)(.*?)(})',
-        re.DOTALL,
-    )
-
-    def _inject(match: re.Match) -> str:  # type: ignore[type-arg]
-        full = match.group(0)
-        if "resourcegroups" not in full.lower():
-            return full
-        if "parent_id" in full:
-            return full  # already has it
+    result = content
+    for start, end, _name, block_text in reversed(_find_azapi_blocks(result)):
+        if "resourcegroups" not in block_text.lower():
+            continue
+        if "parent_id" in block_text:
+            continue
 
-        header = match.group(1)
-        body = match.group(2)
-        closing = match.group(3)
+        header_end = block_text.index("{") + 1
+        header = block_text[:header_end]
+        body_plus_close = block_text[header_end:]
+        body = body_plus_close.rstrip()
+        if body.endswith("}"):
+            body = body[:-1]
+        closing = "}"
 
-        # Insert after the name line
         lines = body.splitlines(keepends=True)
         insert_idx = len(lines)
         for i, line in enumerate(lines):
             if line.strip().startswith("name"):
                 insert_idx = i + 1
                 break
 
-        # Detect indentation
         indent = "  "
-        if insert_idx > 0 and insert_idx <= len(lines):
+        if 0 < insert_idx <= len(lines):
             prev_line = lines[insert_idx - 1]
             leading = len(prev_line) - len(prev_line.lstrip())
             indent = " " * leading
 
         lines.insert(insert_idx, f'{indent}parent_id = "/subscriptions/${{var.subscription_id}}"\n')
-        return header + "".join(lines) + closing
+        new_block = header + "".join(lines) + closing
+        result = result[:start] + new_block + result[end:]
 
-    new_content = block_pattern.sub(_inject, content)
-    if new_content != content:
+    if result != content:
         logger.debug("Added parent_id to resource group azapi_resource")
-    return new_content
+    return result
 
 
 _STRUCTURED_HANDLERS: dict[str, Callable] = {

azext_prototype/stages/build_session.py

@@ -1843,10 +1843,13 @@ def _agent_build_context(self, agent: Any, stage: dict) -> Iterator[Any]:
         layer = stage.get("layer", "")
         self._apply_governor_brief(agent, stage.get("name", ""), stage.get("services", []), layer)
         self._apply_stage_knowledge(agent, stage)
+        svc_types = [s.get("resource_type", "") for s in stage.get("services", []) if s.get("resource_type")]
+        self._context.stage_services = svc_types or None
         try:
             yield agent
         finally:
             agent.set_knowledge_override("")
+            self._context.stage_services = None
 
     def _apply_stage_knowledge(self, agent: Any, stage: dict) -> None:
         """Set stage-specific knowledge on the agent.

tests/agents/test_agents.py

@@ -1,6 +1,6 @@
 """Tests for azext_prototype.agents — registry, loader, base."""
 
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 import pytest
 import yaml
@@ -507,6 +507,39 @@ def test_cloud_architect_injects_azure_api_version_for_bicep(self):
         assert "deployment-language-bicep" in joined
 
 
+# ------------------------------------------------------------------
+# Agent-level service filtering
+# ------------------------------------------------------------------
+
+
+class TestAgentServiceFiltering:
+    """AgentContext must carry stage_services so _apply_governance_check can filter."""
+
+    def test_agent_context_has_stage_services_field(self):
+        from azext_prototype.agents.base import AgentContext
+
+        ctx = AgentContext(project_config={}, project_dir="/tmp", ai_provider=None)
+        assert ctx.stage_services is None, "Default stage_services should be None"
+
+    def test_apply_governance_check_passes_stage_services(self):
+        from azext_prototype.agents.base import AgentContext
+
+        agent = StubAgent()
+        provider = MagicMock()
+        ctx = AgentContext(project_config={"project": {"iac_tool": "terraform"}}, project_dir="/tmp", ai_provider=provider)
+        ctx.stage_services = ["Microsoft.KeyVault/vaults"]
+
+        response = AIResponse(content="resource content", model="test", usage={})
+
+        with patch.object(agent, "validate_response", return_value=[]) as mock_validate:
+            agent._apply_governance_check(response, ctx)
+            mock_validate.assert_called_once()
+            call_kwargs = mock_validate.call_args
+            assert call_kwargs[1]["services"] == ["Microsoft.KeyVault/vaults"], (
+                f"Expected stage_services to be passed through, got: {call_kwargs[1].get('services')}"
+            )
+
+
 # ------------------------------------------------------------------
 # QA checklist content requirements
 # ------------------------------------------------------------------

tests/governance/transforms/test_transforms.py

@@ -393,3 +393,61 @@ def test_cross_file_unused_remote_state_removed(self):
         )
         assert "TFM-TF-001" in ids
         assert "terraform_remote_state" not in result
+
+
+# ------------------------------------------------------------------
+# ReDoS safety: brace-counting replaces nested quantifier regex
+# ------------------------------------------------------------------
+
+
+class TestBraceCountingSafety:
+    """Transform handlers must use brace counting, not nested-quantifier regex."""
+
+    def test_response_export_values_pathological_input(self):
+        """Long line with no newlines must complete in <1 second (no backtracking)."""
+        import time
+
+        # Pathological: very long body with no newlines, followed by closing brace
+        long_body = "x" * 50000
+        content = f'resource "azapi_resource" "kv" {{\n  type = "Microsoft.KeyVault/vaults@2023-07-01"\n  {long_body}\n}}\n'
+
+        start = time.monotonic()
+        result = _add_response_export_values(content)
+        elapsed = time.monotonic() - start
+
+        assert elapsed < 1.0, f"_add_response_export_values took {elapsed:.2f}s on pathological input (ReDoS?)"
+        assert 'response_export_values = ["*"]' in result
+
+    def test_resource_group_parent_id_pathological_input(self):
+        """Long line with no newlines must complete in <1 second (no backtracking)."""
+        import time
+
+        long_body = "x" * 50000
+        content = f'resource "azapi_resource" "rg" {{\n  type = "Microsoft.Resources/resourceGroups@2024-03-01"\n  name = var.rg\n  {long_body}\n}}\n'
+
+        start = time.monotonic()
+        result = _add_resource_group_parent_id(content)
+        elapsed = time.monotonic() - start
+
+        assert elapsed < 1.0, f"_add_resource_group_parent_id took {elapsed:.2f}s on pathological input (ReDoS?)"
+        assert "parent_id" in result
+
+    def test_find_azapi_blocks_nested_braces(self):
+        """Brace counting must handle nested blocks correctly."""
+        from azext_prototype.governance.transforms import _find_azapi_blocks
+
+        content = """resource "azapi_resource" "kv" {
+  type = "Microsoft.KeyVault/vaults@2023-07-01"
+  body = {
+    properties = {
+      tenantId = var.tenant_id
+    }
+  }
+}
+"""
+        blocks = _find_azapi_blocks(content)
+        assert len(blocks) == 1
+        start, end, name, block_text = blocks[0]
+        assert name == "kv"
+        assert block_text.startswith('resource "azapi_resource"')
+        assert block_text.rstrip().endswith("}")