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

autobyteus/tools/__init__.py

@@ -15,23 +15,35 @@ from .tool_config import ToolConfig # Configuration data object, primarily for c
 from .tool_origin import ToolOrigin
 from .tool_category import ToolCategory
 
+# Tool Formatting Registration Support
+# Tool Formatting Registration Support
+from autobyteus.tools.usage.registries.tool_formatting_registry import ToolFormattingRegistry, register_tool_formatter
+from autobyteus.tools.usage.registries.tool_formatter_pair import ToolFormatterPair
+from autobyteus.tools.usage.formatters.base_formatter import BaseSchemaFormatter, BaseExampleFormatter
+
+
 logger = logging.getLogger(__name__)
 
 # --- Re-export specific tools for easier access ---
 
 # Functional tools (decorated functions are now instances)
-from .bash.bash_executor import bash_executor
-from .file.file_reader import file_reader
-from .file.file_writer import file_writer
-from .file.file_editor import file_edit
+from .file.read_file import read_file
+from .file.write_file import write_file
+from .file.patch_file import patch_file
+from .skill.load_skill import load_skill
+
+# Terminal tools (PTY-based stateful terminal)
+from .terminal.tools.run_bash import run_bash
+from .terminal.tools.start_background_process import start_background_process
+from .terminal.tools.get_process_output import get_process_output
+from .terminal.tools.stop_background_process import stop_background_process
 
 # General Class-based tools
 try:
-    from .google_search import GoogleSearch
+    from .search_tool import Search
 except ModuleNotFoundError as import_err:
-    logger.warning("GoogleSearch tool not available: %s", import_err)
-    GoogleSearch = None
-from .timer import Timer
+    logger.warning("Search tool not available: %s", import_err)
+    Search = None
 try:
     from .multimedia.image_tools import GenerateImageTool, EditImageTool
 except ModuleNotFoundError as import_err:
@@ -44,38 +56,19 @@ except ModuleNotFoundError as import_err:
     logger.warning("Media reader tool not available: %s", import_err)
     ReadMediaFile = None
 try:
-    from .download_media_tool import DownloadMediaTool
+    from autobyteus.multimedia.download_media_tool import DownloadMediaTool
 except ModuleNotFoundError as import_err:
     logger.warning("Download media tool not available: %s", import_err)
     DownloadMediaTool = None
 
-# Standalone Browser tools
-try:
-    from .browser.standalone.navigate_to import NavigateTo as StandaloneNavigateTo # Alias to avoid name clash
-    from .browser.standalone.webpage_reader import WebPageReader as StandaloneWebPageReader # Alias
-    from .browser.standalone.webpage_screenshot_taker import WebPageScreenshotTaker as StandaloneWebPageScreenshotTaker # Alias
-    from .browser.standalone.webpage_image_downloader import WebPageImageDownloader
-    from .browser.standalone.web_page_pdf_generator import WebPagePDFGenerator
-except ModuleNotFoundError as import_err:
-    logger.warning('Standalone browser tools not available: %s', import_err)
-    StandaloneNavigateTo = None
-    StandaloneWebPageReader = None
-    StandaloneWebPageScreenshotTaker = None
-    WebPageImageDownloader = None
-    WebPagePDFGenerator = None
-
-# Session-Aware Browser tools
+# Web tools
 try:
-    from .browser.session_aware.browser_session_aware_navigate_to import BrowserSessionAwareNavigateTo
-    from .browser.session_aware.browser_session_aware_web_element_trigger import BrowserSessionAwareWebElementTrigger
-    from .browser.session_aware.browser_session_aware_webpage_reader import BrowserSessionAwareWebPageReader
-    from .browser.session_aware.browser_session_aware_webpage_screenshot_taker import BrowserSessionAwareWebPageScreenshotTaker
+    from .web.read_url_tool import ReadUrl
 except ModuleNotFoundError as import_err:
-    logger.warning('Session-aware browser tools not available: %s', import_err)
-    BrowserSessionAwareNavigateTo = None
-    BrowserSessionAwareWebElementTrigger = None
-    BrowserSessionAwareWebPageReader = None
-    BrowserSessionAwareWebPageScreenshotTaker = None
+    logger.warning("ReadUrl tool not available: %s", import_err)
+    ReadUrl = None
+
+
 
 
 __all__ = [
@@ -90,29 +83,27 @@ __all__ = [
     "ToolCategory",
 
     # Re-exported functional tool instances
-    "bash_executor",
-    "file_reader",
-    "file_writer",
-    "file_edit",
+    "run_bash",
+    "start_background_process",
+    "get_process_output",
+    "stop_background_process",
+    "read_file",
+    "write_file",
+    "patch_file",
+    "load_skill",
 
     # Re-exported general class-based tools
-    "GoogleSearch",
-    "Timer",
+    "Search",
     "GenerateImageTool",
     "EditImageTool",
     "ReadMediaFile",
     "DownloadMediaTool",
 
-    # Re-exported Standalone Browser tools
-    "StandaloneNavigateTo",
-    "StandaloneWebPageReader",
-    "StandaloneWebPageScreenshotTaker",
-    "WebPageImageDownloader",
-    "WebPagePDFGenerator",
-
-    # Re-exported Session-Aware Browser tools
-    "BrowserSessionAwareNavigateTo",
-    "BrowserSessionAwareWebElementTrigger",
-    "BrowserSessionAwareWebPageReader",
-    "BrowserSessionAwareWebPageScreenshotTaker",
+    # Re-exported Web tools
+    "ReadUrl",
+
+    # Tool Formatting
+    "register_tool_formatter",
+    "BaseSchemaFormatter", 
+    "BaseExampleFormatter",
 ]

autobyteus/tools/base_tool.py

@@ -174,6 +174,13 @@ class BaseTool(ABC, EventEmitter, metaclass=ToolMeta):
     async def _execute(self, context: 'AgentContext', **kwargs) -> Any:
         raise NotImplementedError("Subclasses must implement the '_execute' method.")
 
+    async def cleanup(self) -> None:
+        """
+        Lifecycle hook invoked during agent shutdown to release resources held by the tool.
+        Default implementation is a no-op.
+        """
+        return None
+
     @classmethod
     def tool_usage(cls) -> str:
         logger.warning("BaseTool.tool_usage() is deprecated. Tool usage is now generated by formatters.")

autobyteus/tools/file/__init__.py

@@ -0,0 +1,9 @@
+from .patch_file import patch_file
+from .read_file import read_file
+from .write_file import write_file
+
+__all__ = [
+    "patch_file",
+    "read_file",
+    "write_file",
+]

autobyteus/tools/file/patch_file.py

@@ -0,0 +1,149 @@
+import os
+import logging
+from typing import TYPE_CHECKING, List
+
+from pydantic import Field
+
+from autobyteus.tools.functional_tool import tool
+from autobyteus.tools.tool_category import ToolCategory
+from autobyteus.utils.diff_utils import apply_unified_diff, PatchApplicationError
+
+if TYPE_CHECKING:
+    from autobyteus.agent.context import AgentContext
+
+logger = logging.getLogger(__name__)
+
+
+def _resolve_file_path(context: 'AgentContext', path: str) -> str:
+    """Resolves an absolute path for the given input, using the agent workspace when needed."""
+    if os.path.isabs(path):
+        final_path = path
+        logger.debug("patch_file: provided path '%s' is absolute.", path)
+    else:
+        if not context.workspace:
+            error_msg = ("Relative path '%s' provided, but no workspace is configured for agent '%s'. "
+                         "A workspace is required to resolve relative paths.")
+            logger.error(error_msg, path, context.agent_id)
+            raise ValueError(error_msg % (path, context.agent_id))
+        base_path = context.workspace.get_base_path()
+        if not base_path or not isinstance(base_path, str):
+            error_msg = ("Agent '%s' has a configured workspace, but it provided an invalid base path ('%s'). "
+                         "Cannot resolve relative path '%s'.")
+            logger.error(error_msg, context.agent_id, base_path, path)
+            raise ValueError(error_msg % (context.agent_id, base_path, path))
+        final_path = os.path.join(base_path, path)
+        logger.debug("patch_file: resolved relative path '%s' against workspace base '%s' to '%s'.", path, base_path, final_path)
+
+    normalized_path = os.path.normpath(final_path)
+    logger.debug("patch_file: normalized path to '%s'.", normalized_path)
+    return normalized_path
+
+
+@tool(name="patch_file", category=ToolCategory.FILE_SYSTEM)
+async def patch_file(
+    context: 'AgentContext',
+    path: str = Field(..., description="Path to the target file."),
+    patch: str = Field(
+        ...,
+        description=(
+            "Unified diff hunks describing edits to apply. "
+            "Example:\n"
+            "--- a/sample.txt\n"
+            "+++ b/sample.txt\n"
+            "@@ -1,2 +1,2 @@\n"
+            "-old line\n"
+            "+new line\n"
+            " unchanged line"
+        ),
+    ),
+) -> str:
+    """Applies a unified diff patch to update a text file without overwriting unrelated content.
+
+    Args:
+        path: Path to the target file. Relative paths are resolved against the agent workspace when available.
+        patch: Unified diff patch describing the edits to apply.
+
+    Raises:
+        FileNotFoundError: If the file does not exist.
+        PatchApplicationError: If the patch content cannot be applied cleanly.
+        IOError: If file reading or writing fails.
+    """
+    logger.debug("patch_file: requested patch for agent '%s' on path '%s'.", context.agent_id, path)
+    return_path = os.path.normpath(path)
+    
+    # Detailed logging for debugging patch content
+    logger.info("patch_file: ===== PATCH ARGUMENT DEBUG START =====")
+    logger.info("patch_file: raw patch repr: %r", patch)
+    logger.info("patch_file: patch length: %d chars", len(patch) if patch else 0)
+    patch_lines = patch.splitlines(keepends=True) if patch else []
+    for i, line in enumerate(patch_lines, 1):
+        prefix = line[0] if line else '<empty>'
+        logger.info("patch_file: line %d: prefix=%r content=%r", i, prefix, line)
+    logger.info("patch_file: ===== PATCH ARGUMENT DEBUG END =====")
+    
+    final_path = _resolve_file_path(context, path)
+
+    file_exists = os.path.exists(final_path)
+    if not file_exists:
+        raise FileNotFoundError(f"The file at resolved path {final_path} does not exist.")
+
+    try:
+        original_lines: List[str]
+        if file_exists:
+            with open(final_path, 'r', encoding='utf-8') as source:
+                original_lines = source.read().splitlines(keepends=True)
+        else:
+            original_lines = []
+
+        # Log original file content for comparison
+        logger.info("patch_file: ===== ORIGINAL FILE DEBUG START =====")
+        for i, line in enumerate(original_lines, 1):
+            logger.info("patch_file: original line %d: %r", i, line)
+        logger.info("patch_file: ===== ORIGINAL FILE DEBUG END =====")
+
+        patched_lines = None
+        patch_error = None
+        retry_strategies = [
+            (0, False),
+            (1, False),
+            (1, True),
+            (2, True),
+        ]
+        for fuzz_factor, ignore_whitespace in retry_strategies:
+            try:
+                patched_lines = apply_unified_diff(
+                    original_lines,
+                    patch,
+                    fuzz_factor=fuzz_factor,
+                    ignore_whitespace=ignore_whitespace,
+                )
+                if (fuzz_factor, ignore_whitespace) != (0, False):
+                    logger.info(
+                        "patch_file: applied with fuzz=%d ignore_whitespace=%s.",
+                        fuzz_factor,
+                        ignore_whitespace,
+                    )
+                break
+            except PatchApplicationError as patch_err:
+                patch_error = patch_err
+                logger.warning(
+                    "patch_file: patch failed with fuzz=%d ignore_whitespace=%s: %s",
+                    fuzz_factor,
+                    ignore_whitespace,
+                    patch_err,
+                )
+                continue
+        if patched_lines is None:
+            raise patch_error or PatchApplicationError("Patch could not be applied.")
+
+        with open(final_path, 'w', encoding='utf-8') as destination:
+            destination.writelines(patched_lines)
+
+        logger.info("patch_file: successfully applied patch to '%s'.", final_path)
+        return f"File patched successfully at {return_path}"
+    except PatchApplicationError as patch_err:
+        logger.error("patch_file: failed to apply patch to '%s': %s", final_path, patch_err, exc_info=True)
+        raise patch_err
+    except Exception as exc:  # pragma: no cover - general safeguard
+        logger.error("patch_file: unexpected error while patching '%s': %s", final_path, exc, exc_info=True)
+        raise IOError(f"Could not patch file at '{final_path}': {exc}")

autobyteus/tools/file/{file_reader.py → read_file.py}

@@ -1,6 +1,8 @@
 import os
 import logging
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
+
+from pydantic import Field
 
 from autobyteus.tools import tool
 from autobyteus.tools.tool_category import ToolCategory
@@ -10,16 +12,33 @@ if TYPE_CHECKING:
 
 logger = logging.getLogger(__name__)
 
-@tool(name="FileReader", category=ToolCategory.FILE_SYSTEM)
-async def file_reader(context: 'AgentContext', path: str) -> str:
+@tool(name="read_file", category=ToolCategory.FILE_SYSTEM)
+async def read_file(
+    context: 'AgentContext',
+    path: str,
+    start_line: Optional[int] = None,
+    end_line: Optional[int] = None,
+    include_line_numbers: bool = Field(
+        True,
+        description="If true, prefix each returned line with its line number (default).",
+    ),
+) -> str:
     """
-    Reads content from a specified file.
+    Reads content from a specified file. Supports optional 1-based inclusive line ranges via start_line/end_line.
+    Each returned line is prefixed with its line number when include_line_numbers is true.
     'path' is the path to the file. If relative, it must be resolved against a configured agent workspace.
-    Raises ValueError if a relative path is given without a valid workspace.
+    Raises ValueError if a relative path is given without a valid workspace or if line range arguments are invalid.
     Raises FileNotFoundError if the file does not exist.
     Raises IOError if file reading fails for other reasons.
     """
-    logger.debug(f"Functional FileReader tool for agent {context.agent_id}, initial path: {path}")
+    logger.debug(f"Functional read_file tool for agent {context.agent_id}, initial path: {path}")
+
+    if start_line is not None and start_line < 1:
+        raise ValueError(f"start_line must be >= 1 when provided; got {start_line}.")
+    if end_line is not None and end_line < 1:
+        raise ValueError(f"end_line must be >= 1 when provided; got {end_line}.")
+    if start_line is not None and end_line is not None and end_line < start_line:
+        raise ValueError(f"end_line ({end_line}) must be >= start_line ({start_line}).")
     
     final_path: str
     if os.path.isabs(path):
@@ -48,7 +67,19 @@ async def file_reader(context: 'AgentContext', path: str) -> str:
         
     try:
         with open(final_path, 'r', encoding='utf-8') as file:
-            content = file.read()
+            selected_lines = []
+            for line_no, line in enumerate(file, start=1):
+                if start_line is not None and line_no < start_line:
+                    continue
+                if end_line is not None and line_no > end_line:
+                    break
+                if include_line_numbers:
+                    line_text = line.rstrip('\n')
+                    line_suffix = '\n' if line.endswith('\n') else ''
+                    selected_lines.append(f"{line_no}: {line_text}{line_suffix}")
+                else:
+                    selected_lines.append(line)
+            content = ''.join(selected_lines)
         logger.info(f"File successfully read from '{final_path}' for agent '{context.agent_id}'.")
         return content
     except Exception as e:

autobyteus/tools/file/{file_writer.py → write_file.py}

@@ -10,8 +10,8 @@ if TYPE_CHECKING:
 
 logger = logging.getLogger(__name__)
 
-@tool(name="FileWriter", category=ToolCategory.FILE_SYSTEM)
-async def file_writer(context: 'AgentContext', path: str, content: str) -> str:
+@tool(name="write_file", category=ToolCategory.FILE_SYSTEM)
+async def write_file(context: 'AgentContext', path: str, content: str) -> str:
     """
     Creates or overwrites a file with specified content.
     'path' is the path where the file will be written. If relative, it must be resolved against a configured agent workspace.
@@ -20,11 +20,13 @@ async def file_writer(context: 'AgentContext', path: str, content: str) -> str:
     Raises ValueError if a relative path is given without a valid workspace.
     Raises IOError if file writing fails.
     """
-    logger.debug(f"Functional FileWriter tool for agent {context.agent_id}, initial path: {path}")
+    logger.debug(f"Functional write_file tool for agent {context.agent_id}, initial path: {path}")
     
     final_path: str
+    return_path: str
     if os.path.isabs(path):
         final_path = path
+        return_path = final_path
         logger.debug(f"Path '{path}' is absolute. Using it directly.")
     else:
         if not context.workspace:
@@ -39,6 +41,7 @@ async def file_writer(context: 'AgentContext', path: str, content: str) -> str:
             raise ValueError(error_msg)
             
         final_path = os.path.join(base_path, path)
+        return_path = os.path.normpath(path)
         logger.debug(f"Path '{path}' is relative. Resolved to '{final_path}' using workspace base path '{base_path}'.")
 
     try:
@@ -53,7 +56,7 @@ async def file_writer(context: 'AgentContext', path: str, content: str) -> str:
             file.write(content)
             
         logger.info(f"File successfully written to '{final_path}' for agent '{context.agent_id}'.")
-        return f"File created/updated at {final_path}"
+        return f"File created/updated at {return_path}"
     except Exception as e:
         logger.error(f"Error writing file to final path '{final_path}' for agent {context.agent_id}: {e}", exc_info=True)
         raise IOError(f"Could not write file at '{final_path}': {str(e)}")

autobyteus/tools/functional_tool.py

@@ -99,7 +99,7 @@ def _python_type_to_json_schema(py_type: Any) -> Optional[Dict[str, Any]]:
     if py_type is float: return {"type": "number"}
     if py_type is bool: return {"type": "boolean"}
     if py_type is dict: return {"type": "object"}
-    if py_type is list: return {"type": "array", "items": True} 
+    if py_type is list: return {"type": "array", "items": {}} # Use empty dict for 'any'
     
     origin_type = get_origin(py_type)
     if origin_type is Union:
@@ -111,8 +111,8 @@ def _python_type_to_json_schema(py_type: Any) -> Optional[Dict[str, Any]]:
         list_args = get_args(py_type)
         if list_args and len(list_args) == 1:
             item_schema = _python_type_to_json_schema(list_args[0])
-            return {"type": "array", "items": item_schema if item_schema else True}
-        return {"type": "array", "items": True} 
+            return {"type": "array", "items": item_schema if item_schema else {}}
+        return {"type": "array", "items": {}} # Use empty dict for 'any'
     if origin_type is Dict or origin_type is dict: return {"type": "object"}
     logger.debug(f"Could not map Python type {py_type} to a simple JSON schema for array items.")
     return None
@@ -141,18 +141,26 @@ def _get_parameter_type_from_hint(py_type: Any, param_name: str) -> Tuple[Parame
         list_args = get_args(actual_type) 
         if list_args and len(list_args) == 1: 
             array_item_js_schema = _python_type_to_json_schema(list_args[0])
-        if not array_item_js_schema: 
-            array_item_js_schema = True 
+        # FIX: For an untyped list, the item schema should be None, not True.
+        # An empty dict `{}` is a valid JSON schema for 'any'.
+        if array_item_js_schema is None:
+             array_item_js_schema = {}
         return param_type_enum, array_item_js_schema
 
     mapped_type = _TYPE_MAPPING.get(actual_type)
     if mapped_type:
-        item_schema_for_array = True if mapped_type == ParameterType.ARRAY else None
+        item_schema_for_array = {} if mapped_type == ParameterType.ARRAY else None
         return mapped_type, item_schema_for_array
 
     logger.warning(f"Unmapped type hint {py_type} (actual_type: {actual_type}) for param '{param_name}'. Defaulting to ParameterType.STRING.")
     return ParameterType.STRING, None
 
+
+try:
+    from pydantic.fields import FieldInfo
+except ImportError:
+    FieldInfo = None  # type: ignore
+
 def _parse_signature(sig: inspect.Signature, tool_name: str) -> Tuple[TypingList[str], bool, bool, ParameterSchema]:
     func_param_names = []
     expects_context = False
@@ -173,26 +181,51 @@ def _parse_signature(sig: inspect.Signature, tool_name: str) -> Tuple[TypingList
         param_type_enum, item_schema = _get_parameter_type_from_hint(param_type_hint, param_name)
         
         is_required = param_obj.default == inspect.Parameter.empty
-        if get_origin(param_type_hint) is Union and type(None) in get_args(param_type_hint):
-            is_required = False
-
+        default_val = param_obj.default if param_obj.default != inspect.Parameter.empty else None
+        
+        # --- Pydantic Field Extraction Logic ---
         param_desc = f"Parameter '{param_name}' for tool '{tool_name}'."
         param_name_lower = param_name.lower()
         if "path" in param_name_lower or "file" in param_name_lower or "dir" in param_name_lower or "folder" in param_name_lower:
-            param_desc += " This is expected to be a path."
+             param_desc += " This is expected to be a path."
+
+        if FieldInfo and isinstance(param_obj.default, FieldInfo):
+            field_info = param_obj.default
+            
+            # 1. Description
+            if field_info.description:
+                param_desc = field_info.description
             
+            # 2. Default Value & Requiredness
+            # If PydanticUndefined (or similar sentinel), it means required.
+            # Otherwise, use the default value from Field.
+            # Note: Pydantic v1 uses Undefined, v2 uses PydanticUndefined. 
+            # We check if it is the special undefined value via representation or direct check.
+            
+            # Simple heuristic for "Undefined" without importing the specific sentinel
+            if str(field_info.default) == "PydanticUndefined" or field_info.default == ...:
+                is_required = True
+                default_val = None
+            else:
+                is_required = False
+                default_val = field_info.default
+        
+        if get_origin(param_type_hint) is Union and type(None) in get_args(param_type_hint):
+            is_required = False
+
         schema_param = ParameterDefinition(
             name=param_name,
             param_type=param_type_enum,
             description=param_desc,
             required=is_required,
-            default_value=param_obj.default if param_obj.default != inspect.Parameter.empty else None,
+            default_value=default_val,
             array_item_schema=item_schema
         )
         generated_arg_schema.add_parameter(schema_param)
         
     return func_param_names, expects_context, expects_tool_state, generated_arg_schema
 
+
 # --- The refactored @tool decorator ---
 
 def tool(
@@ -215,6 +248,11 @@ def tool(
         
         final_arg_schema = argument_schema if argument_schema is not None else gen_arg_schema
 
+        def _current_description() -> str:
+            """Recompute the description from the latest docstring/override."""
+            latest_doc = inspect.getdoc(func)
+            return description or (latest_doc.split('\n\n')[0] if latest_doc else f"Functional tool: {tool_name}")
+
         def factory(inst_config: Optional[ToolConfig] = None) -> FunctionalTool:
             return FunctionalTool(
                 original_func=func,
@@ -232,12 +270,13 @@ def tool(
         tool_def = ToolDefinition(
             name=tool_name,
             description=tool_desc,
-            argument_schema=final_arg_schema,
-            config_schema=config_schema,
+            argument_schema_provider=lambda: final_arg_schema,
+            config_schema_provider=lambda: config_schema,
             custom_factory=factory,
             tool_class=None,
             origin=ToolOrigin.LOCAL,
-            category=category
+            category=category,
+            description_provider=_current_description
         )
         default_tool_registry.register_tool(tool_def)
         

autobyteus/tools/mcp/__init__.py

@@ -19,6 +19,7 @@ from .types import (
     BaseMcpConfig,
     StdioMcpServerConfig,
     StreamableHttpMcpServerConfig,
+    WebsocketMcpServerConfig,
     McpTransportType,
     McpServerInstanceKey,
 )
@@ -37,6 +38,7 @@ __all__ = [
     "BaseMcpConfig",
     "StdioMcpServerConfig",
     "StreamableHttpMcpServerConfig",
+    "WebsocketMcpServerConfig",
     "McpTransportType",
     "McpServerInstanceKey",
     # Services and Managers

autobyteus/tools/mcp/config_service.py

@@ -9,6 +9,7 @@ from .types import (
     BaseMcpConfig,
     StdioMcpServerConfig,
     StreamableHttpMcpServerConfig,
+    WebsocketMcpServerConfig,
     McpTransportType
 )
 from autobyteus.utils.singleton import SingletonMeta
@@ -49,7 +50,8 @@ class McpConfigService(metaclass=SingletonMeta):
 
         transport_specific_params_key_map = {
             McpTransportType.STDIO: "stdio_params",
-            McpTransportType.STREAMABLE_HTTP: "streamable_http_params"
+            McpTransportType.STREAMABLE_HTTP: "streamable_http_params",
+            McpTransportType.WEBSOCKET: "websocket_params",
         }
 
         if transport_type in transport_specific_params_key_map:
@@ -74,6 +76,8 @@ class McpConfigService(metaclass=SingletonMeta):
                 return StdioMcpServerConfig(**constructor_params)
             elif transport_type == McpTransportType.STREAMABLE_HTTP:
                 return StreamableHttpMcpServerConfig(**constructor_params)
+            elif transport_type == McpTransportType.WEBSOCKET:
+                return WebsocketMcpServerConfig(**constructor_params)
             else:
                 raise ValueError(f"Unsupported McpTransportType '{transport_type}' for server '{server_id}'.")
         except TypeError as e:

autobyteus/tools/mcp/server/__init__.py

@@ -5,6 +5,7 @@ This package contains the core abstractions for managing connections to remote M
 from .base_managed_mcp_server import BaseManagedMcpServer, ServerState
 from .stdio_managed_mcp_server import StdioManagedMcpServer
 from .http_managed_mcp_server import HttpManagedMcpServer
+from .websocket_managed_mcp_server import WebsocketManagedMcpServer
 from .proxy import McpServerProxy
 
 __all__ = [
@@ -12,5 +13,6 @@ __all__ = [
     "ServerState",
     "StdioManagedMcpServer",
     "HttpManagedMcpServer",
+    "WebsocketManagedMcpServer",
     "McpServerProxy",
 ]

autobyteus/tools/mcp/server/http_managed_mcp_server.py

@@ -23,7 +23,7 @@ class HttpManagedMcpServer(BaseManagedMcpServer):
         config = cast(StreamableHttpMcpServerConfig, self._config)
         
         logger.debug(f"Establishing HTTP connection for server '{self.server_id}' to URL: {config.url}")
-        read_stream, write_stream = await self._exit_stack.enter_async_context(
+        read_stream, write_stream, _ = await self._exit_stack.enter_async_context(
             streamablehttp_client(config.url, headers=config.headers)
         )
         session = await self._exit_stack.enter_async_context(ClientSession(read_stream, write_stream))