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

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"]

flock/orchestrator/server_manager.py

@@ -0,0 +1,234 @@
+"""HTTP server management for orchestrator.
+
+Handles service startup with optional dashboard integration.
+Extracted from orchestrator.py to reduce complexity.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from asyncio import Task
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+
+if TYPE_CHECKING:
+    from flock.core.orchestrator import Flock
+
+
+class ServerManager:
+    """Manages HTTP service startup for the orchestrator.
+
+    Handles both standard API mode and dashboard mode with WebSocket support.
+    """
+
+    @staticmethod
+    async def serve(
+        orchestrator: Flock,
+        *,
+        dashboard: bool = False,
+        dashboard_v2: bool = False,
+        host: str = "127.0.0.1",
+        port: int = 8344,
+        blocking: bool = True,
+    ) -> Task[None] | None:
+        """Start HTTP service for the orchestrator.
+
+        Args:
+            orchestrator: The Flock orchestrator instance to serve
+            dashboard: Enable real-time dashboard with WebSocket support (default: False)
+            dashboard_v2: Launch the new dashboard v2 frontend (implies dashboard=True)
+            host: Host to bind to (default: "127.0.0.1")
+            port: Port to bind to (default: 8344)
+            blocking: If True, blocks until server stops. If False, starts server
+                in background and returns task handle (default: True)
+
+        Returns:
+            None if blocking=True, or Task handle if blocking=False
+
+        Examples:
+            # Basic HTTP API (no dashboard) - runs until interrupted
+            await ServerManager.serve(orchestrator)
+
+            # With dashboard (WebSocket + browser launch) - runs until interrupted
+            await ServerManager.serve(orchestrator, dashboard=True)
+
+            # Non-blocking mode - start server in background
+            task = await ServerManager.serve(orchestrator, dashboard=True, blocking=False)
+            # Now you can publish messages and run other logic
+            await orchestrator.publish(my_message)
+            await orchestrator.run_until_idle()
+        """
+        # If non-blocking, start server in background task
+        if not blocking:
+            server_task = asyncio.create_task(
+                ServerManager._serve_impl(
+                    orchestrator,
+                    dashboard=dashboard,
+                    dashboard_v2=dashboard_v2,
+                    host=host,
+                    port=port,
+                )
+            )
+            # Add cleanup callback
+            server_task.add_done_callback(
+                lambda task: ServerManager._cleanup_server_callback(orchestrator, task)
+            )
+            # Store task reference for later cancellation
+            orchestrator._server_task = server_task
+            # Give server a moment to start
+            await asyncio.sleep(0.1)
+            return server_task
+
+        # Blocking mode - run server directly with cleanup
+        try:
+            await ServerManager._serve_impl(
+                orchestrator,
+                dashboard=dashboard,
+                dashboard_v2=dashboard_v2,
+                host=host,
+                port=port,
+            )
+        finally:
+            # In blocking mode, manually cleanup dashboard launcher
+            if (
+                hasattr(orchestrator, "_dashboard_launcher")
+                and orchestrator._dashboard_launcher is not None
+            ):
+                orchestrator._dashboard_launcher.stop()
+                orchestrator._dashboard_launcher = None
+        return None
+
+    @staticmethod
+    def _cleanup_server_callback(orchestrator: Flock, task: Task[None]) -> None:
+        """Cleanup callback when background server task completes."""
+        # Stop dashboard launcher if it was started
+        if (
+            hasattr(orchestrator, "_dashboard_launcher")
+            and orchestrator._dashboard_launcher is not None
+        ):
+            try:
+                orchestrator._dashboard_launcher.stop()
+            except Exception as e:
+                orchestrator._logger.warning(f"Failed to stop dashboard launcher: {e}")
+            finally:
+                orchestrator._dashboard_launcher = None
+
+        # Clear server task reference
+        if hasattr(orchestrator, "_server_task"):
+            orchestrator._server_task = None
+
+        # Log any exceptions from the task
+        try:
+            exc = task.exception()
+            if exc and not isinstance(exc, asyncio.CancelledError):
+                orchestrator._logger.error(f"Server task failed: {exc}", exc_info=exc)
+        except asyncio.CancelledError:
+            pass  # Normal cancellation
+
+    @staticmethod
+    async def _serve_impl(
+        orchestrator: Flock,
+        *,
+        dashboard: bool = False,
+        dashboard_v2: bool = False,
+        host: str = "127.0.0.1",
+        port: int = 8344,
+    ) -> None:
+        """Internal implementation of serve() - actual server logic."""
+        if dashboard_v2:
+            dashboard = True
+
+        if not dashboard:
+            # Standard service without dashboard
+            await ServerManager._serve_standard(orchestrator, host=host, port=port)
+            return
+
+        # Dashboard mode with WebSocket and event collection
+        await ServerManager._serve_dashboard(
+            orchestrator, dashboard_v2=dashboard_v2, host=host, port=port
+        )
+
+    @staticmethod
+    async def _serve_standard(orchestrator: Flock, *, host: str, port: int) -> None:
+        """Serve standard HTTP API without dashboard.
+
+        Args:
+            orchestrator: The Flock orchestrator instance
+            host: Host to bind to
+            port: Port to bind to
+        """
+        from flock.api.service import BlackboardHTTPService
+
+        service = BlackboardHTTPService(orchestrator)
+        await service.run_async(host=host, port=port)
+
+    @staticmethod
+    async def _serve_dashboard(
+        orchestrator: Flock, *, dashboard_v2: bool, host: str, port: int
+    ) -> None:
+        """Serve HTTP API with dashboard and WebSocket support.
+
+        Args:
+            orchestrator: The Flock orchestrator instance
+            dashboard_v2: Whether to use v2 dashboard frontend
+            host: Host to bind to
+            port: Port to bind to
+        """
+        from flock.core import Agent
+        from flock.dashboard.collector import DashboardEventCollector
+        from flock.dashboard.launcher import DashboardLauncher
+        from flock.dashboard.service import DashboardHTTPService
+        from flock.dashboard.websocket import WebSocketManager
+
+        # Create dashboard components
+        websocket_manager = WebSocketManager()
+        event_collector = DashboardEventCollector(store=orchestrator.store)
+        event_collector.set_websocket_manager(websocket_manager)
+        await event_collector.load_persistent_snapshots()
+
+        # Store collector reference for agents added later
+        orchestrator._dashboard_collector = event_collector
+        # Store websocket manager for real-time event emission (Phase 1.2)
+        orchestrator._websocket_manager = websocket_manager
+        # Phase 5A: Set websocket manager on EventEmitter for dashboard updates
+        orchestrator._event_emitter.set_websocket_manager(websocket_manager)
+
+        # Phase 6+7: Set class-level WebSocket broadcast wrapper (dashboard mode)
+        async def _broadcast_wrapper(event):
+            """Isolated broadcast wrapper - no reference chain to orchestrator."""
+            return await websocket_manager.broadcast(event)
+
+        Agent._websocket_broadcast_global = _broadcast_wrapper
+
+        # Inject event collector into all existing agents
+        for agent in orchestrator._agents.values():
+            # Add dashboard collector with priority ordering handled by agent
+            agent._add_utilities([event_collector])
+
+        # Start dashboard launcher (npm process + browser)
+        launcher_kwargs: dict[str, Any] = {"port": port}
+        if dashboard_v2:
+            dashboard_pkg_dir = Path(__file__).parent.parent / "dashboard"
+            launcher_kwargs["frontend_dir"] = dashboard_pkg_dir.parent / "frontend_v2"
+            launcher_kwargs["static_dir"] = dashboard_pkg_dir / "static_v2"
+
+        launcher = DashboardLauncher(**launcher_kwargs)
+        launcher.start()
+
+        # Create dashboard HTTP service
+        service = DashboardHTTPService(
+            orchestrator=orchestrator,
+            websocket_manager=websocket_manager,
+            event_collector=event_collector,
+            use_v2=dashboard_v2,
+        )
+
+        # Store launcher for cleanup
+        orchestrator._dashboard_launcher = launcher
+
+        # Run service (blocking call)
+        # Note: Cleanup is NOT done here - it's handled by:
+        # - ServerManager.serve() finally block (blocking mode)
+        # - ServerManager._cleanup_server_callback() (non-blocking mode)
+        await service.run_async(host=host, port=port)