autobyteus 1.2.0__py3-none-any.whl → 1.2.3__py3-none-any.whl

autobyteus/tools/mcp/server/websocket_managed_mcp_server.py

@@ -0,0 +1,141 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+import ssl
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+from typing import cast
+
+import anyio
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from mcp import ClientSession, types as mcp_types
+from mcp.shared.message import SessionMessage
+from pydantic import ValidationError
+from websockets.asyncio.client import connect as ws_connect
+from websockets.typing import Subprotocol
+
+from .base_managed_mcp_server import BaseManagedMcpServer
+from ..types import WebsocketMcpServerConfig
+
+logger = logging.getLogger(__name__)
+
+INITIALIZE_TIMEOUT = 10  # seconds
+
+
+def _build_ssl_context(config: WebsocketMcpServerConfig) -> ssl.SSLContext | None:
+    """Builds an SSL context when the target URL uses wss://."""
+    if not config.url or not config.url.lower().startswith("wss://"):
+        return None
+
+    if config.verify_tls:
+        context = ssl.create_default_context()
+        if config.ca_file:
+            context.load_verify_locations(cafile=config.ca_file)
+    else:
+        context = ssl._create_unverified_context()
+
+    if config.client_cert:
+        context.load_cert_chain(certfile=config.client_cert, keyfile=config.client_key)
+
+    return context
+
+
+def _normalize_subprotocols(config: WebsocketMcpServerConfig) -> list[Subprotocol]:
+    """Ensures the MCP subprotocol is always negotiated."""
+    provided = [proto for proto in config.subprotocols if proto]
+    lowered = {proto.lower() for proto in provided}
+    if "mcp" not in lowered:
+        provided.append("mcp")
+    return [Subprotocol(proto) for proto in provided]
+
+
+@asynccontextmanager
+async def _websocket_transport(
+    config: WebsocketMcpServerConfig,
+) -> AsyncGenerator[
+    tuple[MemoryObjectReceiveStream[SessionMessage | Exception], MemoryObjectSendStream[SessionMessage]],
+    None,
+]:
+    read_stream_writer, read_stream = anyio.create_memory_object_stream[SessionMessage | Exception](0)
+    write_stream, write_stream_reader = anyio.create_memory_object_stream[SessionMessage](0)
+
+    headers = dict(config.headers)
+    if not headers:
+        headers = None
+
+    connect_kwargs = {
+        "origin": config.origin,
+        "subprotocols": _normalize_subprotocols(config),
+        "additional_headers": headers,
+        "open_timeout": config.open_timeout,
+        "ping_interval": config.ping_interval,
+        "ping_timeout": config.ping_timeout,
+        "ssl": _build_ssl_context(config),
+    }
+    # Remove None values so websockets uses library defaults
+    connect_kwargs = {key: value for key, value in connect_kwargs.items() if value is not None}
+
+    negotiated_protocols = [str(proto) for proto in connect_kwargs.get("subprotocols", [])]
+
+    logger.debug(
+        "Connecting to MCP WebSocket %s (subprotocols=%s, origin=%s)",
+        config.url,
+        negotiated_protocols,
+        config.origin,
+    )
+
+    async with ws_connect(config.url, **connect_kwargs) as websocket:
+
+        async def ws_reader():
+            async with read_stream_writer:
+                async for raw_message in websocket:
+                    payload = raw_message.decode("utf-8") if isinstance(raw_message, bytes) else raw_message
+                    try:
+                        message = mcp_types.JSONRPCMessage.model_validate_json(payload)
+                        await read_stream_writer.send(SessionMessage(message))
+                    except ValidationError as exc:
+                        await read_stream_writer.send(exc)
+
+        async def ws_writer():
+            async with write_stream_reader:
+                async for session_message in write_stream_reader:
+                    payload = session_message.message.model_dump_json(by_alias=True, exclude_none=True)
+                    await websocket.send(payload)
+
+        async with anyio.create_task_group() as tg:
+            tg.start_soon(ws_reader)
+            tg.start_soon(ws_writer)
+            try:
+                yield read_stream, write_stream
+            finally:
+                tg.cancel_scope.cancel()
+
+
+class WebsocketManagedMcpServer(BaseManagedMcpServer):
+    """Manages the lifecycle of a WebSocket-based MCP server connection."""
+
+    def __init__(self, config: WebsocketMcpServerConfig):
+        super().__init__(config)
+
+    async def _create_client_session(self) -> ClientSession:
+        config = cast(WebsocketMcpServerConfig, self._config)
+
+        read_stream, write_stream = await self._exit_stack.enter_async_context(
+            _websocket_transport(config)
+        )
+        session = await self._exit_stack.enter_async_context(ClientSession(read_stream, write_stream))
+
+        try:
+            await asyncio.wait_for(session.initialize(), timeout=INITIALIZE_TIMEOUT)
+        except asyncio.TimeoutError:
+            logger.error(
+                "Timeout occurred while initializing WebSocket session for server '%s'.",
+                self.server_id,
+            )
+            raise ConnectionError(
+                f"Server '{self.server_id}' failed to initialize within the timeout period."
+            ) from None
+
+        logger.debug("ClientSession established and initialized for WebSocket server '%s'.", self.server_id)
+        return session

autobyteus/tools/mcp/server_instance_manager.py

@@ -7,7 +7,12 @@ from contextlib import asynccontextmanager
 from autobyteus.utils.singleton import SingletonMeta
 from autobyteus.agent.context import AgentContextRegistry
 from .config_service import McpConfigService
-from .server import BaseManagedMcpServer, StdioManagedMcpServer, HttpManagedMcpServer
+from .server import (
+    BaseManagedMcpServer,
+    StdioManagedMcpServer,
+    HttpManagedMcpServer,
+    WebsocketManagedMcpServer,
+)
 from .types import McpTransportType, McpServerInstanceKey, BaseMcpConfig, StdioMcpServerConfig
 
 logger = logging.getLogger(__name__)
@@ -29,6 +34,8 @@ class McpServerInstanceManager(metaclass=SingletonMeta):
             return StdioManagedMcpServer(server_config)
         elif server_config.transport_type == McpTransportType.STREAMABLE_HTTP:
             return HttpManagedMcpServer(server_config)
+        elif server_config.transport_type == McpTransportType.WEBSOCKET:
+            return WebsocketManagedMcpServer(server_config)
         else:
             raise NotImplementedError(f"No ManagedMcpServer implementation for transport type '{server_config.transport_type}'.")
 

autobyteus/tools/mcp/tool.py

@@ -41,7 +41,7 @@ class GenericMcpTool(BaseTool):
         self.get_description = self.get_instance_description
         self.get_argument_schema = self.get_instance_argument_schema
         
-        logger.info(f"GenericMcpTool instance created for remote tool '{remote_tool_name}' on server '{self._server_id}'. "
+        logger.info(f"call_remote_mcp_tool instance created for remote tool '{remote_tool_name}' on server '{self._server_id}'. "
                     f"Registered in AutoByteUs as '{self._instance_name}'.")
 
     # --- Getters for instance-specific data ---
@@ -51,7 +51,7 @@ class GenericMcpTool(BaseTool):
 
     # --- Base class methods (class-level, not instance-level) ---
     @classmethod
-    def get_name(cls) -> str: return "GenericMcpTool"
+    def get_name(cls) -> str: return "call_remote_mcp_tool"
     @classmethod
     def get_description(cls) -> str: return "A generic wrapper for executing remote MCP tools."
     @classmethod
@@ -65,7 +65,7 @@ class GenericMcpTool(BaseTool):
         agent_id = context.agent_id
         tool_name_for_log = self.get_instance_name()
         
-        logger.info(f"GenericMcpTool '{tool_name_for_log}': Creating proxy for agent '{agent_id}' and server '{self._server_id}'.")
+        logger.info(f"call_remote_mcp_tool '{tool_name_for_log}': Creating proxy for agent '{agent_id}' and server '{self._server_id}'.")
         
         try:
             # The proxy is created on-demand for each execution.

autobyteus/tools/mcp/tool_registrar.py

@@ -59,6 +59,8 @@ class McpToolRegistrar(metaclass=SingletonMeta):
         if server_config.tool_name_prefix:
             registered_name = f"{server_config.tool_name_prefix.rstrip('_')}_{remote_tool.name}"
 
+        # Note: McpToolFactory is now somewhat redundant as it holds static info,
+        # but we keep it for consistency. It creates a GenericMcpTool which needs this static info.
         tool_factory = McpToolFactory(
             server_id=server_config.server_id,
             remote_tool_name=remote_tool.name,
@@ -70,12 +72,13 @@ class McpToolRegistrar(metaclass=SingletonMeta):
         return ToolDefinition(
             name=registered_name,
             description=actual_desc,
-            argument_schema=actual_arg_schema,
+            # Pass schema providers as lambdas to conform to the new constructor
+            argument_schema_provider=lambda: actual_arg_schema,
+            config_schema_provider=lambda: None,
             origin=ToolOrigin.MCP,
             category=server_config.server_id, # Use server_id as the category
             metadata={"mcp_server_id": server_config.server_id}, # Store origin in generic metadata
             custom_factory=tool_factory.create_tool,
-            config_schema=None,
             tool_class=None
         )
 

autobyteus/tools/mcp/types.py

@@ -10,6 +10,7 @@ class McpTransportType(str, Enum):
     """Enumeration of supported MCP transport types."""
     STDIO = "stdio"
     STREAMABLE_HTTP = "streamable_http"
+    WEBSOCKET = "websocket"
 
 @dataclass(frozen=True)
 class McpServerInstanceKey:
@@ -85,3 +86,63 @@ class StreamableHttpMcpServerConfig(BaseMcpConfig):
             raise ValueError(f"StreamableHttpMcpServerConfig '{self.server_id}' 'token' must be a string if provided.") 
         if not isinstance(self.headers, dict) or not all(isinstance(k, str) and isinstance(v, str) for k, v in self.headers.items()):
             raise ValueError(f"StreamableHttpMcpServerConfig '{self.server_id}' 'headers' must be a Dict[str, str].")
+
+@dataclass
+class WebsocketMcpServerConfig(BaseMcpConfig):
+    """Configuration parameters for an MCP server using a WebSocket transport."""
+
+    url: Optional[str] = None
+    headers: Dict[str, str] = field(default_factory=dict)
+    subprotocols: List[str] = field(default_factory=list)
+    origin: Optional[str] = None
+    open_timeout: Optional[float] = 10.0
+    ping_interval: Optional[float] = None
+    ping_timeout: Optional[float] = None
+    verify_tls: bool = True
+    ca_file: Optional[str] = None
+    client_cert: Optional[str] = None
+    client_key: Optional[str] = None
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.transport_type = McpTransportType.WEBSOCKET
+
+        if self.url is None or not isinstance(self.url, str) or not self.url.strip():
+            raise ValueError(f"WebsocketMcpServerConfig '{self.server_id}' 'url' must be a non-empty string.")
+
+        normalized_url = self.url.strip().lower()
+        if not (normalized_url.startswith("ws://") or normalized_url.startswith("wss://")):
+            raise ValueError(
+                f"WebsocketMcpServerConfig '{self.server_id}' 'url' must start with ws:// or wss://."
+            )
+
+        if not isinstance(self.headers, dict) or not all(isinstance(k, str) and isinstance(v, str) for k, v in self.headers.items()):
+            raise ValueError(f"WebsocketMcpServerConfig '{self.server_id}' 'headers' must be a Dict[str, str].")
+
+        if not isinstance(self.subprotocols, list) or not all(isinstance(item, str) for item in self.subprotocols):
+            raise ValueError(f"WebsocketMcpServerConfig '{self.server_id}' 'subprotocols' must be a list of strings.")
+
+        if self.origin is not None and not isinstance(self.origin, str):
+            raise ValueError(f"WebsocketMcpServerConfig '{self.server_id}' 'origin' must be a string if provided.")
+
+        for field_name in ("open_timeout", "ping_interval", "ping_timeout"):
+            value = getattr(self, field_name)
+            if value is not None and (not isinstance(value, (int, float)) or value <= 0):
+                raise ValueError(
+                    f"WebsocketMcpServerConfig '{self.server_id}' '{field_name}' must be a positive number when provided."
+                )
+
+        if not isinstance(self.verify_tls, bool):
+            raise ValueError(f"WebsocketMcpServerConfig '{self.server_id}' 'verify_tls' must be a boolean.")
+
+        for path_field in ("ca_file", "client_cert", "client_key"):
+            path_value = getattr(self, path_field)
+            if path_value is not None and not isinstance(path_value, str):
+                raise ValueError(
+                    f"WebsocketMcpServerConfig '{self.server_id}' '{path_field}' must be a string path when provided."
+                )
+
+        if self.client_key and not self.client_cert:
+            raise ValueError(
+                f"WebsocketMcpServerConfig '{self.server_id}' requires 'client_cert' when 'client_key' is provided."
+            )

autobyteus/tools/multimedia/__init__.py

@@ -1,10 +1,11 @@
 from .image_tools import GenerateImageTool, EditImageTool
 from .audio_tools import GenerateSpeechTool
 from .media_reader_tool import ReadMediaFile
-
+from .download_media_tool import DownloadMediaTool
 __all__ = [
     "GenerateImageTool",
     "EditImageTool",
     "GenerateSpeechTool",
     "ReadMediaFile",
+    "DownloadMediaTool",
 ]

autobyteus/tools/multimedia/audio_tools.py

@@ -1,14 +1,38 @@
 import os
 import logging
-from typing import Optional, List
+from typing import Optional, List, Any
+from pathlib import Path
 
 from autobyteus.tools.base_tool import BaseTool
 from autobyteus.utils.parameter_schema import ParameterSchema, ParameterDefinition, ParameterType
 from autobyteus.tools.tool_category import ToolCategory
 from autobyteus.multimedia.audio import audio_client_factory, AudioModel, AudioClientFactory
+from autobyteus.multimedia.audio.base_audio_client import BaseAudioClient
+from autobyteus.utils.download_utils import download_file_from_url
+from autobyteus.utils.file_utils import resolve_safe_path
 
 logger = logging.getLogger(__name__)
 
+def _get_workspace_root(context) -> str:
+    if not context.workspace:
+        error_msg = (
+            f"Relative path provided, but no workspace is configured for agent '{context.agent_id}'. "
+            "A workspace is required to resolve relative paths."
+        )
+        logger.error(error_msg)
+        raise ValueError(error_msg)
+
+    base_path = context.workspace.get_base_path()
+    if not base_path or not isinstance(base_path, str):
+        error_msg = (
+            f"Agent '{context.agent_id}' has a configured workspace, but it provided an invalid base path "
+            f"('{base_path}'). Cannot resolve relative paths."
+        )
+        logger.error(error_msg)
+        raise ValueError(error_msg)
+
+    return base_path
+
 
 def _get_configured_model_identifier(env_var: str, default_model: Optional[str] = None) -> str:
     """
@@ -54,22 +78,25 @@ def _build_dynamic_audio_schema(base_params: List[ParameterDefinition], model_en
 
 class GenerateSpeechTool(BaseTool):
     """
-
     An agent tool for generating speech from text using a Text-to-Speech (TTS) model.
     """
     CATEGORY = ToolCategory.MULTIMEDIA
     MODEL_ENV_VAR = "DEFAULT_SPEECH_GENERATION_MODEL"
     DEFAULT_MODEL = "gemini-2.5-flash-tts"
 
+    def __init__(self, config=None):
+        super().__init__(config)
+        self._client: Optional[BaseAudioClient] = None
+
     @classmethod
     def get_name(cls) -> str:
-        return "GenerateSpeech"
+        return "generate_speech"
 
     @classmethod
     def get_description(cls) -> str:
         return (
             "Generates spoken audio from text using the system's default Text-to-Speech (TTS) model. "
-            "Returns a list of local file paths to the generated audio files (.wav) upon success."
+            "Saves the generated audio file (.wav or .mp3) to the specified local file path and returns the path."
         )
 
     @classmethod
@@ -82,25 +109,51 @@ class GenerateSpeechTool(BaseTool):
                     "The text to be converted into spoken audio. For multi-speaker mode, you must format the prompt "
                     "with speaker labels that match the speakers defined in 'speaker_mapping'. "
                     "CRITICAL: Each speaker's dialogue MUST be on a new line. "
-                    "Example: 'Joe: Hello Jane.\\nJane: Hi Joe, how are you?'"
+                    "Example: 'Joe: Hello Jane.\nJane: Hi Joe, how are you?'"
+                ),
+                required=True
+            ),
+            ParameterDefinition(
+                name="output_file_path",
+                param_type=ParameterType.STRING,
+                description=(
+                    "Required. The local file path (relative to workspace) where the generated audio should be saved. "
+                    "Example: 'assets/audio/speech.wav'"
                 ),
                 required=True
             )
         ]
         return _build_dynamic_audio_schema(base_params, cls.MODEL_ENV_VAR, cls.DEFAULT_MODEL)
 
-    async def _execute(self, context, prompt: str, generation_config: Optional[dict] = None) -> List[str]:
+    async def _execute(
+        self,
+        context,
+        prompt: str,
+        output_file_path: str,
+        generation_config: Optional[dict] = None,
+    ) -> Any:
         model_identifier = _get_configured_model_identifier(self.MODEL_ENV_VAR, self.DEFAULT_MODEL)
-        logger.info(f"GenerateSpeechTool executing with configured model '{model_identifier}'.")
-        client = None
-        try:
-            client = audio_client_factory.create_audio_client(model_identifier=model_identifier)
-            response = await client.generate_speech(prompt=prompt, generation_config=generation_config)
-            
-            if not response.audio_urls:
-                raise ValueError("Speech generation failed to return any audio file paths.")
-            
-            return response.audio_urls
-        finally:
-            if client:
-                await client.cleanup()
+        logger.info(f"generate_speech executing with configured model '{model_identifier}'.")
+        if self._client is None:
+            self._client = audio_client_factory.create_audio_client(model_identifier=model_identifier)
+
+        response = await self._client.generate_speech(prompt=prompt, generation_config=generation_config)
+
+        if not response.audio_urls:
+            raise ValueError("Speech generation failed to return any audio file paths.")
+
+        first_url = response.audio_urls[0]
+
+        if not output_file_path:
+            raise ValueError("output_file_path is required but was not provided.")
+
+        # Save to File
+        resolved_path = resolve_safe_path(output_file_path, _get_workspace_root(context))
+        await download_file_from_url(first_url, resolved_path)
+
+        return {"file_path": str(resolved_path)}
+
+    async def cleanup(self) -> None:
+        if self._client:
+            await self._client.cleanup()
+            self._client = None

autobyteus/tools/{download_media_tool.py → multimedia/download_media_tool.py}

@@ -19,18 +19,19 @@ class DownloadMediaTool(BaseTool):
     """
     A unified tool to download any media file (e.g., image, PDF, audio) from a URL.
     """
-    CATEGORY = ToolCategory.WEB
+    CATEGORY = ToolCategory.MULTIMEDIA
 
     @classmethod
     def get_name(cls) -> str:
-        return "DownloadMedia"
+        return "download_media"
 
     @classmethod
     def get_description(cls) -> str:
         return (
-            "Downloads various media files (e.g., images like PNG/JPG, documents like PDF, audio like MP3/WAV) "
-            "from a direct URL and saves them locally. It intelligently determines the correct file extension "
-            "based on the content type. Returns the absolute path to the downloaded file."
+            "Download a media file (image/PDF/audio/etc.) from a direct URL and save it locally. "
+            "The tool picks the correct file extension from the HTTP Content-Type header (or falls back to the URL). "
+            "Files are saved to the agent workspace if you give a relative folder (preferred), or to your default "
+            "Downloads directory when no folder is provided. Returns the absolute path of the saved file."
         )
 
     @classmethod
@@ -62,8 +63,21 @@ class DownloadMediaTool(BaseTool):
             if folder:
                 # Security: prevent path traversal attacks.
                 if ".." in folder:
-                     raise ValueError("Security error: 'folder' path cannot contain '..'.")
-                destination_dir = os.path.abspath(folder)
+                    raise ValueError("Security error: 'folder' path cannot contain '..'.")
+                if not os.path.isabs(folder):
+                    workspace = context.workspace
+                    # Prefer workspace base path when available to keep downloads inside the agent's sandbox.
+                    if workspace and hasattr(workspace, "get_base_path") and callable(getattr(workspace, "get_base_path")):
+                        base_path = os.path.abspath(workspace.get_base_path())
+                        destination_dir = os.path.abspath(os.path.join(base_path, folder))
+                        # Ensure resolved path stays within workspace
+                        if os.path.commonpath([base_path]) != os.path.commonpath([base_path, destination_dir]):
+                            raise ValueError(f"Security error: 'folder' resolves outside workspace: {destination_dir}")
+                    else:
+                        # Fallback: resolve relative folder under the default download directory
+                        destination_dir = os.path.abspath(os.path.join(get_default_download_folder(), folder))
+                else:
+                    destination_dir = os.path.abspath(folder)
             else:
                 destination_dir = get_default_download_folder()