omnibase_infra 0.2.1__py3-none-any.whl → 0.2.3__py3-none-any.whl

omnibase_infra/handlers/handler_vault.py

@@ -27,6 +27,7 @@ from uuid import uuid4
 
 import hvac
 
+from omnibase_core.container import ModelONEXContainer
 from omnibase_core.models.dispatch import ModelHandlerOutput
 from omnibase_infra.enums import (
     EnumHandlerType,
@@ -122,13 +123,18 @@ class HandlerVault(
         - MixinVaultToken: Token management and renewal
     """
 
-    def __init__(self) -> None:
-        """Initialize HandlerVault in uninitialized state.
+    def __init__(self, container: ModelONEXContainer) -> None:
+        """Initialize HandlerVault with ONEX container for dependency injection.
 
-        Note: Circuit breaker is initialized during initialize() call when
-        configuration is available. The mixin's _init_circuit_breaker() method
-        is called there with the actual config values.
+        Args:
+            container: ONEX container for dependency injection.
+
+        Note:
+            Circuit breaker is initialized during initialize() call when
+            configuration is available. The mixin's _init_circuit_breaker() method
+            is called there with the actual config values.
         """
+        self._container = container
         self._client: hvac.Client | None = None
         self._config: ModelVaultHandlerConfig | None = None
         self._initialized: bool = False

omnibase_infra/handlers/mixins/mixin_consul_kv.py

@@ -94,7 +94,7 @@ class MixinConsulKV:
         correlation_id: UUID,
     ) -> T:
         """Execute operation with retry logic - provided by host class."""
-        raise NotImplementedError("Must be provided by implementing class")
+        raise NotImplementedError("Must be provided by implementing class")  # type: ignore[return-value]
 
     def _build_response(
         self,
@@ -103,6 +103,7 @@ class MixinConsulKV:
         input_envelope_id: UUID,
     ) -> ModelHandlerOutput[ModelConsulHandlerResponse]:
         """Build standardized response - provided by host class."""
+        raise NotImplementedError("Must be provided by implementing class")  # type: ignore[return-value]
 
     async def _kv_get(
         self,
@@ -176,7 +177,7 @@ class MixinConsulKV:
             get_func,
             correlation_id,
         )
-        index, data = cast(KVGetResult, result)
+        index, data = cast("KVGetResult", result)
 
         # Handle response - data can be None if key doesn't exist
         if data is None:
@@ -325,7 +326,7 @@ class MixinConsulKV:
             put_func,
             correlation_id,
         )
-        success = cast(bool, result)
+        success = cast("bool", result)
 
         typed_payload = ModelConsulKVPutPayload(
             success=success,

omnibase_infra/handlers/mixins/mixin_consul_service.py

@@ -91,7 +91,7 @@ class MixinConsulService:
         correlation_id: UUID,
     ) -> T:
         """Execute operation with retry logic - provided by host class."""
-        raise NotImplementedError("Must be provided by implementing class")
+        raise NotImplementedError("Must be provided by implementing class")  # type: ignore[return-value]
 
     def _build_response(
         self,
@@ -100,6 +100,7 @@ class MixinConsulService:
         input_envelope_id: UUID,
     ) -> ModelHandlerOutput[ModelConsulHandlerResponse]:
         """Build standardized response - provided by host class."""
+        raise NotImplementedError("Must be provided by implementing class")  # type: ignore[return-value]
 
     async def _register_service(
         self,

omnibase_infra/handlers/registration_storage/handler_registration_storage_postgres.py

@@ -48,6 +48,7 @@ from uuid import UUID, uuid4
 # Import asyncpg at module level to avoid redundant imports inside methods
 import asyncpg
 
+from omnibase_core.container import ModelONEXContainer
 from omnibase_core.enums.enum_node_kind import EnumNodeKind
 from omnibase_infra.enums import EnumInfraTransportType
 from omnibase_infra.errors import (
@@ -152,7 +153,10 @@ class HandlerRegistrationStoragePostgres(MixinAsyncCircuitBreaker):
         handler_type: Returns "postgresql" identifier.
 
     Example:
+        >>> from omnibase_core.container import ModelONEXContainer
+        >>> container = ModelONEXContainer(...)
         >>> handler = HandlerRegistrationStoragePostgres(
+        ...     container=container,
         ...     postgres_adapter=postgres_adapter,
         ...     circuit_breaker_config={"threshold": 5, "reset_timeout": 30.0},
         ... )
@@ -161,6 +165,7 @@ class HandlerRegistrationStoragePostgres(MixinAsyncCircuitBreaker):
 
     def __init__(
         self,
+        container: ModelONEXContainer,
         postgres_adapter: ProtocolPostgresAdapter | None = None,
         dsn: str | None = None,
         host: str = "localhost",
@@ -178,6 +183,7 @@ class HandlerRegistrationStoragePostgres(MixinAsyncCircuitBreaker):
         """Initialize HandlerRegistrationStoragePostgres.
 
         Args:
+            container: ONEX dependency injection container (required).
             postgres_adapter: Optional existing PostgreSQL adapter (ProtocolPostgresAdapter).
                 If not provided, a new asyncpg connection pool will be created.
             dsn: Optional PostgreSQL connection DSN (overrides host/port/etc).
@@ -197,6 +203,7 @@ class HandlerRegistrationStoragePostgres(MixinAsyncCircuitBreaker):
                 table on first connection. Default is False. Production deployments
                 should use database migrations instead of auto-creation.
         """
+        self._container = container
         # Normalize circuit breaker config to ModelCircuitBreakerConfig
         if isinstance(circuit_breaker_config, ModelCircuitBreakerConfig):
             cb_config = circuit_breaker_config

omnibase_infra/handlers/service_discovery/handler_service_discovery_consul.py

@@ -30,6 +30,7 @@ from uuid import NAMESPACE_DNS, UUID, uuid4, uuid5
 
 import consul
 
+from omnibase_core.container import ModelONEXContainer
 from omnibase_infra.enums import EnumInfraTransportType
 from omnibase_infra.errors import (
     InfraConnectionError,
@@ -91,7 +92,10 @@ class HandlerServiceDiscoveryConsul(MixinAsyncCircuitBreaker):
         handler_type: Returns "consul" identifier.
 
     Example:
+        >>> from unittest.mock import MagicMock
+        >>> container = MagicMock(spec=ModelONEXContainer)
         >>> handler = HandlerServiceDiscoveryConsul(
+        ...     container=container,
         ...     consul_client=consul_client,
         ...     circuit_breaker_config=ModelCircuitBreakerConfig(threshold=5),
         ... )
@@ -100,6 +104,7 @@ class HandlerServiceDiscoveryConsul(MixinAsyncCircuitBreaker):
 
     def __init__(
         self,
+        container: ModelONEXContainer,
         consul_client: ProtocolConsulClient | None = None,
         consul_host: str = "localhost",
         consul_port: int = 8500,
@@ -114,6 +119,7 @@ class HandlerServiceDiscoveryConsul(MixinAsyncCircuitBreaker):
         """Initialize HandlerServiceDiscoveryConsul.
 
         Args:
+            container: ONEX container for dependency injection and service resolution.
             consul_client: Optional existing Consul client (ProtocolConsulClient).
                 If not provided, a new python-consul client will be created.
             consul_host: Consul server hostname (default: "localhost").
@@ -129,6 +135,7 @@ class HandlerServiceDiscoveryConsul(MixinAsyncCircuitBreaker):
             max_workers: Thread pool max workers (default: 10).
             timeout_seconds: Operation timeout in seconds (default: 30.0).
         """
+        self._container = container
         # Parse circuit breaker configuration using ModelCircuitBreakerConfig
         if isinstance(circuit_breaker_config, ModelCircuitBreakerConfig):
             cb_config = circuit_breaker_config
@@ -567,15 +574,15 @@ class HandlerServiceDiscoveryConsul(MixinAsyncCircuitBreaker):
             service_infos: list[ModelServiceInfo] = []
             for svc in services:
                 # Cast nested dicts for type safety
-                svc_data = cast(dict[str, object], svc.get("Service", {}))
-                node_data = cast(dict[str, object], svc.get("Node", {}))
+                svc_data = cast("dict[str, object]", svc.get("Service", {}))
+                node_data = cast("dict[str, object]", svc.get("Node", {}))
                 svc_id_str = str(svc_data.get("ID", ""))
                 svc_name = svc_data.get("Service", "")
                 address = svc_data.get("Address", "") or node_data.get("Address", "")
                 port_raw = svc_data.get("Port", 0)
                 port = int(port_raw) if isinstance(port_raw, int | float | str) else 0
-                svc_tags = cast(list[str], svc_data.get("Tags", []))
-                svc_meta = cast(dict[str, str], svc_data.get("Meta", {}))
+                svc_tags = cast("list[str]", svc_data.get("Tags", []))
+                svc_meta = cast("dict[str, str]", svc_data.get("Meta", {}))
 
                 if svc_id_str and svc_name and address and port:
                     # Convert Consul service ID string to UUID
@@ -664,6 +671,303 @@ class HandlerServiceDiscoveryConsul(MixinAsyncCircuitBreaker):
                 context=context,
             ) from e
 
+    async def list_all_services(
+        self,
+        tag_filter: str | None = None,
+        correlation_id: UUID | None = None,
+    ) -> dict[str, list[str]]:
+        """List all registered service names with their tags.
+
+        Uses Consul catalog API (/v1/catalog/services) to retrieve all
+        service names registered in the Consul catalog, along with their tags.
+
+        Args:
+            tag_filter: Optional tag to filter services by. If provided,
+                only services with this tag will be returned.
+            correlation_id: Optional correlation ID for tracing.
+
+        Returns:
+            Dictionary mapping service names to their list of tags.
+            Example: {"my-service": ["web", "api"], "other-service": ["db"]}
+
+        Raises:
+            InfraConnectionError: If connection to Consul fails.
+            InfraTimeoutError: If operation times out.
+            InfraUnavailableError: If circuit breaker is open.
+        """
+        correlation_id = correlation_id or uuid4()
+        start_time = time.monotonic()
+
+        # Guard against use-after-shutdown
+        self._check_not_shutdown("list_all_services", correlation_id)
+
+        # Check circuit breaker
+        async with self._circuit_breaker_lock:
+            await self._check_circuit_breaker(
+                operation="list_all_services",
+                correlation_id=correlation_id,
+            )
+
+        try:
+            # Query Consul catalog for all services
+            client = self._consul_client
+            executor = await self._ensure_executor()
+            loop = asyncio.get_running_loop()
+
+            def _query_catalog() -> tuple[int, dict[str, list[str]]]:
+                # NOTE: client is duck-typed ProtocolConsulClient; mypy cannot verify
+                # catalog.services exists on Optional[Consul] union type.
+                result: tuple[int, dict[str, list[str]]] = client.catalog.services()  # type: ignore[union-attr]  # NOTE: duck-typed client
+                return result
+
+            _, services = await asyncio.wait_for(
+                loop.run_in_executor(executor, _query_catalog),
+                timeout=self._timeout_seconds,
+            )
+
+            # Reset circuit breaker on success
+            async with self._circuit_breaker_lock:
+                await self._reset_circuit_breaker()
+
+            # Apply tag filter if provided
+            if tag_filter is not None:
+                services = {
+                    name: tags for name, tags in services.items() if tag_filter in tags
+                }
+
+            duration_ms = (time.monotonic() - start_time) * 1000
+
+            logger.info(
+                "Listed all services from Consul catalog",
+                extra={
+                    "service_count": len(services),
+                    "tag_filter": tag_filter,
+                    "duration_ms": duration_ms,
+                    "correlation_id": str(correlation_id),
+                },
+            )
+
+            return services
+
+        except TimeoutError as e:
+            async with self._circuit_breaker_lock:
+                await self._record_circuit_failure(
+                    operation="list_all_services",
+                    correlation_id=correlation_id,
+                )
+            raise InfraTimeoutError(
+                f"Consul catalog query timed out after {self._timeout_seconds}s",
+                context=ModelTimeoutErrorContext(
+                    transport_type=EnumInfraTransportType.CONSUL,
+                    operation="list_all_services",
+                    target_name="consul.discovery",
+                    correlation_id=correlation_id,
+                    timeout_seconds=self._timeout_seconds,
+                ),
+            ) from e
+
+        except consul.ConsulException as e:
+            async with self._circuit_breaker_lock:
+                await self._record_circuit_failure(
+                    operation="list_all_services",
+                    correlation_id=correlation_id,
+                )
+            context = ModelInfraErrorContext(
+                transport_type=EnumInfraTransportType.CONSUL,
+                operation="list_all_services",
+                target_name="consul.discovery",
+                correlation_id=correlation_id,
+            )
+            raise InfraConnectionError(
+                "Consul catalog query failed",
+                context=context,
+            ) from e
+
+    async def get_all_service_instances(
+        self,
+        service_name: str,
+        include_unhealthy: bool = True,
+        correlation_id: UUID | None = None,
+    ) -> list[ModelServiceInfo]:
+        """Get all instances of a service, optionally including unhealthy ones.
+
+        Uses Consul health API (/v1/health/service/<name>) to retrieve all
+        instances of a service with their health status information.
+
+        Args:
+            service_name: Name of the service to query.
+            include_unhealthy: If True, return all instances regardless of health.
+                If False, only return healthy (passing) instances.
+            correlation_id: Optional correlation ID for tracing.
+
+        Returns:
+            List of ModelServiceInfo objects representing service instances,
+            including health status, health output, and last check timestamp.
+
+        Raises:
+            InfraConnectionError: If connection to Consul fails.
+            InfraTimeoutError: If operation times out.
+            InfraUnavailableError: If circuit breaker is open.
+        """
+        correlation_id = correlation_id or uuid4()
+        start_time = time.monotonic()
+
+        # Guard against use-after-shutdown
+        self._check_not_shutdown("get_all_service_instances", correlation_id)
+
+        # Check circuit breaker
+        async with self._circuit_breaker_lock:
+            await self._check_circuit_breaker(
+                operation="get_all_service_instances",
+                correlation_id=correlation_id,
+            )
+
+        try:
+            # Query Consul health API for service instances
+            client = self._consul_client
+            executor = await self._ensure_executor()
+            loop = asyncio.get_running_loop()
+
+            def _query_health() -> tuple[int, list[dict[str, object]]]:
+                # NOTE: client is duck-typed ProtocolConsulClient; mypy cannot verify
+                # health.service exists on Optional[Consul] union type.
+                # When passing=False, we get ALL instances
+                # When passing=True, we only get healthy instances
+                result: tuple[int, list[dict[str, object]]] = client.health.service(  # type: ignore[union-attr]  # NOTE: duck-typed client
+                    service_name,
+                    passing=not include_unhealthy,
+                )
+                return result
+
+            _, services = await asyncio.wait_for(
+                loop.run_in_executor(executor, _query_health),
+                timeout=self._timeout_seconds,
+            )
+
+            # Reset circuit breaker on success
+            async with self._circuit_breaker_lock:
+                await self._reset_circuit_breaker()
+
+            # Convert to ModelServiceInfo list with full health details
+            service_infos: list[ModelServiceInfo] = []
+            for svc in services:
+                # Cast nested dicts for type safety
+                svc_data = cast("dict[str, object]", svc.get("Service", {}))
+                node_data = cast("dict[str, object]", svc.get("Node", {}))
+                checks = cast("list[dict[str, object]]", svc.get("Checks", []))
+
+                svc_id_str = str(svc_data.get("ID", ""))
+                svc_name = svc_data.get("Service", "")
+                address = svc_data.get("Address", "") or node_data.get("Address", "")
+                port_raw = svc_data.get("Port", 0)
+                port = int(port_raw) if isinstance(port_raw, int | float | str) else 0
+                svc_tags = cast("list[str]", svc_data.get("Tags", []))
+                svc_meta = cast("dict[str, str]", svc_data.get("Meta", {}))
+
+                if svc_id_str and svc_name and address and port:
+                    # Convert Consul service ID string to UUID
+                    try:
+                        svc_id = UUID(svc_id_str)
+                    except ValueError:
+                        # Non-UUID service ID from Consul - generate deterministic UUID5
+                        logger.warning(
+                            "Non-UUID service ID from Consul, using deterministic UUID5 conversion",
+                            extra={
+                                "original_id": svc_id_str,
+                                "correlation_id": str(correlation_id),
+                            },
+                        )
+                        svc_id = uuid5(NAMESPACE_ONEX_SERVICE, svc_id_str)
+
+                    # Determine health status from checks
+                    health_status = EnumHealthStatus.UNKNOWN
+                    health_output: str | None = None
+                    last_check_at: datetime | None = None
+
+                    # Find the service-specific health check
+                    for check in checks:
+                        check_service_id = check.get("ServiceID", "")
+                        if check_service_id == svc_id_str or not check_service_id:
+                            status_str = str(check.get("Status", "")).lower()
+                            if status_str == "passing":
+                                health_status = EnumHealthStatus.HEALTHY
+                            elif status_str in ("critical", "warning"):
+                                health_status = EnumHealthStatus.UNHEALTHY
+                            else:
+                                health_status = EnumHealthStatus.UNKNOWN
+
+                            health_output = str(check.get("Output", "")) or None
+
+                            # Parse CreateIndex as a proxy for last check time
+                            # (Consul doesn't directly expose last check time in this API)
+                            # In production, you might use the check's CreateIndex or ModifyIndex
+                            break
+
+                    service_infos.append(
+                        ModelServiceInfo(
+                            service_id=svc_id,
+                            service_name=str(svc_name),
+                            address=str(address),
+                            port=port,
+                            tags=tuple(svc_tags or []),
+                            health_status=health_status,
+                            health_output=health_output,
+                            last_check_at=last_check_at,
+                            metadata=svc_meta or {},
+                            registered_at=datetime.now(UTC),
+                            correlation_id=correlation_id,
+                        )
+                    )
+
+            duration_ms = (time.monotonic() - start_time) * 1000
+
+            logger.info(
+                "Retrieved service instances from Consul",
+                extra={
+                    "service_name": service_name,
+                    "instance_count": len(service_infos),
+                    "include_unhealthy": include_unhealthy,
+                    "duration_ms": duration_ms,
+                    "correlation_id": str(correlation_id),
+                },
+            )
+
+            return service_infos
+
+        except TimeoutError as e:
+            async with self._circuit_breaker_lock:
+                await self._record_circuit_failure(
+                    operation="get_all_service_instances",
+                    correlation_id=correlation_id,
+                )
+            raise InfraTimeoutError(
+                f"Consul health query timed out after {self._timeout_seconds}s",
+                context=ModelTimeoutErrorContext(
+                    transport_type=EnumInfraTransportType.CONSUL,
+                    operation="get_all_service_instances",
+                    target_name="consul.discovery",
+                    correlation_id=correlation_id,
+                    timeout_seconds=self._timeout_seconds,
+                ),
+            ) from e
+
+        except consul.ConsulException as e:
+            async with self._circuit_breaker_lock:
+                await self._record_circuit_failure(
+                    operation="get_all_service_instances",
+                    correlation_id=correlation_id,
+                )
+            context = ModelInfraErrorContext(
+                transport_type=EnumInfraTransportType.CONSUL,
+                operation="get_all_service_instances",
+                target_name="consul.discovery",
+                correlation_id=correlation_id,
+            )
+            raise InfraConnectionError(
+                "Consul health query failed",
+                context=context,
+            ) from e
+
     async def health_check(
         self,
         correlation_id: UUID | None = None,

omnibase_infra/handlers/service_discovery/models/model_service_info.py

@@ -45,6 +45,8 @@ class ModelServiceInfo(BaseModel):
         health_status: Current health status of the service.
         metadata: Additional key-value metadata for the service.
         health_check_url: Optional URL for health checks (handler extension).
+        health_output: Output message from the last health check (handler extension).
+        last_check_at: Timestamp of the last health check (handler extension).
         registered_at: Timestamp when the service was registered (handler extension).
         correlation_id: Correlation ID for tracing (handler extension).
     """
@@ -86,6 +88,14 @@ class ModelServiceInfo(BaseModel):
         default=None,
         description="Optional URL for health checks",
     )
+    health_output: str | None = Field(
+        default=None,
+        description="Output message from the last health check",
+    )
+    last_check_at: datetime | None = Field(
+        default=None,
+        description="Timestamp of the last health check",
+    )
     registered_at: datetime | None = Field(
         default=None,
         description="Timestamp when the service was registered",

omnibase_infra/mixins/mixin_async_circuit_breaker.py

@@ -101,6 +101,7 @@ import logging
 import time
 from uuid import UUID, uuid4
 
+from omnibase_core.types import JsonType
 from omnibase_infra.enums import EnumCircuitState, EnumInfraTransportType
 from omnibase_infra.errors import (
     InfraUnavailableError,
@@ -580,7 +581,7 @@ class MixinAsyncCircuitBreaker:
         self._circuit_breaker_failures = 0
         self._circuit_breaker_open_until = 0.0
 
-    def _get_circuit_breaker_state(self) -> dict[str, object]:
+    def _get_circuit_breaker_state(self) -> dict[str, JsonType]:
         """Return current circuit breaker state for introspection.
 
         This method encapsulates circuit breaker internals for safe access
@@ -638,7 +639,7 @@ class MixinAsyncCircuitBreaker:
             cb_state = "closed"
             seconds_until_half_open = None
 
-        result: dict[str, object] = {
+        result: dict[str, JsonType] = {
             "initialized": cb_initialized,
             "state": cb_state,
             "failures": cb_failures,