flock-core 0.4.511__py3-none-any.whl → 0.4.513__py3-none-any.whl

flock/core/mcp/types/types.py

@@ -1,9 +1,13 @@
 """Types for Flock's MCP functionality."""
 
+import importlib
+import inspect
+import os
 from collections.abc import Awaitable, Callable
 from contextlib import AbstractAsyncContextManager
-from typing import Any
+from typing import Any, Literal
 
+import httpx
 from anyio.streams.memory import (
     MemoryObjectReceiveStream,
     MemoryObjectSendStream,
@@ -36,6 +40,7 @@ from mcp.types import (
 from pydantic import AnyUrl, BaseModel, ConfigDict, Field
 
 from flock.core.mcp.util.helpers import get_default_env
+from flock.core.serialization.serializable import Serializable
 
 
 class ServerNotification(_MCPServerNotification):
@@ -80,17 +85,40 @@ class MCPRoot(_MCPRoot):
     """Wrapper for mcp.types.Root."""
 
 
-class ServerParameters(BaseModel):
+class ServerParameters(BaseModel, Serializable):
     """Base Type for server parameters."""
 
     model_config = ConfigDict(
         arbitrary_types_allowed=True,
     )
 
+    transport_type: Literal["stdio", "websockets", "sse", "streamable_http"] = Field(
+        ...,
+        description="which type of transport these connection params are used for."
+    )
+
+    def to_dict(self, path_type: str = "relative"):
+        """Serialize."""
+        return self.model_dump(
+            exclude_defaults=False,
+            exclude_none=True,
+            mode="json"
+        )
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]):
+        """Deserialize."""
+        return cls(**data)
+
 
 class StdioServerParameters(_MCPStdioServerParameters, ServerParameters):
     """Base Type for Stdio Server parameters."""
 
+    transport_type: Literal["stdio"] = Field(
+        default="stdio",
+        description="Use stdio params."
+    )
+
     env: dict[str, str] | None = Field(
         default_factory=get_default_env,
         description="Environment for the MCP Server.",
@@ -100,12 +128,121 @@ class StdioServerParameters(_MCPStdioServerParameters, ServerParameters):
 class WebsocketServerParameters(ServerParameters):
     """Base Type for Websocket Server params."""
 
+    transport_type: Literal["websockets"] = Field(
+        default="websockets",
+        description="Use websocket params."
+    )
+
     url: str | AnyUrl = Field(..., description="Url the server listens at.")
 
+class StreamableHttpServerParameters(ServerParameters):
+    """Base Type for StreamableHttp params."""
+
+    transport_type: Literal["streamable_http"] = Field(
+        default="streamable_http",
+        description="Use streamable http params."
+    )
+
+    url: str | AnyUrl = Field(
+        ...,
+        description="The url the server listens at."
+    )
+
+    headers: dict[str, Any] | None = Field(
+        default=None,
+        description="Additional headers to pass to the client."
+    )
+
+    timeout: float | int = Field(
+        default=5,
+        description="Http Timeout",
+    )
+
+    sse_read_timeout: float | int = Field(
+        default=60*5,
+        description="How long the client will wait before disconnecting from the server."
+    )
+
+    terminate_on_close: bool = Field(
+        default=True,
+        description="Terminate connection on close"
+    )
+
+    auth: httpx.Auth | None = Field(
+        default=None,
+        description="Httpx Auth Scheme"
+    )
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]):
+        """Deserialize the object from a dict."""
+        # find and import the concrete implementation for
+        # the auth object
+        auth_obj: httpx.Auth | None = None
+        auth_impl = data.pop("auth", None)
+        if auth_impl:
+            # find the concrete implementation
+            impl = auth_impl.pop("implementation", None)
+            params = auth_impl.pop("params", None)
+            if impl:
+                mod = importlib.import_module(impl["module_path"])
+                real_cls = getattr(mod, impl["classname"])
+                if params:
+                    auth_obj = real_cls(**{k: v for k, v in params.items()})
+                else:
+                    # assume that the implementation handles it.
+                    auth_obj = real_cls()
+            else:
+                raise ValueError("No concrete implementation for auth provided.")
+
+        data["auth"] = auth_obj
+        return cls(**{k: v for k, v in data.items()})
+
+    def to_dict(self, path_type = "relative"):
+        """Serialize the object."""
+        exclude = ["auth"]
+
+        data = self.model_dump(
+            exclude=exclude,
+            exclude_defaults=False,
+            exclude_none=True,
+        )
+
+        # inject implentation info for auth
+        if self.auth is not None:
+            file_path = inspect.getsourcefile(type(self.auth))
+            if path_type == "relative":
+                file_path = os.path.relpath(file_path)
+            try:
+                # params should be primitive types, keeping with the
+                # declarative approach of flock.
+                params = {
+                    k: getattr(self.auth, k)
+                    for k in getattr(self.auth, "__dict__", {})
+                    if not k.startswith("_")
+                }
+            except Exception:
+                params = None
+
+        data["auth"] = {
+            "implementation": {
+                "class_name": type(self.auth).__name__,
+                "module_path": type(self.auth).__module__,
+                "file_path": file_path,
+            },
+            "params": params,
+        }
+
+        return data
 
 class SseServerParameters(ServerParameters):
     """Base Type for SSE Server params."""
 
+    transport_type: Literal["sse"] = Field(
+        default="sse",
+        description="Use sse server params."
+    )
+
     url: str | AnyUrl = Field(..., description="The url the server listens at.")
 
     headers: dict[str, Any] | None = Field(
@@ -119,6 +256,72 @@ class SseServerParameters(ServerParameters):
         description="How long the client will wait before disconnecting from the server.",
     )
 
+    auth: httpx.Auth | None = Field(
+        default=None,
+        description="Httpx Auth Scheme."
+    )
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]):
+        """Deserialize the object from a dict."""
+        # find and import the concrete implementation for
+        # the auth object.
+        auth_obj: httpx.Auth | None = None
+        auth_impl = data.pop("auth", None) # get the specs for the auth class
+        if auth_impl:
+            # find the concrete implementation
+            impl = auth_impl.pop("implementation", None)
+            params = auth_impl.pop("params", None)
+            if impl:
+                mod = importlib.import_module(impl["module_path"])
+                real_cls = getattr(mod, impl["class_name"])
+                if params:
+                    auth_obj = real_cls(**{k: v for k, v in params.items()})
+                else:
+                    # assume that implementation handles it
+                    auth_obj = real_cls()
+            else:
+                raise ValueError("No concrete implementation for auth provided.")
+
+        data["auth"] = auth_obj
+        return cls(**{k: v for k, v in data.items()})
+
+    def to_dict(self, path_type = "relative"):
+        """Serialize the object."""
+        exclude = ["auth"]
+
+        data = self.model_dump(
+            exclude=exclude,
+            exclude_defaults=False,
+            exclude_none=True,
+        )
+
+        # inject implentation info for auth
+        if self.auth is not None:
+            file_path = inspect.getsourcefile(type(self.auth))
+            if path_type == "relative":
+                file_path = os.path.relpath(file_path)
+            try:
+                # params should be primitive types, keeping with the
+                # declarative approach of flock.
+                params = {
+                    k: getattr(self.auth, k)
+                    for k in getattr(self.auth, "__dict__", {})
+                    if not k.startswith("_")
+                }
+            except Exception:
+                params = None
+
+        data["auth"] = {
+            "implementation": {
+                "class_name": type(self.auth).__name__,
+                "module_path": type(self.auth).__module__,
+                "file_path": file_path,
+            },
+            "params": params,
+        }
+
+        return data
 
 MCPCLientInitFunction = Callable[
     ...,

flock/mcp/servers/sse/flock_sse_server.py

@@ -4,12 +4,13 @@ import copy
 from contextlib import AbstractAsyncContextManager
 from typing import Any, Literal
 
+import httpx
 from anyio.streams.memory import (
     MemoryObjectReceiveStream,
     MemoryObjectSendStream,
 )
 from mcp.client.sse import sse_client
-from mcp.types import JSONRPCMessage
+from mcp.shared.message import SessionMessage
 from opentelemetry import trace
 from pydantic import Field
 
@@ -63,45 +64,51 @@ class FlockSSEClient(FlockMCPClientBase):
         additional_params: dict[str, Any] | None = None,
     ) -> AbstractAsyncContextManager[
         tuple[
-            MemoryObjectReceiveStream[JSONRPCMessage | Exception],
-            MemoryObjectSendStream[JSONRPCMessage],
+            MemoryObjectReceiveStream[SessionMessage | Exception],
+            MemoryObjectSendStream[SessionMessage],
         ]
     ]:
         """Return an async context manager whose __aenter__ method yields (read_stream, send_stream)."""
         # avoid modifying the config of the client as a side-effect.
         param_copy = copy.deepcopy(params)
 
-        if self.additional_params:
+        if additional_params:
             override_headers = bool(
-                self.additional_params.get("override_headers", False)
+                additional_params.get("override_headers", False)
             )
-            if "headers" in self.additional_params:
+            if "headers" in additional_params:
                 if override_headers:
-                    param_copy.headers = self.additional_params.get(
+                    param_copy.headers = additional_params.get(
                         "headers", params.headers
                     )
                 else:
                     param_copy.headers.update(
-                        self.additional_params.get("headers", {})
+                        additional_params.get("headers", {})
                     )
-            if "read_timeout_seconds" in self.additional_params:
-                param_copy.timeout = self.additional_params.get(
+            if "read_timeout_seconds" in additional_params:
+                param_copy.timeout =  additional_params.get(
                     "read_timeout_seconds", params.timeout
                 )
 
-            if "sse_read_timeout" in self.additional_params:
-                param_copy.sse_read_timeout = self.additional_params.get(
+            if "sse_read_timeout" in additional_params:
+                param_copy.sse_read_timeout = additional_params.get(
                     "sse_read_timeout",
                     params.sse_read_timeout,
                 )
-            if "url" in self.additional_params:
-                param_copy.url = self.additional_params.get(
+            if "url" in additional_params:
+                param_copy.url = additional_params.get(
                     "url",
                     params.url,
                 )
 
+            if "auth" in additional_params and isinstance(
+                additional_params.get("auth"), httpx.Auth
+            ):
+                param_copy.auth = additional_params.get("auth", param_copy.auth)
+
         return sse_client(
             url=param_copy.url,
+            auth=param_copy.auth,
             headers=param_copy.headers,
             timeout=float(param_copy.timeout),
             sse_read_timeout=float(param_copy.sse_read_timeout),

flock/mcp/servers/streamable_http/flock_streamable_http_server.py

@@ -0,0 +1,169 @@
+"""This module provides the Flock Streamable-Http functionality."""
+
+import copy
+from collections.abc import Callable
+from contextlib import AbstractAsyncContextManager
+from datetime import timedelta
+from typing import Any, Literal
+
+import httpx
+from anyio.streams.memory import (
+    MemoryObjectReceiveStream,
+    MemoryObjectSendStream,
+)
+from mcp.client.streamable_http import streamablehttp_client
+from mcp.shared.message import SessionMessage
+from opentelemetry import trace
+from pydantic import Field
+
+from flock.core.logging.logging import get_logger
+from flock.core.mcp.flock_mcp_server import FlockMCPServerBase
+from flock.core.mcp.mcp_client import FlockMCPClientBase
+from flock.core.mcp.mcp_client_manager import FlockMCPClientManagerBase
+from flock.core.mcp.mcp_config import (
+    FlockMCPConfigurationBase,
+    FlockMCPConnectionConfigurationBase,
+)
+from flock.core.mcp.types.types import (
+    StreamableHttpServerParameters,
+)
+
+logger = get_logger("mcp.streamable_http.server")
+tracer = trace.get_tracer(__name__)
+
+GetSessionIdCallback = Callable[[], str | None]
+
+
+class FlockStreamableHttpConnectionConfig(FlockMCPConnectionConfigurationBase):
+    """Concrete ConnectionConfig for a StreamableHttpClient."""
+
+    # Only thing we need to override here is the concrete transport_type
+    # and connection parameter fields.
+    transport_type: Literal["streamable_http"] = Field(
+        default="streamable_http",
+        description="Use the streamable_http Transport type.",
+    )
+
+    connection_parameters: StreamableHttpServerParameters = Field(
+        ..., description="Streamable HTTP Server Connection Parameters."
+    )
+
+
+class FlockStreamableHttpConfig(FlockMCPConfigurationBase):
+    """Configuration for Streamable HTTP Clients."""
+
+    # The only thing we need to override here is the
+    # concrete connection config.
+    # The rest is generic enough to handle everything else.
+    connection_config: FlockStreamableHttpConnectionConfig = Field(
+        ..., description="Concrete StreamableHttp Connection Configuration."
+    )
+
+
+class FlockStreamableHttpClient(FlockMCPClientBase):
+    """Client for StreamableHttpServers."""
+
+    config: FlockStreamableHttpConfig = Field(
+        ..., description="Client configuration."
+    )
+
+    async def create_transport(
+        self,
+        params: StreamableHttpServerParameters,
+        additional_params: dict[str, Any] | None = None,
+    ) -> AbstractAsyncContextManager[
+    tuple[
+        MemoryObjectReceiveStream[SessionMessage | Exception],
+        MemoryObjectSendStream[SessionMessage],
+        GetSessionIdCallback,
+    ],
+    None,
+]:
+        """Return an async context manager whose __aenter__ method yields (read_stream, send_stream)."""
+        param_copy = copy.deepcopy(params)
+
+        if additional_params:
+            override_headers = bool(
+                additional_params.get("override_headers", False)
+            )
+
+            if "headers" in additional_params:
+                if override_headers:
+                    param_copy.headers = additional_params.get(
+                        "headers", params.headers
+                    )
+                else:
+                    param_copy.headers.update(
+                        additional_params.get("headers", {})
+                    )
+            if "auth" in additional_params and isinstance(
+                additional_params.get("auth"), httpx.Auth
+            ):
+                param_copy.auth = additional_params.get("auth", param_copy.auth)
+
+            if "read_timeout_seconds" in additional_params:
+                param_copy.timeout = additional_params.get(
+                    "read_timeout_seconds", params.timeout
+                )
+
+            if "sse_read_timeout" in additional_params:
+                param_copy.sse_read_timeout = additional_params.get(
+                    "sse_read_timeout",
+                    params.sse_read_timeout,
+                )
+            if "url" in additional_params:
+                param_copy.url = additional_params.get(
+                    "url",
+                    params.url,
+                )
+
+            if "terminate_on_close" in additional_params:
+                param_copy.terminate_on_close = bool(
+                    additional_params.get("terminate_on_close", True)
+                )
+
+        timeout_http = timedelta(seconds=param_copy.timeout)
+        sse_timeout = timedelta(seconds=param_copy.sse_read_timeout)
+
+        return streamablehttp_client(
+            url=param_copy.url,
+            headers=param_copy.headers,
+            timeout=timeout_http,
+            sse_read_timeout=sse_timeout,
+            terminate_on_close=param_copy.terminate_on_close,
+            auth=param_copy.auth,
+        )
+
+
+class FlockStreamableHttpClientManager(FlockMCPClientManagerBase):
+    """Manager for handling StreamableHttpClients."""
+
+    client_config: FlockStreamableHttpConfig = Field(
+        ..., description="Configuration for clients."
+    )
+
+    async def make_client(
+        self, additional_params: dict[str, Any] | None = None
+    ) -> FlockStreamableHttpClient:
+        """Create a new client instance."""
+        new_client = FlockStreamableHttpClient(
+            config=self.client_config,
+            additional_params=additional_params,
+        )
+        return new_client
+
+
+class FlockStreamableHttpServer(FlockMCPServerBase):
+    """Class which represents a MCP Server using the streamable Http Transport type."""
+
+    config: FlockStreamableHttpConfig = Field(
+        ..., description="Config for the server."
+    )
+
+    async def initialize(self) -> FlockStreamableHttpClientManager:
+        """Called when initializing the server."""
+        client_manager = FlockStreamableHttpClientManager(
+            client_config=self.config
+        )
+
+        return client_manager

flock/mcp/servers/websockets/flock_websocket_server.py

@@ -9,7 +9,7 @@ from anyio.streams.memory import (
     MemoryObjectSendStream,
 )
 from mcp.client.websocket import websocket_client
-from mcp.types import JSONRPCMessage
+from mcp.shared.message import SessionMessage
 from opentelemetry import trace
 from pydantic import Field
 
@@ -69,8 +69,8 @@ class FlockWSClient(FlockMCPClientBase):
         additional_params: dict[str, Any] | None = None,
     ) -> AbstractAsyncContextManager[
         tuple[
-            MemoryObjectReceiveStream[JSONRPCMessage | Exception],
-            MemoryObjectSendStream[JSONRPCMessage],
+            MemoryObjectReceiveStream[SessionMessage | Exception],
+            MemoryObjectSendStream[SessionMessage],
         ]
     ]:
         """Return an async context manager whose __aenter__ method yields a read_stream and a send_stream."""

flock/webapp/app/main.py

@@ -1,4 +1,5 @@
 # src/flock/webapp/app/main.py
+import asyncio
 import json
 import os  # Added import
 import shutil
@@ -8,6 +9,7 @@ import urllib.parse
 import uuid
 from contextlib import asynccontextmanager
 from pathlib import Path
+from typing import Any
 
 import markdown2  # Import markdown2
 from fastapi import (
@@ -24,13 +26,13 @@ from fastapi.responses import HTMLResponse, RedirectResponse
 from fastapi.staticfiles import StaticFiles
 from fastapi.templating import Jinja2Templates
 from pydantic import BaseModel
-from typing import Any
 
 from flock.core.api.endpoints import create_api_router
 from flock.core.api.run_store import RunStore
 
 # Import core Flock components and API related modules
 from flock.core.flock import Flock  # For type hinting
+from flock.core.flock_scheduler import FlockScheduler
 from flock.core.logging.logging import get_logger  # For logging
 from flock.core.util.spliter import parse_schema
 
@@ -44,7 +46,6 @@ from flock.webapp.app.api import (
 from flock.webapp.app.config import (
     DEFAULT_THEME_NAME,
     FLOCK_FILES_DIR,
-    SHARED_LINKS_DB_PATH,
     THEMES_DIR,
     get_current_theme_name,
 )
@@ -72,7 +73,7 @@ from flock.webapp.app.services.flock_service import (
 from flock.webapp.app.services.sharing_models import SharedLinkConfig
 from flock.webapp.app.services.sharing_store import (
     SharedLinkStoreInterface,
-    SQLiteSharedLinkStore,
+    create_shared_link_store,
 )
 from flock.webapp.app.theme_mapper import alacritty_to_pico
 
@@ -138,23 +139,21 @@ async def lifespan(app: FastAPI):
     logger.info("FastAPI application starting up...")
     # Flock instance and RunStore are expected to be set on app.state
     # by `start_unified_server` in `webapp/run.py` *before* uvicorn starts the app.
-    # The call to `set_global_flock_services` also happens there.
-
-    # Initialize and set the SharedLinkStore
+    # The call to `set_global_flock_services` also happens there.    # Initialize and set the SharedLinkStore
     try:
-        logger.info(f"Initializing SharedLinkStore with DB path: {SHARED_LINKS_DB_PATH}")
-        shared_link_store = SQLiteSharedLinkStore(db_path=str(SHARED_LINKS_DB_PATH))
+        logger.info("Initializing SharedLinkStore using factory...")
+        shared_link_store = create_shared_link_store()
         await shared_link_store.initialize() # Create tables if they don't exist
         set_global_shared_link_store(shared_link_store)
         logger.info("SharedLinkStore initialized and set globally.")
     except Exception as e:
-        logger.error(f"Failed to initialize SharedLinkStore: {e}", exc_info=True)    # Configure chat features with clear precedence:
+        logger.error(f"Failed to initialize SharedLinkStore: {e}", exc_info=True)# Configure chat features with clear precedence:
     # 1. Value set by start_unified_server (programmatic)
     # 2. Environment variables (standalone mode)
     programmatic_chat_enabled = getattr(app.state, "chat_enabled", None)
     env_start_mode = os.environ.get("FLOCK_START_MODE")
     env_chat_enabled = os.environ.get("FLOCK_CHAT_ENABLED", "false").lower() == "true"
-    
+
     if programmatic_chat_enabled is not None:
         # Programmatic setting takes precedence (from start_unified_server)
         should_enable_chat_routes = programmatic_chat_enabled
@@ -235,11 +234,66 @@ async def lifespan(app: FastAPI):
             logger.info(f"Lifespan: Added {len(pending_endpoints)} custom API routes to main app.")
         else:
             logger.warning("Lifespan: Pending custom endpoints found, but no Flock instance in app.state. Cannot add custom routes.")
+
+    # --- Add Scheduler Startup Logic ---
+    flock_instance_from_state: Flock | None = getattr(app.state, "flock_instance", None)
+    if flock_instance_from_state:
+        # Create and start the scheduler
+        scheduler = FlockScheduler(flock_instance_from_state)
+        app.state.flock_scheduler = scheduler  # Store for access during shutdown
+
+        scheduler_loop_task = await scheduler.start() # Start returns the task
+        if scheduler_loop_task:
+            app.state.flock_scheduler_task = scheduler_loop_task # Store the task
+            logger.info("FlockScheduler background task started.")
+        else:
+            app.state.flock_scheduler_task = None
+            logger.info("FlockScheduler initialized, but no scheduled agents found or loop not started.")
+    else:
+        app.state.flock_scheduler = None
+        app.state.flock_scheduler_task = None
+        logger.warning("No Flock instance found in app.state; FlockScheduler not started.")
+    # --- End Scheduler Startup Logic ---
+
     yield
     logger.info("FastAPI application shutting down...")
 
-app = FastAPI(title="Flock Web UI & API", lifespan=lifespan)
+     # --- Add Scheduler Shutdown Logic ---
+    logger.info("FastAPI application initiating shutdown...")
+    scheduler_to_stop: FlockScheduler | None = getattr(app.state, "flock_scheduler", None)
+    scheduler_task_to_await: asyncio.Task | None = getattr(app.state, "flock_scheduler_task", None)
 
+    if scheduler_to_stop:
+        logger.info("Attempting to stop FlockScheduler...")
+        await scheduler_to_stop.stop() # Signal the scheduler loop to stop
+
+        if scheduler_task_to_await and not scheduler_task_to_await.done():
+            logger.info("Waiting for FlockScheduler task to complete...")
+            try:
+                await asyncio.wait_for(scheduler_task_to_await, timeout=10.0) # Wait for graceful exit
+                logger.info("FlockScheduler task completed gracefully.")
+            except asyncio.TimeoutError:
+                logger.warning("FlockScheduler task did not complete in time, cancelling.")
+                scheduler_task_to_await.cancel()
+                try:
+                    await scheduler_task_to_await # Await cancellation
+                except asyncio.CancelledError:
+                    logger.info("FlockScheduler task cancelled.")
+            except Exception as e:
+                logger.error(f"Error during FlockScheduler task finalization: {e}", exc_info=True)
+        elif scheduler_task_to_await and scheduler_task_to_await.done():
+            logger.info("FlockScheduler task was already done.")
+        else:
+            logger.info("FlockScheduler instance found, but no running task was stored to await.")
+    else:
+        logger.info("No active FlockScheduler found to stop.")
+
+    logger.info("FastAPI application finished shutdown sequence.")
+    # --- End Scheduler Shutdown Logic ---
+
+app = FastAPI(title="Flock Web UI & API", lifespan=lifespan, docs_url="/docs",
+    openapi_url="/openapi.json", root_path=os.getenv("FLOCK_ROOT_PATH", ""))
+logger.info("FastAPI booting complete.")
 BASE_DIR = Path(__file__).resolve().parent.parent
 app.mount("/static", StaticFiles(directory=str(BASE_DIR / "static")), name="static")
 templates = Jinja2Templates(directory=str(BASE_DIR / "templates"))
@@ -826,6 +880,7 @@ async def ui_create_flock_action(request: Request, flock_name: str = Form(...),
     context = get_base_context_web(request, success=success_msg_text, ui_mode=ui_mode_query)
     return templates.TemplateResponse("partials/_execution_view_container.html", context, headers=response_headers)
 
+
 # --- Settings Page & Endpoints ---
 @app.get("/ui/settings", response_class=HTMLResponse, tags=["UI Pages"])
 async def page_settings(request: Request, error: str = None, success: str = None, ui_mode: str = Query("standalone")):