Add ConsumerSupervisor for background task management

netengine/cli/main.py

@@ -78,6 +78,21 @@ async def _up(spec_file: str, mock: bool, skip_migrations: bool) -> None:
     orchestrator = Orchestrator(spec, mock_mode=mock)
     await orchestrator.execute_phases()
 
+    # Start background consumers (DNS updates, etc.)
+    if orchestrator.consumer_supervisor.consumers:
+        logger.info("All phases complete. Starting background consumers.")
+        await orchestrator.start_consumers()
+        # Keep the event loop alive for background consumers
+        logger.info("Background consumers started. Keeping event loop alive (Ctrl+C to exit).")
+        try:
+            await asyncio.sleep(float("inf"))
+        except KeyboardInterrupt:
+            logger.info("Interrupted by user")
+            await orchestrator.consumer_supervisor.stop_all()
+            logger.info("Consumers stopped")
+    else:
+        logger.info("All phases complete. No background consumers to start.")
+
 
 @cli.command()
 def status() -> None:

netengine/core/consumer_supervisor.py

@@ -0,0 +1,63 @@
+import asyncio
+import logging
+from typing import Any, Callable, Coroutine, Dict
+
+logger = logging.getLogger(__name__)
+
+
+class ConsumerSupervisor:
+    """Manages long-running consumer tasks with automatic restart on failure."""
+
+    def __init__(self) -> None:
+        self.tasks: Dict[str, asyncio.Task[Any]] = {}
+        self.consumers: Dict[str, Callable[[], Coroutine[Any, Any, None]]] = {}
+
+    def register(self, name: str, consumer_coro: Callable[[], Coroutine[Any, Any, None]]) -> None:
+        """Register a consumer coroutine function."""
+        self.consumers[name] = consumer_coro
+
+    async def start_all(self) -> None:
+        """Start all registered consumers."""
+        for name, consumer_func in self.consumers.items():
+            await self.start_consumer(name, consumer_func)
+
+    async def start_consumer(
+        self, name: str, consumer_func: Callable[[], Coroutine[Any, Any, None]]
+    ) -> None:
+        """Start a single consumer with automatic restart on failure."""
+
+        async def supervised_consumer() -> None:
+            while True:
+                try:
+                    logger.info(f"Starting consumer: {name}")
+                    await consumer_func()
+                except asyncio.CancelledError:
+                    logger.info(f"Consumer {name} cancelled")
+                    break
+                except Exception as e:
+                    logger.error(f"Consumer {name} crashed: {e}. Restarting in 5s...")
+                    await asyncio.sleep(5)
+
+        task: asyncio.Task[None] = asyncio.create_task(supervised_consumer())
+        self.tasks[name] = task
+        logger.info(f"Consumer {name} started")
+
+    async def stop_all(self) -> None:
+        """Stop all consumers gracefully."""
+        for name, task in self.tasks.items():
+            logger.info(f"Stopping consumer: {name}")
+            task.cancel()
+            try:
+                await task
+            except asyncio.CancelledError:
+                pass
+
+    def get_status(self) -> Dict[str, str]:
+        """Get status of all consumers."""
+        status: Dict[str, str] = {}
+        for name, task in self.tasks.items():
+            if task.done():
+                status[name] = "crashed" if task.exception() else "completed"
+            else:
+                status[name] = "running"
+        return status

netengine/core/orchestrator.py

@@ -4,6 +4,7 @@
 
 from pydantic import ValidationError
 
+from netengine.core.consumer_supervisor import ConsumerSupervisor
 from netengine.core.state import RuntimeState
 from netengine.handlers._base import BasePhaseHandler
 from netengine.handlers.context import PhaseContext
@@ -70,12 +71,16 @@ def __init__(self, spec: NetEngineSpec | dict[str, Any], mock_mode: Optional[boo
                 logger.warning(f"Docker unavailable, falling back to mock mode: {exc}")
                 effective_mock = True
 
+        # Initialize consumer supervisor for background tasks
+        self.consumer_supervisor = ConsumerSupervisor()
+
         self.context = PhaseContext(
             spec=self.spec,
             runtime_state=self.runtime_state,
             logger=logger,
             docker_client=docker_client,
             mock_mode=effective_mock,
+            consumer_supervisor=self.consumer_supervisor,
         )
 
     @staticmethod
@@ -136,6 +141,10 @@ async def execute_phases(self, up_to_phase: int = 8) -> None:
                 self.runtime_state.save()
                 raise
 
+    async def start_consumers(self) -> None:
+        """Start all registered consumer tasks."""
+        await self.consumer_supervisor.start_all()
+
     def _mark_phase_complete(self, phase_num: int, handler: BasePhaseHandler) -> None:
         """Record user-facing phase completion for a completed handler.
 

netengine/errors.py

@@ -1,10 +1,17 @@
 # TODO: NetEngine exceptions, error handling/reporting/etc module
 
-from .logging import get_logger
+from typing import Any
+
 from omegaconf import DictConfig as Config
 
+
 class BaseNetEngineException(Exception):  # TODO: Comprehensive Engine Exception Base
-    def __init__(self, message: str = "An unknown NetEngine exception occurred.", *args, **kwargs):
+    def __init__(
+        self,
+        message: str = "An unknown NetEngine exception occurred.",
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
         """
 
         :param message:
@@ -13,20 +20,21 @@ def __init__(self, message: str = "An unknown NetEngine exception occurred.", *a
         """
         self._msg = message
         self._code: int | str | None = None
-        self._log_rules: Config | dict = Config({
-            "log_on_init": True,
-            "at_lvl": "TRACE",
-            "with_msg": self.message,
-        })
-        self._log_xt: Config | dict = Config({})
+        self._log_rules: Config | dict[str, Any] = Config(
+            {
+                "log_on_init": True,
+                "at_lvl": "TRACE",
+                "with_msg": self.message,
+            }
+        )
+        self._log_xt: Config | dict[str, Any] = Config({})
 
         if kwargs:
-            for k, v in kwargs:
+            for k, v in kwargs.items():
                 self._log_xt.update({k: v})
 
         super().__init__(self.message)
 
     @property
     def message(self) -> str:
         return self._msg or "An unknown NetEngine exception occurred."
-

netengine/handlers/context.py

@@ -4,11 +4,14 @@
 import os
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Any, Optional
+from typing import TYPE_CHECKING, Any, Optional
 
 from netengine.core.state import RuntimeState
 from netengine.spec.models import NetEngineSpec
 
+if TYPE_CHECKING:
+    from netengine.core.consumer_supervisor import ConsumerSupervisor
+
 # Default directory for CoreDNS Corefile and zone files.
 # Overridden by NETENGINE_ZONE_DIR env var.
 DEFAULT_ZONE_DIR = str(Path.cwd() / "data" / "coredns")
@@ -31,6 +34,7 @@ class PhaseContext:
     kubernetes_client: Any = None
     supabase_client: Any = None
     pgmq_client: Any = None
+    consumer_supervisor: Optional["ConsumerSupervisor"] = None
 
     # Phase-specific config
     phase_name: Optional[str] = None

netengine/handlers/oidc_handler.py

@@ -143,6 +143,26 @@ async def create_user(
         users = await self._admin_request("GET", f"realms/{realm}/users?username={username}")
         return users[0]["id"]
 
+    async def add_token_mapper(
+        self,
+        realm: str,
+        client_id: str,
+        mapper_name: str,
+        protocol_mapper_type: str,
+        config: dict,
+    ) -> str:
+        """Add a protocol mapper (token claim) to a client."""
+        payload = {
+            "name": mapper_name,
+            "protocolMapper": protocol_mapper_type,
+            "protocol": "openid-connect",
+            "config": config,
+        }
+        await self._admin_request(
+            "POST", f"realms/{realm}/clients/{client_id}/protocol-mappers/models", json=payload
+        )
+        return mapper_name
+
     def _generate_secret(self) -> str:
         import secrets
 

netengine/phases/phase_platform_identity.py

@@ -83,12 +83,40 @@ async def execute(self, context: PhaseContext) -> None:
             email="admin@platform.internal",
             password=admin_password,
         )
+
+        # Create platform client for API authentication
+        client_id = await oidc.create_client(
+            realm="platform",
+            client_id="platform-api",
+            name="Platform API",
+            redirect_uris=["https://api.platform.internal/callback"],
+            public=False,
+        )
+
+        # Add token mapper to include org claim in JWT
+        await oidc.add_token_mapper(
+            realm="platform",
+            client_id=client_id,
+            mapper_name="org-claim-mapper",
+            protocol_mapper_type="oidc-usermodel-property-mapper",
+            config={
+                "user.attribute": "org",
+                "claim.name": "org",
+                "jsonType.label": "String",
+                "id.token.claim": "true",
+                "access.token.claim": "true",
+                "userinfo.token.claim": "true",
+            },
+        )
+
         context.runtime_state.platform_realm_id = realm_id
         context.runtime_state.admin_user_id = user_id
+        context.runtime_state.platform_client_id = client_id
         context.runtime_state.identity_platform_output = {
             "keycloak_container_id": container_id,
             "platform_realm_id": realm_id,
             "admin_user_id": user_id,
+            "platform_client_id": client_id,
             "deployed_at": datetime.utcnow().isoformat(),
         }
         context.runtime_state.phase_completed["4"] = True

netengine/phases/phase_registries.py

@@ -54,9 +54,14 @@ async def execute(self, context: PhaseContext) -> None:
                 value=tld["listen_ip"],
             )
 
-        # 5. Wire pgmq consumers (stub – in production, run a loop)
-        # We'll set up a consumer for DNS updates in a background task.
-        asyncio.create_task(self._consume_dns_updates(context))
+        # 5. Wire pgmq consumers
+        # Register consumer with supervisor to run after all phases complete
+        if context.consumer_supervisor:
+            # Create a wrapper that captures the context
+            async def dns_consumer():
+                await self._consume_dns_updates(context)
+
+            context.consumer_supervisor.register("dns_updates", dns_consumer)
 
         # 6. Update state
         context.runtime_state.world_registry_output = {
@@ -92,26 +97,38 @@ async def should_skip(self, context: PhaseContext) -> bool:
 
     async def _consume_dns_updates(self, context: PhaseContext):
         """pgmq consumer: domain.registered -> DNSHandler.add_zone_record."""
+        logger = context.logger
         pgmq = PGMQClient()
         dns = DNSHandler()
         while True:
-            msg = await pgmq.receive("dns_updates")
-            if not msg:
-                await asyncio.sleep(1)
-                continue
             try:
-                envelope = EventEnvelope(**json.loads(msg["message"]))
-                payload = envelope.payload
-                # Add zone record for the domain (e.g., acme.internal -> IP)
-                # For MVP, we add an A record pointing to a placeholder or to the AND gateway.
-                # In real use, the IP would come from the AND allocation.
-                await dns.add_zone_record(
-                    context=context,
-                    zone=payload["domain"],
-                    record_type="A",
-                    name="@",
-                    value="10.0.0.1",  # placeholder – would be replaced with actual IP from AND handler
-                )
-                await pgmq.delete("dns_updates", msg["msg_id"])
+                msg = await pgmq.receive("dns_updates")
+                if not msg:
+                    await asyncio.sleep(1)
+                    continue
+
+                try:
+                    envelope = EventEnvelope(**json.loads(msg["message"]))
+                    payload = envelope.payload
+                    logger.info(f"Processing DNS update for domain: {payload.get('domain')}")
+
+                    # Add zone record for the domain (e.g., acme.internal -> IP)
+                    # For MVP, we add an A record pointing to a placeholder or to the AND gateway.
+                    # In real use, the IP would come from the AND allocation.
+                    await dns.add_zone_record(
+                        context=context,
+                        zone=payload["domain"],
+                        record_type="A",
+                        name="@",
+                        value="10.0.0.1",  # placeholder – would be replaced with actual IP from AND handler
+                    )
+                    await pgmq.delete("dns_updates", msg["msg_id"])
+                    logger.info(
+                        f"Successfully processed DNS update for domain: {payload.get('domain')}"
+                    )
+                except Exception as e:
+                    logger.error(f"Error processing DNS update: {e}")
+                    await pgmq.archive_to_dlq("dns_updates", msg["msg_id"], str(e))
             except Exception as e:
-                await pgmq.archive_to_dlq("dns_updates", msg["msg_id"], str(e))
+                logger.error(f"Error in DNS update consumer loop: {e}")
+                await asyncio.sleep(5)  # backoff before retrying

tests/test_cli.py

@@ -20,6 +20,8 @@ def test_up_invokes_execute_phases_with_example_spec():
     with patch("netengine.cli.main.Orchestrator") as mock_orchestrator_class:
         mock_orchestrator = mock_orchestrator_class.return_value
         mock_orchestrator.execute_phases = AsyncMock()
+        # Mock consumer_supervisor as empty (no consumers registered)
+        mock_orchestrator.consumer_supervisor.consumers = {}
 
         result = CliRunner().invoke(cli_main.cli, ["up", str(spec_file)])