diff --git a/src/agentic_cli/cli/workflow_controller.py b/src/agentic_cli/cli/workflow_controller.py index 5d27e2f..d446449 100644 --- a/src/agentic_cli/cli/workflow_controller.py +++ b/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() diff --git a/src/agentic_cli/workflow/base_manager.py b/src/agentic_cli/workflow/base_manager.py index fb088c1..58a46bc 100644 --- a/src/agentic_cli/workflow/base_manager.py +++ b/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: diff --git a/tests/test_workflow_controller.py b/tests/test_workflow_controller.py index 9005ed4..63bc99b 100644 --- a/tests/test_workflow_controller.py +++ b/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() diff --git a/tests/workflow/test_base_manager_init_lock.py b/tests/workflow/test_base_manager_init_lock.py new file mode 100644 index 0000000..7207905 --- /dev/null +++ b/tests/workflow/test_base_manager_init_lock.py @@ -0,0 +1,78 @@ +"""Concurrent initialize_services() must run the init body exactly once. + +The guard at the top of BaseWorkflowManager.initialize_services() was +check-then-act: a user message arriving while background init is mid-flight +(manager._ensure_initialized → initialize_services) raced the background +call and ran the whole body twice (duplicate registry refresh, duplicate +service creation). An asyncio lock serializes them. +""" + +import asyncio +from unittest.mock import AsyncMock, MagicMock + +from agentic_cli.workflow.base_manager import BaseWorkflowManager + + +class _CountingManager(BaseWorkflowManager): + """Minimal concrete manager counting _do_initialize runs.""" + + def __init__(self, settings): + super().__init__(agent_configs=[], settings=settings) + self.do_init_calls = 0 + + def _get_state_tools(self): + return [] + + @property + def backend_type(self) -> str: + return "test" + + async def _do_initialize(self) -> None: + self.do_init_calls += 1 + # Yield so a concurrent initialize_services() can interleave + await asyncio.sleep(0.02) + + async def process(self, message, user_id, session_id=None): + raise NotImplementedError + + async def reinitialize(self, model=None, preserve_sessions=True): + pass + + async def cleanup(self): + pass + + +def _settings(): + s = MagicMock() + s.app_name = "test-app" + s.google_api_key = None + s.anthropic_api_key = None + return s + + +def _manager(): + m = _CountingManager(_settings()) + m._model_registry = MagicMock(refresh=AsyncMock()) + m._ensure_managers_initialized = lambda: None + return m + + +async def test_concurrent_initialize_services_runs_once(): + m = _manager() + + await asyncio.gather( + m.initialize_services(validate=False), + m.initialize_services(validate=False), + ) + + assert m.do_init_calls == 1 + assert m.is_initialized + + +async def test_sequential_initialize_services_is_idempotent(): + m = _manager() + + await m.initialize_services(validate=False) + await m.initialize_services(validate=False) + + assert m.do_init_calls == 1