fix(workflow): controller readiness, atomic orchestrator swap, and shutdown cleanup

src/agentic_cli/cli/workflow_controller.py

@@ -40,9 +40,14 @@ class WorkflowController:
 
     Encapsulates:
     - Background initialization in ThreadPoolExecutor
-    - Readiness checking (blocking and non-blocking)
-    - Reinitialization when model/settings change
-    - Cleanup of pending init tasks
+    - Readiness checking (blocking and non-blocking); the manager is
+      published only after initialize_services() succeeds, so is_ready
+      implies a fully-initialized manager
+    - Reinitialization when model/settings change (an orchestrator swap
+      initializes the replacement first, swaps atomically, then cleans up
+      the old manager)
+    - close(): idempotent shutdown — cancels pending init and cleans up
+      the live manager; invoked on app exit via background_init()
 
     Example:
         controller = WorkflowController(
@@ -141,27 +146,40 @@ async def _background_init(self) -> None:
         Creates the workflow manager and calls initialize_services() to
         preload LLM, build graph, and set up checkpointing. This avoids
         lag on the first user message.
+
+        The manager is published to ``self._workflow`` only after
+        initialize_services() succeeds, so ``is_ready`` /
+        ``ensure_initialized()`` never report a partially-initialized or
+        failed manager as ready.
         """
         loop = asyncio.get_running_loop()
 
         def _create_workflow() -> "BaseWorkflowManager":
             return self._create_fn()
 
+        manager: "BaseWorkflowManager | None" = None
         try:
             logger.debug("background_init_starting")
 
             # Step 1: Create workflow manager (sync, in thread pool)
-            self._workflow = await loop.run_in_executor(
+            manager = await loop.run_in_executor(
                 self._init_executor, _create_workflow
             )
 
             # Step 2: Initialize services (async - builds graph, loads LLM, etc.)
-            await self._workflow.initialize_services()
+            await manager.initialize_services()
 
-            logger.info("background_init_complete", model=self._workflow.model)
+            self._workflow = manager
+            logger.info("background_init_complete", model=manager.model)
 
+        except asyncio.CancelledError:
+            if manager is not None:
+                await self._cleanup_manager(manager)
+            raise
         except Exception as e:
             self._init_error = e
+            if manager is not None:
+                await self._cleanup_manager(manager)
             logger.debug("background_init_failed", error=str(e))
 
     async def ensure_initialized(
@@ -245,13 +263,22 @@ async def reinitialize(self, model: str | None = None) -> None:
                 old_model=self._workflow.model,
                 new_model=model,
             )
-            self._workflow = create_workflow_manager_from_settings(
+            # Initialize the replacement fully before swapping so a failed
+            # init leaves the working manager in place; clean up whichever
+            # manager ends up unused.
+            new_workflow = create_workflow_manager_from_settings(
                 agent_configs=self._agent_configs,
                 settings=self._settings,
                 app_name=self._app_name,
                 model=model,
             )
-            await self._workflow.initialize_services()
+            try:
+                await new_workflow.initialize_services()
+            except Exception:
+                await self._cleanup_manager(new_workflow)
+                raise
+            old_workflow, self._workflow = self._workflow, new_workflow
+            await self._cleanup_manager(old_workflow)
         else:
             await self._workflow.reinitialize(model=model, preserve_sessions=True)
 
@@ -265,6 +292,25 @@ async def cancel_init(self) -> None:
                 pass
         self._init_executor.shutdown(wait=False)
 
+    @staticmethod
+    async def _cleanup_manager(manager: "BaseWorkflowManager") -> None:
+        """Best-effort manager cleanup; a failing cleanup is logged, not raised."""
+        try:
+            await manager.cleanup()
+        except Exception as e:
+            logger.warning("workflow_manager_cleanup_failed", error=str(e))
+
+    async def close(self) -> None:
+        """Release the controller: cancel pending init, clean up the manager.
+
+        Idempotent — safe to call multiple times. Invoked from application
+        shutdown (the background_init context manager exit).
+        """
+        await self.cancel_init()
+        manager, self._workflow = self._workflow, None
+        if manager is not None:
+            await self._cleanup_manager(manager)
+
     def update_status_bar(self, ui: "ThinkingPromptSession") -> None:
         """Update UI status bar with current workflow status.
 
@@ -315,12 +361,12 @@ async def wait_and_update() -> None:
         try:
             yield
         finally:
-            # Cancel init task if still running
-            await self.cancel_init()
-            # Also cancel status update task if still running
+            # Cancel status update task if still running
             if not update_task.done():
                 update_task.cancel()
                 try:
                     await update_task
                 except asyncio.CancelledError:
                     pass
+            # Cancel pending init and clean up the manager (app shutdown)
+            await self.close()

src/agentic_cli/workflow/base_manager.py

@@ -13,6 +13,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import contextlib
 from abc import ABC, abstractmethod
 from typing import Any, AsyncGenerator, Awaitable, Callable, Iterator, TYPE_CHECKING
@@ -90,6 +91,10 @@ def __init__(
         self._settings = settings or get_settings()
         self._app_name = app_name or self._settings.app_name
         self._initialized = False
+        # Serializes initialize_services(): background init and a first user
+        # message (process → _ensure_initialized) may call it concurrently,
+        # and the _initialized guard alone is check-then-act.
+        self._init_lock = asyncio.Lock()
         self._on_event = on_event
 
         # Model resolution (lazy)
@@ -578,39 +583,42 @@ def model(self) -> str:
 
     async def initialize_services(self, validate: bool = True) -> None:
         """Initialize backend services asynchronously.
+
+        Concurrency-safe and idempotent: concurrent callers serialize on the
+        init lock; late arrivals see ``_initialized`` and return.
+
         Args:
             validate: If True, validate settings before initialization.
         Raises:
             SettingsValidationError: If settings validation fails.
         """
-        if self._initialized:
-            return
+        async with self._init_lock:
+            if self._initialized:
+                return
 
-        from agentic_cli.config import validate_settings
+            from agentic_cli.config import validate_settings
 
-        if validate:
-            validate_settings(self._settings)
+            if validate:
+                validate_settings(self._settings)
 
-        self._settings.export_api_keys_to_env()
+            self._settings.export_api_keys_to_env()
 
-        # Refresh model registry from APIs
-        await self._model_registry.refresh(
-            google_api_key=self._settings.google_api_key,
-            anthropic_api_key=self._settings.anthropic_api_key,
-        )
-        self._settings.set_model_registry(self._model_registry)
-
-        # Create services BEFORE backend init so _build_tools() can
-        # produce factory-bound tools during agent/graph creation.
-        # Offloaded to a worker thread because constructors here may
-        # load heavy dependencies (e.g. the sentence-transformers model
-        # inside EmbeddingService) that would otherwise block the event
-        # loop — which keeps the prompt unresponsive at startup.
-        import asyncio as _asyncio
-
-        await _asyncio.to_thread(self._ensure_managers_initialized)
-        await self._do_initialize()
-        self._initialized = True
+            # Refresh model registry from APIs
+            await self._model_registry.refresh(
+                google_api_key=self._settings.google_api_key,
+                anthropic_api_key=self._settings.anthropic_api_key,
+            )
+            self._settings.set_model_registry(self._model_registry)
+
+            # Create services BEFORE backend init so _build_tools() can
+            # produce factory-bound tools during agent/graph creation.
+            # Offloaded to a worker thread because constructors here may
+            # load heavy dependencies (e.g. the sentence-transformers model
+            # inside EmbeddingService) that would otherwise block the event
+            # loop — which keeps the prompt unresponsive at startup.
+            await asyncio.to_thread(self._ensure_managers_initialized)
+            await self._do_initialize()
+            self._initialized = True
 
     @abstractmethod
     async def _do_initialize(self) -> None:

tests/test_workflow_controller.py

@@ -313,3 +313,149 @@ async def test_reinitialize_migrates_stale_langgraph_to_adk(self):
         old_workflow.reinitialize.assert_not_called()
         new_workflow.initialize_services.assert_awaited_once()
         assert controller._workflow is new_workflow
+
+
+# --- Lifecycle: readiness, atomic swap, close() ---
+
+
+def _make_lifecycle_controller(orchestrator=OrchestratorType.ADK):
+    configs = [AgentConfig(name="test", prompt="Test")]
+    return WorkflowController(configs, _make_settings(orchestrator=orchestrator))
+
+
+def _blocked_init_workflow():
+    """Fake manager whose initialize_services blocks until released."""
+    import asyncio
+
+    wf = _FakeADKWorkflow()
+    started, release = asyncio.Event(), asyncio.Event()
+
+    async def _slow_init():
+        started.set()
+        await release.wait()
+
+    wf.initialize_services = _slow_init
+    return wf, started, release
+
+
+class TestControllerLifecycle:
+    """is_ready / ensure_initialized must reflect completed service init."""
+
+    async def test_not_ready_until_services_initialized(self):
+        import asyncio
+
+        controller = _make_lifecycle_controller()
+        wf, started, release = _blocked_init_workflow()
+        controller._create_fn = lambda: wf
+
+        await controller.start_background_init()
+        await asyncio.wait_for(started.wait(), timeout=5)
+
+        assert controller.is_ready is False
+        with pytest.raises(RuntimeError):
+            controller.workflow
+
+        release.set()
+        await controller._init_task
+        assert controller.is_ready is True
+        assert controller.workflow is wf
+
+    async def test_failed_service_init_is_not_ready_and_cleans_up(self):
+        controller = _make_lifecycle_controller()
+        wf = _FakeADKWorkflow()
+        wf.initialize_services = AsyncMock(side_effect=RuntimeError("boom"))
+        controller._create_fn = lambda: wf
+
+        await controller.start_background_init()
+        await controller._init_task
+
+        assert controller.is_ready is False
+        assert isinstance(controller.init_error, RuntimeError)
+        assert await controller.ensure_initialized() is False
+        wf.cleanup.assert_awaited_once()
+
+    async def test_ensure_initialized_waits_for_inflight_services(self):
+        import asyncio
+
+        controller = _make_lifecycle_controller()
+        wf, started, release = _blocked_init_workflow()
+        controller._create_fn = lambda: wf
+
+        await controller.start_background_init()
+        await asyncio.wait_for(started.wait(), timeout=5)
+
+        ensure_task = asyncio.create_task(controller.ensure_initialized())
+        await asyncio.sleep(0.05)
+        assert not ensure_task.done()
+
+        release.set()
+        assert await ensure_task is True
+
+
+class TestOrchestratorSwapLifecycle:
+    """Swap must initialize the replacement first, then replace atomically."""
+
+    def _controller_needing_swap(self):
+        # Live ADK manager while settings demand LangGraph → swap required
+        controller = _make_lifecycle_controller(
+            orchestrator=OrchestratorType.LANGGRAPH
+        )
+        old = _FakeADKWorkflow()
+        controller._workflow = old
+        return controller, old
+
+    async def test_swap_failure_keeps_old_manager(self):
+        controller, old = self._controller_needing_swap()
+        new = _FakeLangGraphWorkflow()
+        new.initialize_services = AsyncMock(side_effect=RuntimeError("init failed"))
+
+        with patch(
+            "agentic_cli.cli.workflow_controller.create_workflow_manager_from_settings",
+            return_value=new,
+        ):
+            with pytest.raises(RuntimeError, match="init failed"):
+                await controller.reinitialize()
+
+        assert controller._workflow is old
+        old.cleanup.assert_not_awaited()
+        new.cleanup.assert_awaited_once()
+
+    async def test_swap_success_replaces_then_cleans_old(self):
+        controller, old = self._controller_needing_swap()
+        new = _FakeLangGraphWorkflow()
+
+        with patch(
+            "agentic_cli.cli.workflow_controller.create_workflow_manager_from_settings",
+            return_value=new,
+        ):
+            await controller.reinitialize()
+
+        assert controller._workflow is new
+        old.cleanup.assert_awaited_once()
+
+
+class TestControllerClose:
+    """close() releases the manager and is safe to call repeatedly."""
+
+    async def test_close_cleans_manager_and_is_idempotent(self):
+        controller = _make_lifecycle_controller()
+        wf = _FakeADKWorkflow()
+        controller._workflow = wf
+
+        await controller.close()
+        wf.cleanup.assert_awaited_once()
+        assert controller.is_ready is False
+
+        await controller.close()
+        wf.cleanup.assert_awaited_once()  # still once — idempotent
+
+    async def test_background_init_cm_closes_manager_on_exit(self):
+        controller = _make_lifecycle_controller()
+        wf = _FakeADKWorkflow()
+        controller._create_fn = lambda: wf
+        ui = MagicMock()
+
+        async with controller.background_init(ui):
+            assert await controller.ensure_initialized() is True
+
+        wf.cleanup.assert_awaited_once()