flock-core 0.5.11__py3-none-any.whl → 0.5.21__py3-none-any.whl

flock/orchestrator/lifecycle_manager.py

@@ -0,0 +1,226 @@
+"""Lifecycle management for background tasks and cleanup.
+
+Phase 5A: Extracted from orchestrator.py to isolate background task coordination.
+
+This module handles background tasks for batch timeouts and correlation cleanup,
+reducing orchestrator complexity and centralizing async task management.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from asyncio import Task
+from typing import TYPE_CHECKING, Any
+
+
+if TYPE_CHECKING:
+    from flock.orchestrator.batch_accumulator import BatchEngine
+    from flock.orchestrator.correlation_engine import CorrelationEngine
+
+
+class LifecycleManager:
+    """Manages background tasks for batch and correlation lifecycle.
+
+    This module centralizes all background task management for:
+    - Correlation group cleanup (time-based expiry)
+    - Batch timeout checking (timeout-based flushing)
+
+    Phase 5A: Extracted to reduce orchestrator complexity and improve testability.
+    """
+
+    def __init__(
+        self,
+        *,
+        correlation_engine: CorrelationEngine,
+        batch_engine: BatchEngine,
+        cleanup_interval: float = 0.1,
+    ):
+        """Initialize LifecycleManager with engines and intervals.
+
+        Args:
+            correlation_engine: Engine managing correlation groups
+            batch_engine: Engine managing batch accumulation
+            cleanup_interval: How often to check for expiry (seconds, default: 0.1)
+        """
+        self._correlation_engine = correlation_engine
+        self._batch_engine = batch_engine
+        self._cleanup_interval = cleanup_interval
+
+        # Background tasks
+        self._correlation_cleanup_task: Task[Any] | None = None
+        self._batch_timeout_task: Task[Any] | None = None
+
+        # Callback for batch timeout flushing (set by orchestrator)
+        self._batch_timeout_callback: Any | None = None
+
+        self._logger = logging.getLogger(__name__)
+
+    async def start_correlation_cleanup(self) -> None:
+        """Start background correlation cleanup loop if not already running.
+
+        This ensures expired correlation groups are periodically discarded.
+        Called when there are pending correlations during run_until_idle.
+        """
+        if (
+            self._correlation_cleanup_task is None
+            or self._correlation_cleanup_task.done()
+        ):
+            self._correlation_cleanup_task = asyncio.create_task(
+                self._correlation_cleanup_loop()
+            )
+
+    def set_batch_timeout_callback(self, callback: Any) -> None:
+        """Set the callback to invoke when batches timeout.
+
+        Args:
+            callback: Async function to call when timeout checking. Should handle
+                flushing expired batches and scheduling agent tasks.
+        """
+        self._batch_timeout_callback = callback
+
+    async def start_batch_timeout_checker(self) -> None:
+        """Start background batch timeout checker loop if not already running.
+
+        This ensures timeout-expired batches are periodically flushed.
+        Called when there are pending batches during run_until_idle.
+        """
+        if self._batch_timeout_task is None or self._batch_timeout_task.done():
+            self._batch_timeout_task = asyncio.create_task(
+                self._batch_timeout_checker_loop()
+            )
+
+    async def shutdown(self) -> None:
+        """Cancel and cleanup all background tasks.
+
+        Called during orchestrator shutdown to ensure clean resource cleanup.
+        """
+        # Cancel correlation cleanup task if running
+        if self._correlation_cleanup_task and not self._correlation_cleanup_task.done():
+            self._correlation_cleanup_task.cancel()
+            try:
+                await self._correlation_cleanup_task
+            except asyncio.CancelledError:
+                pass
+
+        # Cancel batch timeout checker if running
+        if self._batch_timeout_task and not self._batch_timeout_task.done():
+            self._batch_timeout_task.cancel()
+            try:
+                await self._batch_timeout_task
+            except asyncio.CancelledError:
+                pass
+
+    # Background Loops ─────────────────────────────────────────────────────
+
+    async def _correlation_cleanup_loop(self) -> None:
+        """Background task that periodically cleans up expired correlation groups.
+
+        Runs continuously until all correlation groups are cleared or orchestrator shuts down.
+        Checks every 100ms for time-based expired correlations and discards them.
+        """
+        try:
+            while True:
+                await asyncio.sleep(self._cleanup_interval)
+                self._cleanup_expired_correlations()
+
+                # Stop if no correlation groups remain
+                if not self._correlation_engine.correlation_groups:
+                    self._correlation_cleanup_task = None
+                    break
+        except asyncio.CancelledError:
+            # Clean shutdown
+            self._correlation_cleanup_task = None
+            raise
+
+    def _cleanup_expired_correlations(self) -> None:
+        """Clean up all expired correlation groups across all subscriptions.
+
+        Called periodically by background task to enforce time-based correlation windows.
+        Discards incomplete correlations that have exceeded their time window.
+        """
+        # Get all active subscription keys
+        for agent_name, subscription_index in list(
+            self._correlation_engine.correlation_groups.keys()
+        ):
+            self._correlation_engine.cleanup_expired(agent_name, subscription_index)
+
+    async def _batch_timeout_checker_loop(self) -> None:
+        """Background task that periodically checks for batch timeouts.
+
+        Runs continuously until all batches are cleared or orchestrator shuts down.
+        Checks every 100ms for expired batches and flushes them via callback.
+        """
+        try:
+            while True:
+                await asyncio.sleep(self._cleanup_interval)
+
+                # Call the timeout callback to check and flush expired batches
+                if self._batch_timeout_callback:
+                    await self._batch_timeout_callback()
+
+                # Stop if no batches remain
+                if not self._batch_engine.batches:
+                    self._batch_timeout_task = None
+                    break
+        except asyncio.CancelledError:
+            # Clean shutdown
+            self._batch_timeout_task = None
+            raise
+
+    # Helper Methods ───────────────────────────────────────────────────────
+
+    async def check_batch_timeouts(self, orchestrator_callback: Any) -> None:
+        """Check all batches for timeout expiry and invoke callback for expired batches.
+
+        This method is called periodically by the background timeout checker
+        or manually (in tests) to enforce timeout-based batching.
+
+        Args:
+            orchestrator_callback: Async function to call for each expired batch.
+                Signature: async def callback(agent_name: str, subscription_index: int,
+                                              artifacts: list[Artifact]) -> None
+        """
+        expired_batches = self._batch_engine.check_timeouts()
+
+        for agent_name, subscription_index in expired_batches:
+            # Flush the expired batch
+            artifacts = self._batch_engine.flush_batch(agent_name, subscription_index)
+
+            if artifacts is not None:
+                # Invoke orchestrator callback to schedule task
+                await orchestrator_callback(agent_name, subscription_index, artifacts)
+
+    async def flush_all_batches(self, orchestrator_callback: Any) -> None:
+        """Flush all partial batches (for shutdown - ensures zero data loss).
+
+        Args:
+            orchestrator_callback: Async function to call for each flushed batch.
+                Signature: async def callback(agent_name: str, subscription_index: int,
+                                              artifacts: list[Artifact]) -> None
+        """
+        all_batches = self._batch_engine.flush_all()
+
+        for agent_name, subscription_index, artifacts in all_batches:
+            # Invoke orchestrator callback to schedule task
+            await orchestrator_callback(agent_name, subscription_index, artifacts)
+
+    # Properties ───────────────────────────────────────────────────────────
+
+    @property
+    def has_pending_correlations(self) -> bool:
+        """Check if there are any pending correlation groups."""
+        return any(
+            groups and any(group.waiting_artifacts for group in groups.values())
+            for groups in self._correlation_engine.correlation_groups.values()
+        )
+
+    @property
+    def has_pending_batches(self) -> bool:
+        """Check if there are any pending batches."""
+        return any(
+            accumulator.artifacts for accumulator in self._batch_engine.batches.values()
+        )
+
+
+__all__ = ["LifecycleManager"]

flock/orchestrator/mcp_manager.py

@@ -0,0 +1,202 @@
+"""MCP (Model Context Protocol) server management for orchestrator.
+
+This module handles MCP server registration and client manager lifecycle.
+Implements lazy connection establishment pattern (AD005).
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+
+if TYPE_CHECKING:
+    from flock.mcp import (
+        FlockMCPClientManager,
+        FlockMCPConfiguration,
+        ServerParameters,
+    )
+
+
+class MCPManager:
+    """Manages MCP server registration and client connections.
+
+    Architecture Decision: AD001 - Two-Level Architecture
+    MCP servers are registered at orchestrator level and assigned to agents.
+
+    Architecture Decision: AD005 - Lazy Connection Establishment
+    Connections are established only when get_mcp_manager() is first called.
+
+    Attributes:
+        _configs: Dict mapping server names to their configurations
+        _client_manager: Lazy-initialized MCP client manager instance
+    """
+
+    def __init__(self) -> None:
+        """Initialize the MCP manager with empty configuration."""
+        self._configs: dict[str, FlockMCPConfiguration] = {}
+        self._client_manager: FlockMCPClientManager | None = None
+
+    def add_mcp(
+        self,
+        name: str,
+        connection_params: ServerParameters,
+        *,
+        enable_tools_feature: bool = True,
+        enable_prompts_feature: bool = True,
+        enable_sampling_feature: bool = True,
+        enable_roots_feature: bool = True,
+        mount_points: list[str] | None = None,
+        tool_whitelist: list[str] | None = None,
+        read_timeout_seconds: float = 300,
+        max_retries: int = 3,
+        **kwargs,
+    ) -> None:
+        """Register an MCP server configuration.
+
+        Args:
+            name: Unique identifier for this MCP server
+            connection_params: Server connection parameters
+            enable_tools_feature: Enable tool execution
+            enable_prompts_feature: Enable prompt templates
+            enable_sampling_feature: Enable LLM sampling requests
+            enable_roots_feature: Enable filesystem roots
+            mount_points: Optional list of filesystem mount points
+            tool_whitelist: Optional list of tool names to allow
+            read_timeout_seconds: Timeout for server communications
+            max_retries: Connection retry attempts
+
+        Raises:
+            ValueError: If server name already registered
+        """
+        if name in self._configs:
+            raise ValueError(f"MCP server '{name}' is already registered.")
+
+        # Import configuration types
+        from flock.mcp import (
+            FlockMCPConfiguration,
+            FlockMCPConnectionConfiguration,
+            FlockMCPFeatureConfiguration,
+        )
+
+        # Detect transport type
+        from flock.mcp.types import (
+            SseServerParameters,
+            StdioServerParameters,
+            StreamableHttpServerParameters,
+            WebsocketServerParameters,
+        )
+
+        if isinstance(connection_params, StdioServerParameters):
+            transport_type = "stdio"
+        elif isinstance(connection_params, WebsocketServerParameters):
+            transport_type = "websockets"
+        elif isinstance(connection_params, SseServerParameters):
+            transport_type = "sse"
+        elif isinstance(connection_params, StreamableHttpServerParameters):
+            transport_type = "streamable_http"
+        else:
+            transport_type = "custom"
+
+        # Process mount points (convert paths to URIs)
+        mcp_roots = None
+        if mount_points:
+            from pathlib import Path as PathLib
+
+            from flock.mcp.types import MCPRoot
+
+            mcp_roots = []
+            for path in mount_points:
+                # Normalize the path
+                if path.startswith("file://"):
+                    # Already a file URI
+                    uri = path
+                    # Extract path from URI for name
+                    path_str = path.replace("file://", "")
+                # the test:// path-prefix is used by testing servers such as the mcp-everything server.
+                elif path.startswith("test://"):
+                    # Already a test URI
+                    uri = path
+                    # Extract path from URI for name
+                    path_str = path.replace("test://", "")
+                else:
+                    # Convert to absolute path and create URI
+                    abs_path = PathLib(path).resolve()
+                    uri = f"file://{abs_path}"
+                    path_str = str(abs_path)
+
+                # Extract a meaningful name (last component of path)
+                name_component = (
+                    PathLib(path_str).name
+                    or path_str.rstrip("/").split("/")[-1]
+                    or "root"
+                )
+                mcp_roots.append(MCPRoot(uri=uri, name=name_component))
+
+        # Build configuration
+        connection_config = FlockMCPConnectionConfiguration(
+            max_retries=max_retries,
+            connection_parameters=connection_params,
+            transport_type=transport_type,
+            read_timeout_seconds=read_timeout_seconds,
+            mount_points=mcp_roots,
+        )
+
+        feature_config = FlockMCPFeatureConfiguration(
+            tools_enabled=enable_tools_feature,
+            prompts_enabled=enable_prompts_feature,
+            sampling_enabled=enable_sampling_feature,
+            roots_enabled=enable_roots_feature,
+            tool_whitelist=tool_whitelist,
+        )
+
+        mcp_config = FlockMCPConfiguration(
+            name=name,
+            connection_config=connection_config,
+            feature_config=feature_config,
+        )
+
+        self._configs[name] = mcp_config
+
+    def get_mcp_manager(self) -> FlockMCPClientManager:
+        """Get or create the MCP client manager.
+
+        Architecture Decision: AD005 - Lazy Connection Establishment
+        Connections are established only when this method is first called.
+
+        Returns:
+            FlockMCPClientManager instance
+
+        Raises:
+            RuntimeError: If no MCP servers registered
+        """
+        if not self._configs:
+            raise RuntimeError("No MCP servers registered. Call add_mcp() first.")
+
+        if self._client_manager is None:
+            from flock.mcp import FlockMCPClientManager
+
+            self._client_manager = FlockMCPClientManager(self._configs)
+
+        return self._client_manager
+
+    async def cleanup(self) -> None:
+        """Clean up MCP connections.
+
+        Called during orchestrator shutdown to properly close all MCP connections.
+        """
+        if self._client_manager is not None:
+            await self._client_manager.cleanup_all()
+            self._client_manager = None
+
+    @property
+    def configs(self) -> dict[str, FlockMCPConfiguration]:
+        """Get the dictionary of MCP configurations."""
+        return self._configs
+
+    @property
+    def has_configs(self) -> bool:
+        """Check if any MCP servers are registered."""
+        return bool(self._configs)
+
+
+__all__ = ["MCPManager"]

flock/orchestrator/scheduler.py

@@ -0,0 +1,189 @@
+"""Agent scheduling engine."""
+
+from __future__ import annotations
+
+import asyncio
+from asyncio import Task
+from typing import TYPE_CHECKING, Any
+
+from flock.components.orchestrator import ScheduleDecision
+
+
+if TYPE_CHECKING:
+    from flock.agent import Agent
+    from flock.core import Flock
+    from flock.core.artifacts import Artifact
+    from flock.core.visibility import AgentIdentity
+    from flock.orchestrator import ComponentRunner
+
+
+class AgentScheduler:
+    """Schedules agents for execution based on artifact subscriptions.
+
+    Responsibilities:
+    - Match artifacts to agent subscriptions
+    - Run scheduling hooks via ComponentRunner
+    - Create agent execution tasks
+    - Manage task lifecycle
+    - Track processed artifacts for deduplication
+    """
+
+    def __init__(self, orchestrator: Flock, component_runner: ComponentRunner):
+        """Initialize scheduler.
+
+        Args:
+            orchestrator: Flock orchestrator instance
+            component_runner: Runner for executing component hooks
+        """
+        self._orchestrator = orchestrator
+        self._component_runner = component_runner
+        self._tasks: set[Task[Any]] = set()
+        self._processed: set[tuple[str, str]] = set()
+        self._logger = orchestrator._logger
+
+    async def schedule_artifact(self, artifact: Artifact) -> None:
+        """Schedule agents for an artifact using component hooks.
+
+        Args:
+            artifact: Published artifact to match against subscriptions
+        """
+        # Initialize components on first artifact
+        if not self._component_runner.is_initialized:
+            await self._component_runner.run_initialize(self._orchestrator)
+
+        # Component hook - artifact published (can transform or block)
+        artifact = await self._component_runner.run_artifact_published(
+            self._orchestrator, artifact
+        )
+        if artifact is None:
+            return  # Artifact blocked by component
+
+        for agent in self._orchestrator.agents:
+            identity = agent.identity
+            for subscription in agent.subscriptions:
+                if not subscription.accepts_events():
+                    continue
+
+                # Check prevent_self_trigger
+                if agent.prevent_self_trigger and artifact.produced_by == agent.name:
+                    continue  # Skip - agent produced this artifact
+
+                # Visibility check
+                if not self._check_visibility(artifact, identity):
+                    continue
+
+                # Subscription match check
+                if not subscription.matches(artifact):
+                    continue
+
+                # Component hook - before schedule (circuit breaker, deduplication)
+                decision = await self._component_runner.run_before_schedule(
+                    self._orchestrator, artifact, agent, subscription
+                )
+                if decision == ScheduleDecision.SKIP:
+                    continue
+                if decision == ScheduleDecision.DEFER:
+                    continue
+
+                # Component hook - collect artifacts (AND gates, correlation, batching)
+                collection = await self._component_runner.run_collect_artifacts(
+                    self._orchestrator, artifact, agent, subscription
+                )
+                if not collection.complete:
+                    continue  # Still collecting
+
+                artifacts = collection.artifacts
+
+                # Component hook - before agent schedule (final validation)
+                artifacts = await self._component_runner.run_before_agent_schedule(
+                    self._orchestrator, agent, artifacts
+                )
+                if artifacts is None:
+                    continue  # Scheduling blocked
+
+                # Schedule agent task
+                is_batch_execution = subscription.batch is not None
+                task = self.schedule_task(agent, artifacts, is_batch=is_batch_execution)
+
+                # Component hook - agent scheduled (notification)
+                await self._component_runner.run_agent_scheduled(
+                    self._orchestrator, agent, artifacts, task
+                )
+
+    def schedule_task(
+        self, agent: Agent, artifacts: list[Artifact], is_batch: bool = False
+    ) -> Task[Any]:
+        """Schedule agent task and return the task handle.
+
+        Args:
+            agent: Agent to execute
+            artifacts: Input artifacts
+            is_batch: Whether this is batch execution
+
+        Returns:
+            Asyncio task handle
+        """
+        task = asyncio.create_task(
+            self._orchestrator._run_agent_task(agent, artifacts, is_batch=is_batch)
+        )
+        self._tasks.add(task)
+        task.add_done_callback(self._tasks.discard)
+        return task
+
+    def record_agent_run(self, agent: Agent) -> None:
+        """Record agent run metric.
+
+        Args:
+            agent: Agent that ran
+        """
+        self._orchestrator.metrics["agent_runs"] += 1
+
+    def mark_processed(self, artifact: Artifact, agent: Agent) -> None:
+        """Mark artifact as processed by agent.
+
+        Args:
+            artifact: Processed artifact
+            agent: Agent that processed it
+        """
+        key = (str(artifact.id), agent.name)
+        self._processed.add(key)
+
+    def seen_before(self, artifact: Artifact, agent: Agent) -> bool:
+        """Check if artifact was already processed by agent.
+
+        Args:
+            artifact: Artifact to check
+            agent: Agent to check
+
+        Returns:
+            True if already processed
+        """
+        key = (str(artifact.id), agent.name)
+        return key in self._processed
+
+    def _check_visibility(self, artifact: Artifact, identity: AgentIdentity) -> bool:
+        """Check if artifact is visible to agent.
+
+        Args:
+            artifact: Artifact to check
+            identity: Agent identity
+
+        Returns:
+            True if visible
+        """
+        try:
+            return artifact.visibility.allows(identity)
+        except AttributeError:  # pragma: no cover - fallback
+            return True
+
+    @property
+    def pending_tasks(self) -> set[Task[Any]]:
+        """Get set of pending agent tasks.
+
+        Returns:
+            Set of asyncio tasks
+        """
+        return self._tasks
+
+
+__all__ = ["AgentScheduler"]