flock-core 0.5.0b1__py3-none-any.whl → 0.5.0b3__py3-none-any.whl

flock/core/flock_factory.py

@@ -16,13 +16,13 @@ from flock.components.utility.metrics_utility_component import (
 # New unified components imported locally to avoid circular imports
 from flock.core.config.flock_agent_config import FlockAgentConfig
 from flock.core.config.scheduled_agent_config import ScheduledAgentConfig
-from flock.core.flock_agent import FlockAgent, SignatureType
+from flock.core.flock_agent import DynamicStr, FlockAgent
 from flock.core.logging.formatters.themes import OutputTheme
-from flock.core.mcp.flock_mcp_server import FlockMCPServerBase
+from flock.core.mcp.flock_mcp_server import FlockMCPServer
 from flock.core.mcp.mcp_config import (
-    FlockMCPCachingConfigurationBase,
-    FlockMCPCallbackConfigurationBase,
-    FlockMCPFeatureConfigurationBase,
+    FlockMCPCachingConfiguration,
+    FlockMCPCallbackConfiguration,
+    FlockMCPFeatureConfiguration,
 )
 from flock.core.mcp.types.types import (
     FlockListRootsMCPCallback,
@@ -207,7 +207,7 @@ class FlockFactory:
         tool_result_cache_ttl=100,
         description: str | Callable[..., str] | None = None,
         alert_latency_threshold_ms: int = 30000,
-    ) -> FlockMCPServerBase:
+    ) -> FlockMCPServer:
         """Create a default MCP Server with common modules.
 
         Allows for creating one of the three default-implementations provided
@@ -246,19 +246,19 @@ class FlockFactory:
                     continue  # ignore
 
         # build generic configs
-        feature_config = FlockMCPFeatureConfigurationBase(
+        feature_config = FlockMCPFeatureConfiguration(
             roots_enabled=enable_roots_feature,
             tools_enabled=enable_tools_feature,
             prompts_enabled=enable_prompts_feature,
             sampling_enabled=enable_sampling_feature,
         )
-        callback_config = FlockMCPCallbackConfigurationBase(
+        callback_config = FlockMCPCallbackConfiguration(
             sampling_callback=sampling_callback,
             list_roots_callback=list_roots_callback,
             logging_callback=logging_callback,
             message_handler=message_handler,
         )
-        caching_config = FlockMCPCachingConfigurationBase(
+        caching_config = FlockMCPCachingConfiguration(
             tool_cache_max_size=tool_cache_size,
             tool_cache_max_ttl=tool_cache_ttl,
             resource_contents_cache_max_size=resource_contents_cache_size,
@@ -388,12 +388,12 @@ class FlockFactory:
     @staticmethod
     def create_default_agent(
         name: str,
-        description: str | Callable[..., str] | None = None,
+        description: DynamicStr| None = None,
         model: str | Callable[..., str] | None = None,
-        input: SignatureType = None,
-        output: SignatureType = None,
+        input: DynamicStr = None,
+        output: DynamicStr = None,
         tools: list[Callable[..., Any] | Any] | None = None,
-        servers: list[str | FlockMCPServerBase] | None = None,
+        servers: list[str | FlockMCPServer] | None = None,
         use_cache: bool = True,
         enable_rich_tables: bool = False,
         output_theme: OutputTheme = OutputTheme.abernathy,
@@ -408,6 +408,7 @@ class FlockFactory:
         write_to_file: bool = False,
         stream: bool = False,
         include_thought_process: bool = False,
+        next_agent: DynamicStr | None = None,
         temporal_activity_config: TemporalActivityConfig | None = None,
     ) -> FlockAgent:
         """Creates a default FlockAgent using unified component architecture.
@@ -479,7 +480,7 @@ class FlockFactory:
             components=[evaluator, output_component, metrics_component],
             config=FlockAgentConfig(write_to_file=write_to_file,
                                     wait_for_input=wait_for_input),
-            next_agent=None,  # No next agent by default
+            next_agent=next_agent,
             temporal_activity_config=temporal_activity_config,
         )
 
@@ -489,14 +490,15 @@ class FlockFactory:
     def create_scheduled_agent(
         name: str,
         schedule_expression: str,  # e.g., "every 1h", "0 0 * * *"
-        description: str | Callable[..., str] | None = None,
-        model: str | Callable[..., str] | None = None,
-        output: SignatureType = None,  # Input might be implicit or none
+        description: DynamicStr | None = None,
+        model: str | None = None,
+        output: DynamicStr | None = None,  # Input might be implicit or none
         tools: list[Callable[..., Any] | Any] | None = None,
-        servers: list[str | FlockMCPServerBase] | None = None,
+        servers: list[str | FlockMCPServer] | None = None,
         use_cache: bool = False,  # Whether to cache results
         temperature: float = 0.7,  # Temperature for model responses
         # ... other common agent params from create_default_agent ...
+        next_agent: DynamicStr | None = None,
         temporal_activity_config: TemporalActivityConfig
         | None = None,  # If you want scheduled tasks to be Temporal activities
         **kwargs,  # Forward other standard agent params
@@ -522,6 +524,7 @@ class FlockFactory:
                 temporal_activity_config=temporal_activity_config,
                 use_cache=use_cache,
                 temperature=temperature,
+                next_agent=next_agent,
                 **kwargs,
             )
         )

flock/core/flock_server_manager.py

@@ -6,13 +6,13 @@ from contextlib import AsyncExitStack
 from anyio import Lock
 from pydantic import BaseModel, ConfigDict, Field
 
-from flock.core.mcp.flock_mcp_server import FlockMCPServerBase
+from flock.core.mcp.flock_mcp_server import FlockMCPServer
 
 
 class FlockServerManager(BaseModel):
     """Async-context-manager to start/stop a set of Flock MCP servers."""
 
-    servers: list[FlockMCPServerBase] | None = Field(
+    servers: list[FlockMCPServer] | None = Field(
         ..., exclude=True, description="The servers to manage."
     )
 
@@ -32,7 +32,7 @@ class FlockServerManager(BaseModel):
 
     def __init__(
         self,
-        servers: list[FlockMCPServerBase] | None = None,
+        servers: list[FlockMCPServer] | None = None,
         stack: AsyncExitStack | None = None,
         lock: asyncio.Lock | None = None,
     ) -> None:
@@ -43,7 +43,7 @@ class FlockServerManager(BaseModel):
             lock=lock,
         )
 
-    def add_server_sync(self, server: FlockMCPServerBase) -> None:
+    def add_server_sync(self, server: FlockMCPServer) -> None:
         """Add a server to be managed by the ServerManager.
 
         Note:
@@ -57,7 +57,7 @@ class FlockServerManager(BaseModel):
 
         self.servers.append(server)
 
-    def remove_server_sync(self, server: FlockMCPServerBase) -> None:
+    def remove_server_sync(self, server: FlockMCPServer) -> None:
         """Remove a server from the list of managed servers.
 
         Note:
@@ -71,7 +71,7 @@ class FlockServerManager(BaseModel):
 
     # -- For future use: Allow adding and removal of servers during runtime ---
     async def add_server_during_runtime(
-        self, server: FlockMCPServerBase
+        self, server: FlockMCPServer
     ) -> None:
         """Add a server to the manager and, if already running, start it immediately."""
         if self.lock is None:
@@ -88,13 +88,13 @@ class FlockServerManager(BaseModel):
             await self.stack.enter_async_context(server)
 
     async def remove_server_during_runtime(
-        self, server: FlockMCPServerBase
+        self, server: FlockMCPServer
     ) -> None:
         """Tear down and remove a server from the manager at runtime."""
         if self.lock is None:
             self.lock = asyncio.Lock()
 
-        retrieved_server: FlockMCPServerBase | None = None
+        retrieved_server: FlockMCPServer | None = None
 
         async with self.lock:
             if not self.servers or server not in self.servers:

flock/core/mcp/flock_mcp_server.py

@@ -17,9 +17,9 @@ from pydantic import (
 
 from flock.core.component.agent_component_base import AgentComponent
 from flock.core.logging.logging import get_logger
-from flock.core.mcp.flock_mcp_tool_base import FlockMCPToolBase
-from flock.core.mcp.mcp_client_manager import FlockMCPClientManagerBase
-from flock.core.mcp.mcp_config import FlockMCPConfigurationBase
+from flock.core.mcp.flock_mcp_tool import FlockMCPTool
+from flock.core.mcp.mcp_client_manager import FlockMCPClientManager
+from flock.core.mcp.mcp_config import FlockMCPConfiguration
 from flock.core.serialization.serializable import Serializable
 from flock.core.serialization.serialization_utils import (
     deserialize_component,
@@ -28,7 +28,7 @@ from flock.core.serialization.serialization_utils import (
 
 logger = get_logger("mcp.server")
 tracer = trace.get_tracer(__name__)
-T = TypeVar("T", bound="FlockMCPServerBase")
+T = TypeVar("T", bound="FlockMCPServer")
 
 LoggingLevel = Literal[
     "debug",
@@ -42,7 +42,7 @@ LoggingLevel = Literal[
 ]
 
 
-class FlockMCPServerBase(BaseModel, Serializable, ABC):
+class FlockMCPServer(BaseModel, Serializable, ABC):
     """Base class for all Flock MCP Server Types.
 
     Servers serve as an abstraction-layer between the underlying MCPClientSession
@@ -64,7 +64,7 @@ class FlockMCPServerBase(BaseModel, Serializable, ABC):
     2. Using FlockMCPServerConfig.with_fields() to create a config class.
     """
 
-    config: FlockMCPConfigurationBase = Field(
+    config: FlockMCPConfiguration = Field(
         ..., description="Config for clients connecting to the server."
     )
 
@@ -82,7 +82,7 @@ class FlockMCPServerBase(BaseModel, Serializable, ABC):
     # --- Underlying ConnectionManager ---
     # (Manages a pool of ClientConnections and does the actual talking to the MCP Server)
     # (Excluded from Serialization)
-    client_manager: FlockMCPClientManagerBase | None = Field(
+    client_manager: FlockMCPClientManager | None = Field(
         default=None,
         exclude=True,
         description="Underlying Connection Manager. Handles the actual underlying connections to the server.",
@@ -134,7 +134,7 @@ class FlockMCPServerBase(BaseModel, Serializable, ABC):
         return [c for c in self.components.values() if c.config.enabled]
 
     @abstractmethod
-    async def initialize(self) -> FlockMCPClientManagerBase:
+    async def initialize(self) -> FlockMCPClientManager:
         """Called when initializing the server."""
         pass
 
@@ -211,7 +211,7 @@ class FlockMCPServerBase(BaseModel, Serializable, ABC):
                         additional_params=additional_params
                     )
                     result: list[
-                        FlockMCPToolBase
+                        FlockMCPTool
                     ] = await self.client_manager.get_tools(
                         agent_id=agent_id,
                         run_id=run_id,
@@ -385,7 +385,7 @@ class FlockMCPServerBase(BaseModel, Serializable, ABC):
                 span.record_exception(module_error)
 
     # --- Async Methods ---
-    async def __aenter__(self) -> "FlockMCPServerBase":
+    async def __aenter__(self) -> "FlockMCPServer":
         """Enter the asynchronous context for the server."""
         # Spin up the client-manager
         with tracer.start_as_current_span("server.__aenter__") as span:
@@ -614,7 +614,7 @@ class FlockMCPServerBase(BaseModel, Serializable, ABC):
                 config_cls = config_field.annotation
             except (AttributeError, KeyError):
                 # fallback if Pydantic v1 or missing
-                config_cls = FlockMCPConfigurationBase
+                config_cls = FlockMCPConfiguration
             config_object = config_cls.from_dict(config_data)
             data["config"] = config_object
 

flock/core/mcp/{flock_mcp_tool_base.py → flock_mcp_tool.py}

@@ -13,7 +13,7 @@ from flock.core.logging.logging import get_logger
 logger = get_logger("mcp.tool")
 tracer = trace.get_tracer(__name__)
 
-T = TypeVar("T", bound="FlockMCPToolBase")
+T = TypeVar("T", bound="FlockMCPTool")
 
 TYPE_MAPPING = {
     "string": str,
@@ -25,7 +25,7 @@ TYPE_MAPPING = {
 }
 
 
-class FlockMCPToolBase(BaseModel):
+class FlockMCPTool(BaseModel):
     """Base Class for MCP Tools for Flock."""
 
     name: str = Field(..., description="Name of the tool")

flock/core/mcp/mcp_client.py

@@ -37,8 +37,8 @@ from pydantic import (
 )
 
 from flock.core.logging.logging import get_logger
-from flock.core.mcp.flock_mcp_tool_base import FlockMCPToolBase
-from flock.core.mcp.mcp_config import FlockMCPConfigurationBase
+from flock.core.mcp.flock_mcp_tool import FlockMCPTool
+from flock.core.mcp.mcp_config import FlockMCPConfiguration
 from flock.core.mcp.types.factories import (
     default_flock_mcp_list_roots_callback_factory,
     default_flock_mcp_logging_callback_factory,
@@ -61,7 +61,7 @@ tracer = trace.get_tracer(__name__)
 GetSessionIdCallback = Callable[[], str | None]
 
 
-class FlockMCPClientBase(BaseModel, ABC):
+class FlockMCPClient(BaseModel, ABC):
     """Wrapper for mcp ClientSession.
 
     Class will attempt to re-establish connection if possible.
@@ -71,7 +71,7 @@ class FlockMCPClientBase(BaseModel, ABC):
     """
 
     # --- Properties ---
-    config: FlockMCPConfigurationBase = Field(
+    config: FlockMCPConfiguration = Field(
         ..., description="The config for this client instance."
     )
 
@@ -225,7 +225,7 @@ class FlockMCPClientBase(BaseModel, ABC):
 
     def __init__(
         self,
-        config: FlockMCPConfigurationBase,
+        config: FlockMCPConfiguration,
         lock: Lock | None = None,
         tool_cache: TTLCache | None = None,
         tool_result_cache: TTLCache | None = None,
@@ -374,24 +374,24 @@ class FlockMCPClientBase(BaseModel, ABC):
         self,
         agent_id: str,
         run_id: str,
-    ) -> list[FlockMCPToolBase]:
+    ) -> list[FlockMCPTool]:
         """Gets a list of available tools from the server."""
 
         @cached(cache=self.tool_cache, key=cache_key_generator)
         async def _get_tools_cached(
             agent_id: str,
             run_id: str,
-        ) -> list[FlockMCPToolBase]:
+        ) -> list[FlockMCPTool]:
             if not self.config.feature_config.tools_enabled:
                 return []
 
-            async def _get_tools_internal() -> list[FlockMCPToolBase]:
+            async def _get_tools_internal() -> list[FlockMCPTool]:
                 # TODO: Crash
                 response: ListToolsResult = await self.session.list_tools()
                 flock_tools = []
 
                 for tool in response.tools:
-                    converted_tool = FlockMCPToolBase.from_mcp_tool(
+                    converted_tool = FlockMCPTool.from_mcp_tool(
                         tool,
                         agent_id=agent_id,
                         run_id=run_id,

flock/core/mcp/mcp_client_manager.py

@@ -13,22 +13,22 @@ from pydantic import (
 )
 
 from flock.core.logging.logging import get_logger
-from flock.core.mcp.flock_mcp_tool_base import FlockMCPToolBase
+from flock.core.mcp.flock_mcp_tool import FlockMCPTool
 from flock.core.mcp.mcp_client import (
-    FlockMCPClientBase,
+    FlockMCPClient,
 )
-from flock.core.mcp.mcp_config import FlockMCPConfigurationBase
+from flock.core.mcp.mcp_config import FlockMCPConfiguration
 
 logger = get_logger("mcp.client_manager")
 tracer = trace.get_tracer(__name__)
 
-TClient = TypeVar("TClient", bound="FlockMCPClientBase")
+TClient = TypeVar("TClient", bound="FlockMCPClient")
 
 
-class FlockMCPClientManagerBase(BaseModel, ABC, Generic[TClient]):
+class FlockMCPClientManager(BaseModel, ABC, Generic[TClient]):
     """Handles a Pool of MCPClients of type TClient."""
 
-    client_config: FlockMCPConfigurationBase = Field(
+    client_config: FlockMCPConfiguration = Field(
         ..., description="Configuration for clients."
     )
 
@@ -38,7 +38,7 @@ class FlockMCPClientManagerBase(BaseModel, ABC, Generic[TClient]):
         exclude=True,
     )
 
-    clients: dict[str, dict[str, FlockMCPClientBase]] = Field(
+    clients: dict[str, dict[str, FlockMCPClient]] = Field(
         default_factory=dict,
         exclude=True,
         description="Internal Store for the clients.",
@@ -154,7 +154,7 @@ class FlockMCPClientManagerBase(BaseModel, ABC, Generic[TClient]):
         agent_id: str,
         run_id: str,
         additional_params: dict[str, Any] | None = None,
-    ) -> list[FlockMCPToolBase]:
+    ) -> list[FlockMCPTool]:
         """Retrieves a list of tools for the agents to act on."""
         with tracer.start_as_current_span("client_manager.get_tools") as span:
             span.set_attribute("agent_id", agent_id)
@@ -165,7 +165,7 @@ class FlockMCPClientManagerBase(BaseModel, ABC, Generic[TClient]):
                     run_id=run_id,
                     additional_params=additional_params,
                 )
-                tools: list[FlockMCPToolBase] = await client.get_tools(
+                tools: list[FlockMCPTool] = await client.get_tools(
                     agent_id=agent_id, run_id=run_id
                 )
                 return tools

flock/core/mcp/mcp_config.py

@@ -36,14 +36,14 @@ LoggingLevel = Literal[
 ]
 
 
-A = TypeVar("A", bound="FlockMCPCallbackConfigurationBase")
-B = TypeVar("B", bound="FlockMCPConnectionConfigurationBase")
-C = TypeVar("C", bound="FlockMCPConfigurationBase")
-D = TypeVar("D", bound="FlockMCPCachingConfigurationBase")
-E = TypeVar("E", bound="FlockMCPFeatureConfigurationBase")
+A = TypeVar("A", bound="FlockMCPCallbackConfiguration")
+B = TypeVar("B", bound="FlockMCPConnectionConfiguration")
+C = TypeVar("C", bound="FlockMCPConfiguration")
+D = TypeVar("D", bound="FlockMCPCachingConfiguration")
+E = TypeVar("E", bound="FlockMCPFeatureConfiguration")
 
 
-class FlockMCPCachingConfigurationBase(BaseModel, Serializable):
+class FlockMCPCachingConfiguration(BaseModel, Serializable):
     """Configuration for Caching in Clients."""
 
     tool_cache_max_size: float = Field(
@@ -110,7 +110,7 @@ class FlockMCPCachingConfigurationBase(BaseModel, Serializable):
         )
 
 
-class FlockMCPCallbackConfigurationBase(BaseModel, Serializable):
+class FlockMCPCallbackConfiguration(BaseModel, Serializable):
     """Base Configuration Class for Callbacks for Clients."""
 
     sampling_callback: FlockSamplingMCPCallback | None = Field(
@@ -188,7 +188,7 @@ class FlockMCPCallbackConfigurationBase(BaseModel, Serializable):
         )
 
 
-class FlockMCPConnectionConfigurationBase(BaseModel, Serializable):
+class FlockMCPConnectionConfiguration(BaseModel, Serializable):
     """Base Configuration Class for Connection Parameters for a client."""
 
     max_retries: int = Field(
@@ -298,7 +298,7 @@ class FlockMCPConnectionConfigurationBase(BaseModel, Serializable):
         )
 
 
-class FlockMCPFeatureConfigurationBase(BaseModel, Serializable):
+class FlockMCPFeatureConfiguration(BaseModel, Serializable):
     """Base Configuration Class for switching MCP Features on and off."""
 
     roots_enabled: bool = Field(
@@ -346,7 +346,7 @@ class FlockMCPFeatureConfigurationBase(BaseModel, Serializable):
         )
 
 
-class FlockMCPConfigurationBase(BaseModel, Serializable):
+class FlockMCPConfiguration(BaseModel, Serializable):
     """Base Configuration Class for MCP Clients.
 
     Each Client should implement their own config
@@ -357,22 +357,22 @@ class FlockMCPConfigurationBase(BaseModel, Serializable):
         ..., description="Name of the server the client connects to."
     )
 
-    connection_config: FlockMCPConnectionConfigurationBase = Field(
+    connection_config: FlockMCPConnectionConfiguration = Field(
         ..., description="MCP Connection Configuration for a client."
     )
 
-    caching_config: FlockMCPCachingConfigurationBase = Field(
-        default_factory=FlockMCPCachingConfigurationBase,
+    caching_config: FlockMCPCachingConfiguration = Field(
+        default_factory=FlockMCPCachingConfiguration,
         description="Configuration for the internal caches of the client.",
     )
 
-    callback_config: FlockMCPCallbackConfigurationBase = Field(
-        default_factory=FlockMCPCallbackConfigurationBase,
+    callback_config: FlockMCPCallbackConfiguration = Field(
+        default_factory=FlockMCPCallbackConfiguration,
         description="Callback configuration for the client.",
     )
 
-    feature_config: FlockMCPFeatureConfigurationBase = Field(
-        default_factory=FlockMCPFeatureConfigurationBase,
+    feature_config: FlockMCPFeatureConfiguration = Field(
+        default_factory=FlockMCPFeatureConfiguration,
         description="Feature configuration for the client.",
     )
 
@@ -425,7 +425,7 @@ class FlockMCPConfigurationBase(BaseModel, Serializable):
                 config_cls = config_field.annotation
             except (AttributeError, KeyError):
                 # fallback
-                config_cls = FlockMCPConnectionConfigurationBase
+                config_cls = FlockMCPConnectionConfiguration
             instance_data["connection_config"] = config_cls.from_dict(connection_config)
         else:
             raise ValueError(f"connection_config MUST be specified for '{data.get('name', 'unknown_server')}")
@@ -436,10 +436,10 @@ class FlockMCPConfigurationBase(BaseModel, Serializable):
                 config_cls = config_field.annotation
             except (AttributeError, KeyError):
                 # fallback
-                config_cls = FlockMCPCachingConfigurationBase
+                config_cls = FlockMCPCachingConfiguration
             instance_data["caching_config"] = config_cls.from_dict(caching_config)
         else:
-            instance_data["caching_config"] = FlockMCPCachingConfigurationBase()
+            instance_data["caching_config"] = FlockMCPCachingConfiguration()
 
         if feature_config:
             try:
@@ -447,10 +447,10 @@ class FlockMCPConfigurationBase(BaseModel, Serializable):
                 config_cls = config_field.annotation
             except (AttributeError, KeyError):
                 # fallback
-                config_cls = FlockMCPFeatureConfigurationBase
+                config_cls = FlockMCPFeatureConfiguration
             instance_data["feature_config"] = config_cls.from_dict(feature_config)
         else:
-            instance_data["feature_config"] = FlockMCPFeatureConfigurationBase()
+            instance_data["feature_config"] = FlockMCPFeatureConfiguration()
 
         if callback_config:
             try:
@@ -458,10 +458,10 @@ class FlockMCPConfigurationBase(BaseModel, Serializable):
                 config_cls = config_field.annotation
             except (AttributeError, KeyError):
                 # fallback
-                config_cls = FlockMCPCallbackConfigurationBase
+                config_cls = FlockMCPCallbackConfiguration
             instance_data["callback_config"] = config_cls.from_dict(callback_config)
         else:
-            instance_data["callback_config"] = FlockMCPCallbackConfigurationBase()
+            instance_data["callback_config"] = FlockMCPCallbackConfiguration()
 
         return cls(**{k: v for k, v in instance_data.items()})
 

flock/core/orchestration/flock_execution.py

@@ -95,7 +95,7 @@ class FlockExecution:
         # Import here to allow forward reference resolution
         from flock.core.flock_agent import FlockAgent as ConcreteFlockAgent
         from flock.core.mcp.flock_mcp_server import (
-            FlockMCPServerBase as ConcreteFlockServer,
+            FlockMCPServer as ConcreteFlockServer,
         )
 
         with tracer.start_as_current_span("flock.run_async") as span:
@@ -131,7 +131,7 @@ class FlockExecution:
             span.set_attribute("input", str(run_input))
             span.set_attribute("run_id", effective_run_id)
             span.set_attribute("enable_temporal", self.flock.enable_temporal)
-            
+
             logger.info(
                 f"Initiating Flock run '{self.flock.name}'. Start Agent: '{start_agent_name}'. Temporal: {self.flock.enable_temporal}."
             )
@@ -169,7 +169,7 @@ class FlockExecution:
                 logger.error(f"Flock run '{self.flock.name}' failed: {e}", exc_info=True)
                 span.record_exception(e)
                 span.set_status(trace.Status(trace.StatusCode.ERROR, str(e)))
-                
+
                 # Return a consistent error structure
                 error_output = {
                     "error": str(e),

flock/core/orchestration/flock_initialization.py

@@ -3,18 +3,18 @@
 
 import os
 import uuid
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 from opentelemetry.baggage import get_baggage, set_baggage
 
-from flock.core.registry import get_registry
 from flock.core.logging.logging import get_logger
+from flock.core.registry import get_registry
 from flock.core.util.cli_helper import init_console
 
 if TYPE_CHECKING:
     from flock.core.flock import Flock
     from flock.core.flock_agent import FlockAgent
-    from flock.core.mcp.flock_mcp_server import FlockMCPServerBase
+    from flock.core.mcp.flock_mcp_server import FlockMCPServer
 
 logger = get_logger("flock.initialization")
 
@@ -28,7 +28,7 @@ class FlockInitialization:
     def setup(
         self,
         agents: list["FlockAgent"] | None = None,
-        servers: list["FlockMCPServerBase"] | None = None,
+        servers: list["FlockMCPServer"] | None = None,
     ) -> None:
         """Handle all initialization side effects and setup."""
         # Register passed servers first (agents may depend on them)
@@ -64,10 +64,10 @@ class FlockInitialization:
             enable_temporal=self.flock.enable_temporal,
         )
 
-    def _register_servers(self, servers: list["FlockMCPServerBase"]) -> None:
+    def _register_servers(self, servers: list["FlockMCPServer"]) -> None:
         """Register servers with the Flock instance."""
         from flock.core.mcp.flock_mcp_server import (
-            FlockMCPServerBase as ConcreteFlockMCPServer,
+            FlockMCPServer as ConcreteFlockMCPServer,
         )
 
         for server in servers:

flock/core/orchestration/flock_server_manager.py

@@ -3,12 +3,14 @@
 
 from typing import TYPE_CHECKING
 
-from flock.core.flock_server_manager import FlockServerManager as InternalServerManager
+from flock.core.flock_server_manager import (
+    FlockServerManager as InternalServerManager,
+)
 from flock.core.logging.logging import get_logger
 
 if TYPE_CHECKING:
     from flock.core.flock import Flock
-    from flock.core.mcp.flock_mcp_server import FlockMCPServerBase
+    from flock.core.mcp.flock_mcp_server import FlockMCPServer
 
 logger = get_logger("flock.server_manager")
 
@@ -21,12 +23,12 @@ class FlockServerManager:
         # Use the existing internal server manager
         self._internal_mgr = InternalServerManager()
 
-    def add_server(self, server: "FlockMCPServerBase") -> "FlockMCPServerBase":
+    def add_server(self, server: "FlockMCPServer") -> "FlockMCPServer":
         """Adds a server instance to this Flock configuration and registry as well as set it up to be managed by internal manager."""
-        from flock.core.registry import get_registry
         from flock.core.mcp.flock_mcp_server import (
-            FlockMCPServerBase as ConcreteFlockMCPServer,
+            FlockMCPServer as ConcreteFlockMCPServer,
         )
+        from flock.core.registry import get_registry
 
         registry = get_registry()
 
@@ -60,6 +62,6 @@ class FlockServerManager:
         return await self._internal_mgr.__aexit__(exc_type, exc_val, exc_tb)
 
     @property
-    def servers(self) -> dict[str, "FlockMCPServerBase"]:
+    def servers(self) -> dict[str, "FlockMCPServer"]:
         """Returns the dictionary of servers managed by this Flock instance."""
         return self.flock._servers