fix: initial merge conflict resolutions

Author: ikemilian-lewis

src/orb/api/routers/providers.py

@@ -82,17 +82,16 @@ def _get_schema_for_provider_type(provider_type: str) -> list[dict[str, Any]]:
     Raises ``KeyError`` when the provider type is not registered.
     """
     from orb.providers.registry.provider_registry import get_provider_registry
-    from orb.providers.registry.types import ProviderRegistration
 
     registry = get_provider_registry()
-    reg = registry._get_type_registration(provider_type)  # raises ValueError if missing
-    if not isinstance(reg, ProviderRegistration) or reg.strategy_class is None:
+    strategy_class = registry.get_strategy_class(provider_type)
+    if strategy_class is None:
         return []
 
     try:
         # Call the classmethod directly — no instance needed.
         # get_ui_column_schema only constructs UIColumnDescriptor objects; no I/O.
-        schema = reg.strategy_class.get_ui_column_schema()
+        schema = strategy_class.get_ui_column_schema()
         return [col.to_dict() for col in schema]
     except Exception as exc:
         logger.warning(

src/orb/application/commands/request_creation_handlers.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import asyncio
 from typing import Any
 
 from orb.application.base.handlers import BaseCommandHandler
@@ -11,6 +12,7 @@
     CreateReturnRequestCommand,
 )
 from orb.application.ports.query_bus_port import QueryBusPort
+from orb.application.services.provider_validation_service import ProviderValidationService
 from orb.application.services.provisioning_orchestration_service import (
     ProvisioningOrchestrationService,
 )
@@ -25,7 +27,9 @@
     LoggingPort,
     ProviderSelectionPort,
 )
+from orb.domain.request.aggregate import Request
 from orb.domain.request.request_identifiers import RequestId
+from orb.domain.request.request_types import RequestStatus
 from orb.domain.request.value_objects import RequestType
 from orb.domain.template.factory import TemplateFactoryPort
 from orb.domain.template.template_aggregate import Template
@@ -50,6 +54,7 @@ def __init__(
         query_bus: QueryBusPort,  # QueryBus is required for template lookup
         provider_selection_port: ProviderSelectionPort,
         provisioning_service: ProvisioningOrchestrationService,
+        provider_validation_service: ProviderValidationService,
     ) -> None:
         """Initialize the instance."""
         super().__init__(logger, event_publisher, error_handler)
@@ -59,7 +64,6 @@ def __init__(
         self._provider_selection_port = provider_selection_port
 
         # Initialize services
-        from orb.application.services.provider_validation_service import ProviderValidationService
         from orb.application.services.request_creation_service import RequestCreationService
         from orb.application.services.request_status_management_service import (
             RequestStatusManagementService,
@@ -68,9 +72,7 @@ def __init__(
         self._request_creation_service = RequestCreationService(logger)
         self._provisioning_service = provisioning_service
         self._status_service = RequestStatusManagementService(uow_factory, logger)
-        self._provider_validation_service = ProviderValidationService(
-            container, logger, provider_selection_port
-        )
+        self._provider_validation_service = provider_validation_service
 
     async def validate_command(self, command: CreateRequestCommand) -> None:
         """Validate create request command."""
@@ -148,8 +150,6 @@ async def _load_template(self, template_id: str) -> Template:
 
     def _handle_dry_run(self, request: Any) -> Any:
         """Handle dry-run request."""
-        from orb.domain.request.value_objects import RequestStatus
-
         self.logger.info(
             "Skipping actual provisioning for request %s (dry-run mode)",
             request.request_id,
@@ -167,8 +167,6 @@ async def _persist_and_publish(self, request: Any) -> None:
         error is logged at ERROR level but NOT re-raised because the request was
         already successfully persisted.
         """
-        import asyncio
-
         with self.uow_factory.create_unit_of_work() as uow:
             events = uow.requests.save(request)
 
@@ -268,8 +266,6 @@ async def execute_command(self, command: CreateReturnRequestCommand) -> None:
         self.logger.info("Creating return request for machines: %s", command.machine_ids)
 
         try:
-            from orb.domain.request.aggregate import Request
-
             domain_config = self._container.get(DomainConfigurationService)
             prefix = domain_config.get_return_request_prefix()
             force_return = command.force_return or False
@@ -297,12 +293,17 @@ async def execute_command(self, command: CreateReturnRequestCommand) -> None:
             created_requests: list[str] = []
             pending_deprovision: list[tuple[list[str], Any, str]] = []
 
-            for (provider_type, provider_name), machine_ids in provider_groups.items():
+            for (
+                provider_type,
+                provider_name,
+                provider_api,
+            ), machine_ids in provider_groups.items():
                 return_request_id = str(RequestId.generate(RequestType.RETURN, prefix=prefix))
                 request = Request.create_return_request(
                     machine_ids=machine_ids,
                     provider_type=provider_type,
                     provider_name=provider_name,
+                    provider_api=provider_api,
                     metadata=command.metadata or {},
                     request_id=return_request_id,
                 )
@@ -330,9 +331,24 @@ async def execute_command(self, command: CreateReturnRequestCommand) -> None:
             command.processed_machines = valid_machines
             command.skipped_machines = skipped_machines
 
-            # Await deprovisioning sequentially — one per provider group.
+            # Synchronous deprovisioning: await provider termination before
+            # returning. Integration tests + the CLI rely on the handler
+            # reaching a terminal status by the time it returns; UI-side
+            # responsiveness is addressed at the API layer (the HTTP
+            # caller can return early via a 202 + background task at the
+            # router level, not here in the command handler).
             for machine_ids, request, provider_name in pending_deprovision:
-                await self._execute_deprovisioning_for_request(machine_ids, request, provider_name)
+                try:
+                    await self._execute_deprovisioning_for_request(
+                        machine_ids, request, provider_name
+                    )
+                except Exception as deprov_err:
+                    self.logger.error(
+                        "Deprovisioning failed for request %s: %s",
+                        request.request_id,
+                        deprov_err,
+                        exc_info=True,
+                    )
 
         except Exception as e:
             self.logger.error(
@@ -436,7 +452,6 @@ async def _update_request_to_in_progress(self, request: Any) -> None:
         try:
             from orb.application.dto.commands import UpdateRequestStatusCommand
             from orb.application.ports.command_bus_port import CommandBusPort
-            from orb.domain.request.request_types import RequestStatus
 
             update_command = UpdateRequestStatusCommand(
                 request_id=str(request.request_id),
@@ -481,7 +496,6 @@ async def _execute_deprovisioning_for_request(
                 if skipped_ids:
                     from orb.application.dto.commands import UpdateRequestStatusCommand
                     from orb.application.ports.command_bus_port import CommandBusPort
-                    from orb.domain.request.request_types import RequestStatus
 
                     skipped_str = ", ".join(skipped_ids)
                     update_command = UpdateRequestStatusCommand(
@@ -502,11 +516,8 @@ async def _execute_deprovisioning_for_request(
                         skipped_str,
                     )
                 else:
-                    await self._update_request_to_in_progress_with_message(
-                        request,
-                        "Termination initiated, waiting for provider confirmation",
-                    )
-                    self.logger.info("Termination initiated for request %s", request.request_id)
+                    await self._update_request_to_terminating(request)
+                    self.logger.info("Termination accepted for request %s", request.request_id)
             else:
                 await self._update_request_to_failed(request, provisioning_result.get("errors", []))
 
@@ -533,7 +544,9 @@ def _update_machines_to_pending(self, machine_ids: list[str]) -> None:
                     )
                     uow.machines.save(updated_machine)
 
-    def _persist_return_follow_up_context(self, request: Any, provisioning_result: dict[str, Any]) -> None:
+    def _persist_return_follow_up_context(
+        self, request: Any, provisioning_result: dict[str, Any]
+    ) -> None:
         """Persist durable provider follow-up metadata needed after this process exits."""
         provider_data = provisioning_result.get("provider_data")
         if not isinstance(provider_data, dict) or not provider_data:
@@ -546,69 +559,87 @@ def _persist_return_follow_up_context(self, request: Any, provisioning_result: d
             for event in events or []:
                 self.event_publisher.publish(event)  # type: ignore[union-attr]
 
-    async def _update_request_to_failed(self, request: Any, errors: list[str]) -> None:
-        """Update request status to failed."""
-        try:
-            from orb.application.dto.commands import UpdateRequestStatusCommand
-            from orb.application.ports.command_bus_port import CommandBusPort
-            from orb.domain.request.request_types import RequestStatus
-
-            error_message = "; ".join(errors) if errors else "Deprovisioning failed"
+    async def _update_request_status(
+        self, request: Request, status: RequestStatus, message: str
+    ) -> None:
+        """Update request status via the command bus with a UoW fallback for FAILED.
 
-            update_command = UpdateRequestStatusCommand(
-                request_id=str(request.request_id),
-                status=RequestStatus.FAILED,
-                message=f"Return request failed: {error_message}",
-            )
+        On command-bus failure the error is logged.  For FAILED status a direct
+        UoW write is attempted so the request always reaches a terminal state and
+        callers do not poll forever.  For all other statuses the caller decides
+        how to handle the failure (e.g. _update_request_to_terminating re-raises
+        so the deprovisioning loop can apply its own fallback).
+        """
+        from orb.application.dto.commands import UpdateRequestStatusCommand
+        from orb.application.ports.command_bus_port import CommandBusPort
 
+        update_command = UpdateRequestStatusCommand(
+            request_id=str(request.request_id),
+            status=status,
+            message=message,
+        )
+        try:
             command_bus = self._container.get(CommandBusPort)
             await command_bus.execute(update_command)
-
-            self.logger.info("Updated request %s status to failed", request.request_id)
-
+            self.logger.info("Updated request %s status to %s", request.request_id, status)
         except Exception as update_error:
             self.logger.error(
-                "Failed to update request status: %s",
+                "Failed to update request %s status to %s: %s",
+                request.request_id,
+                status,
                 update_error,
                 exc_info=True,
             )
-            # Force-write terminal status directly to prevent permanent stuck state
-            try:
-                from orb.domain.request.request_types import RequestStatus
+            if status == RequestStatus.FAILED:
+                # Force-write terminal status directly to prevent permanent stuck state
+                try:
+                    with self.uow_factory.create_unit_of_work() as uow:
+                        stuck_request = uow.requests.get_by_id(request.request_id)
+                        if stuck_request:
+                            stuck_request = stuck_request.update_status(
+                                RequestStatus.FAILED,
+                                "System error: failed to update status after double failure",
+                                force=True,
+                            )
+                            uow.requests.save(stuck_request)
+                except Exception as final_error:
+                    self.logger.critical(
+                        "CRITICAL: Failed to mark request %s as failed after double failure. "
+                        "Request is stuck in IN_PROGRESS. Manual intervention required. Error: %s",
+                        request.request_id,
+                        final_error,
+                    )
+            else:
+                raise
 
-                with self.uow_factory.create_unit_of_work() as uow:
-                    stuck_request = uow.requests.get_by_id(request.request_id)
-                    if stuck_request:
-                        stuck_request = stuck_request.update_status(
-                            RequestStatus.FAILED,
-                            "System error: failed to update status after double failure",
-                            force=True,
-                        )
-                        uow.requests.save(stuck_request)
-            except Exception as final_error:
-                self.logger.critical(
-                    "CRITICAL: Failed to mark request %s as failed after double failure. "
-                    "Request is stuck in IN_PROGRESS. Manual intervention required. Error: %s",
-                    request.request_id,
-                    final_error,
-                )
-                # Nothing more we can do
+    async def _update_request_to_failed(self, request: Any, errors: list[str]) -> None:
+        """Update request status to failed."""
+        error_message = "; ".join(errors) if errors else "Deprovisioning failed"
+        await self._update_request_status(
+            request, RequestStatus.FAILED, f"Return request failed: {error_message}"
+        )
+
+    async def _update_request_to_terminating(self, request: Any) -> None:
+        """Set return request to IN_PROGRESS after termination is accepted by provider.
+
+        Termination is asynchronous — the provider accepted the request but instances
+        are still ``shutting-down``.  Using IN_PROGRESS here lets the background sync
+        poll AWS and transition to COMPLETED only when all instances reach ``terminated``.
+        """
+        await self._update_request_status(
+            request,
+            RequestStatus.IN_PROGRESS,
+            "Termination accepted: waiting for instances to reach terminated state",
+        )
 
     async def _update_request_to_in_progress_with_message(self, request: Any, message: str) -> None:
         """Persist an in-progress return message after submit without claiming completion."""
         try:
-            from orb.application.dto.commands import UpdateRequestStatusCommand
-            from orb.application.ports.command_bus_port import CommandBusPort
-            from orb.domain.request.request_types import RequestStatus
-
-            update_command = UpdateRequestStatusCommand(
-                request_id=str(request.request_id),
-                status=RequestStatus.IN_PROGRESS,
-                message=message,
+            await self._update_request_status(
+                request,
+                RequestStatus.IN_PROGRESS,
+                message,
             )
-            command_bus = self._container.get(CommandBusPort)
-            await command_bus.execute(update_command)
-            self.logger.info("Updated request %s status to in_progress", request.request_id)
         except Exception as update_error:
             self.logger.error(
                 "Failed to update request status to in_progress: %s",

src/orb/application/services/deprovisioning_orchestrator.py

@@ -195,6 +195,7 @@ async def _process_resource_group(
                     "resource_mapping": {iid: (resource_id, 1) for iid in instance_ids},
                     "request_id": origin_request_id,
                     "request_metadata": self._build_request_metadata(request),
+                    "request": request,
                 },
                 context={
                     "correlation_id": str(request.request_id),