aixtools 0.1.10__py3-none-any.whl → 0.2.0__py3-none-any.whl

aixtools/mcp/client.py

@@ -1,14 +1,22 @@
 """MCP server utilities with caching and robust error handling."""
 
 import asyncio
-from typing import Any
+import logging
+from contextlib import asynccontextmanager
+from datetime import timedelta
+from typing import Any, AsyncGenerator
 
 import anyio
+import httpx
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 from cachebox import TTLCache
+from fastmcp.client.logging import LogMessage
 from mcp import types as mcp_types
+from mcp.client import streamable_http
 from mcp.shared.exceptions import McpError
+from mcp.shared.message import SessionMessage
 from pydantic_ai import RunContext, exceptions
-from pydantic_ai.mcp import MCPServerStreamableHTTP, ToolResult
+from pydantic_ai.mcp import MCPServerStdio, MCPServerStreamableHTTP, ToolResult
 from pydantic_ai.toolsets.abstract import ToolsetTool
 
 from aixtools.context import SessionIdTuple
@@ -16,11 +24,55 @@ from aixtools.logging.logging_config import get_logger
 
 MCP_TOOL_CACHE_TTL = 300  # 5 minutes
 DEFAULT_MCP_CONNECTION_TIMEOUT = 30
+DEFAULT_MCP_READ_TIMEOUT = float(60 * 5)  # 5 minutes
 CACHE_KEY = "TOOL_LIST"
 
 logger = get_logger(__name__)
 
 
+# Default log_handler for MCP clients
+LOGGING_LEVEL_MAP = logging.getLevelNamesMapping()
+
+
+async def default_mcp_log_handler(message: LogMessage):
+    """
+    Handles incoming logs from the MCP server and forwards them
+    to the standard Python logging system.
+    """
+    msg = message.data.get("msg")
+    extra = message.data.get("extra")
+
+    # Convert the MCP log level to a Python log level
+    level = LOGGING_LEVEL_MAP.get(message.level.upper(), logging.INFO)
+
+    # Log the message using the standard logging library
+    logger.log(level, msg, extra=extra)
+
+
+def get_mcp_client(
+    url: str | None = None,
+    command: str | None = None,
+    args: list[str] = None,
+    log_handler: callable = default_mcp_log_handler,  # type: ignore
+) -> MCPServerStreamableHTTP | MCPServerStdio:
+    """
+    Create an MCP client instance based on the provided URL or command.
+    By providing a log_handler, incoming logs from the MCP server can be shown, which improves debugging.
+
+    Args:
+        url (str | None): The URL of the MCP server.
+        command (str | None): The command to start a local MCP server (STDIO MCP).
+        args (list[str] | None): Additional arguments for the command (STDIO MCP).
+    """
+    if args is None:
+        args = []
+    if url:
+        return MCPServerStreamableHTTP(url="http://127.0.0.1:8089/mcp/", log_handler=log_handler)
+    if command:
+        return MCPServerStdio(command=command, args=args, log_handler=log_handler)
+    raise ValueError("Either url or command must be provided to create MCP client.")
+
+
 def get_mcp_headers(session_id_tuple: SessionIdTuple) -> dict[str, str] | None:
     """
     Generate headers for MCP server requests.
@@ -145,6 +197,23 @@ class CachedMCPServerStreamableHTTP(MCPServerStreamableHTTP):
                 logger.warning("MCP %s: %s exception %s: %s", self.url, func.__name__, type(exc), exc)
             return fallback(exc)
 
+    @property
+    def _transport_client(self):
+        """Override base transport client with wrapper logging and suppressing exceptions"""
+        return patched_streamablehttp_client
+
+    @asynccontextmanager
+    async def client_streams(self):
+        """Override base client_streams with wrapper logging and suppressing exceptions"""
+        try:
+            async with super().client_streams() as streams:  # pylint: disable=contextmanager-generator-missing-cleanup
+                try:
+                    yield streams
+                except Exception as exc:  # pylint: disable=broad-except
+                    logger.error("MCP %s: client_streams; %s: %s", self.url, type(exc).__name__, exc)
+        except Exception as exc:  # pylint: disable=broad-except
+            logger.error("MCP %s: client_streams: %s: %s", self.url, type(exc).__name__, exc)
+
     async def __aenter__(self):
         """Enter the context of the cached MCP server with complete cancellation isolation."""
         async with self._isolation_lock:
@@ -272,3 +341,80 @@ class CachedMCPServerStreamableHTTP(MCPServerStreamableHTTP):
             raise exceptions.ModelRetry(text)
 
         return content[0] if len(content) == 1 else content
+
+
+class PatchedStreamableHTTPTransport(streamable_http.StreamableHTTPTransport):
+    """Patched StreamableHTTPTransport with exception suppression for _handle_post_request."""
+
+    async def _handle_post_request(self, ctx: streamable_http.RequestContext) -> None:
+        """Patched _handle_post_request with proper error handling."""
+        try:
+            await super()._handle_post_request(ctx)
+        except Exception as exc:  # pylint: disable=broad-except
+            logger.error("MCP %s: _handle_post_request %s: %s", self.url, type(exc).__name__, exc)
+
+
+@asynccontextmanager
+async def patched_streamablehttp_client(  # noqa: PLR0913, pylint: disable=too-many-arguments,too-many-positional-arguments,too-many-locals
+    url: str,
+    headers: dict[str, str] | None = None,
+    timeout: float | timedelta = 30,
+    sse_read_timeout: float | timedelta = DEFAULT_MCP_READ_TIMEOUT,
+    terminate_on_close: bool = True,
+    httpx_client_factory: streamable_http.McpHttpClientFactory = streamable_http.create_mcp_http_client,
+    auth: httpx.Auth | None = None,
+) -> AsyncGenerator[
+    tuple[
+        MemoryObjectReceiveStream[SessionMessage | Exception],
+        MemoryObjectSendStream[SessionMessage],
+        streamable_http.GetSessionIdCallback,
+    ],
+    None,
+]:
+    """Patched version of `streamablehttp_client` with exception suppression."""
+    try:
+        transport = PatchedStreamableHTTPTransport(url, headers, timeout, sse_read_timeout, auth)
+
+        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)
+        async with anyio.create_task_group() as tg:
+            try:
+                async with httpx_client_factory(
+                    headers=transport.request_headers,
+                    timeout=httpx.Timeout(transport.timeout, read=transport.sse_read_timeout),
+                    auth=transport.auth,
+                ) as client:
+                    # Define callbacks that need access to tg
+                    def start_get_stream() -> None:
+                        tg.start_soon(transport.handle_get_stream, client, read_stream_writer)
+
+                    tg.start_soon(
+                        transport.post_writer,
+                        client,
+                        write_stream_reader,
+                        read_stream_writer,
+                        write_stream,
+                        start_get_stream,
+                        tg,
+                    )
+
+                    try:
+                        yield (
+                            read_stream,
+                            write_stream,
+                            transport.get_session_id,
+                        )
+                    except GeneratorExit:
+                        logger.warning("patched_streamablehttp_client: GeneratorExit caught, closing streams.")
+                    finally:
+                        if transport.session_id and terminate_on_close:
+                            await transport.terminate_session(client)
+                        tg.cancel_scope.cancel()
+            finally:
+                await read_stream_writer.aclose()
+                await write_stream.aclose()
+    except Exception as exc:  # pylint: disable=broad-except
+        if str(exc) == "Attempted to exit cancel scope in a different task than it was entered in":
+            logger.warning("MCP %s: patched_streamablehttp_client: enter/exit cancel scope task mismatch.", url)
+        else:
+            logger.error("MCP %s: patched_streamablehttp_client: %s: %s", url, type(exc).__name__, exc)

aixtools/server/__init__.py

@@ -13,10 +13,6 @@ from .utils import (
     get_session_id_tuple,
     run_in_thread,
 )
-from .workspace_privacy import (
-    is_session_private,
-    set_session_private,
-)
 
 __all__ = [
     "get_workspace_path",
@@ -24,6 +20,4 @@ __all__ = [
     "container_to_host_path",
     "host_to_container_path",
     "run_in_thread",
-    "is_session_private",
-    "set_session_private",
 ]

aixtools/server/path.py

@@ -2,47 +2,82 @@
 Workspace path handling for user sessions.
 """
 
-from pathlib import Path, PurePath, PurePosixPath
+import posixpath
+from pathlib import Path, PurePosixPath
 
 from fastmcp import Context
 
 from ..utils.config import DATA_DIR
 from .utils import get_session_id_tuple
 
-WORKSPACES_ROOT_DIR = DATA_DIR / "workspaces"  # Path on the host where workspaces are stored
+WORKSPACES_ROOT_DIR = (DATA_DIR / "workspaces").resolve()  # Path on the host where workspaces are stored
 CONTAINER_WORKSPACE_PATH = PurePosixPath("/workspace")  # Path inside the sandbox container where workspace is mounted
 
 
-def get_workspace_path(service_name: str = None, *, in_sandbox: bool = False, ctx: Context | tuple = None) -> PurePath:
+def get_workspace_path(ctx: Context | tuple | None = None) -> Path:
     """
-    Get the workspace path for a specific service (e.g. MCP server).
-    If `service_name` is None, then the path to entire workspace folder (as mounted to a sandbox) is returned.
-    If `in_sandbox` is True, it returns a path in sandbox, e.g.: `/workspace/mcp_repl`.
-    If `in_sandbox` is False, it returns the path based on user and session IDs in the format:
-    `<DATA_DIR>/workspaces/<user_id>/<session_id>/<service_name>`, where `DATA_DIR` should come from
-    the environment variables, e.g.:
-    `/data/workspaces/foo-user/bar-session/mcp_repl`.
-    The `ctx` is used to get user and session IDs tuple. It can be passed directly or via HTTP headers from `Context`.
-    If `ctx` is None, the current FastMCP request HTTP headers are used.
+    Get the workspace path (in the host) for a specific service (e.g. MCP server).
+    Returns the path based on user and session IDs in the format:
+
+        <DATA_DIR>/workspaces/<user_id>/<session_id>
+
+    where `DATA_DIR` should come from the environment variables
+    Example workspace path:
+
+        /data/workspaces/foo-user/bar-session
+
+    The `ctx` is used to get user and session IDs tuple. It can be passed directly
+    or via HTTP headers from `Context`. If `ctx` is None, the current FastMCP
+    request HTTP headers are used.
 
     Args:
         ctx: The FastMCP context, which contains the user session.
-        service_name: The name of the service (e.g. "mcp_server").
-        in_sandbox: If True, use a sandbox path; otherwise, use user/session-based path.
 
-    Returns: The workspace path as a PurePath object.
+    Returns: The workspace path as a Path object.
+    """
+    user_id, session_id = ctx if isinstance(ctx, tuple) else get_session_id_tuple(ctx)
+    return WORKSPACES_ROOT_DIR / user_id / session_id
+
+
+def get_workspace_path_sandbox() -> PurePosixPath:
+    """
+    Get the workspace path in the sandbox container.
+
+    We return PurePosixPath to ensure compatibility with Linux containers.
+
+    The paths inside the sandbox cannot be resolved (because they don't exist
+    on the host), so we use PurePosixPath instead of Path. Also Path could be
+    a WindowsPath on Windows hosts, which would be incorrect for Linux containers.
+
+    Returns: The workspace path as a PurePosixPath object.
     """
-    if in_sandbox:
-        path = CONTAINER_WORKSPACE_PATH
-    else:
-        user_id, session_id = ctx if isinstance(ctx, tuple) else get_session_id_tuple(ctx)
-        path = WORKSPACES_ROOT_DIR / user_id / session_id
-    if service_name:
-        path = path / service_name
-    return path
+    return CONTAINER_WORKSPACE_PATH
 
 
-def container_to_host_path(path: PurePosixPath, *, ctx: Context | tuple = None) -> Path | None:
+def path_normalize(p: PurePosixPath) -> PurePosixPath:
+    """
+    Normalize a PurePosixPath (remove redundant separators and up-level references).
+    """
+    return PurePosixPath(posixpath.normpath(p.as_posix()))
+
+
+def path_chroot(path: Path, old_root: Path, new_root: Path) -> Path:
+    """
+    Change the root of a given path from old_root to new_root.
+    If the path is not absolute (e.g. 'my_file.txt', './my_file.txt', 'my_dir/file.txt')
+    we treat it as relative to the 'new_root'
+    """
+    if not Path(path).is_absolute():
+        new_path = Path(new_root / path).resolve()
+        new_root = Path(new_root).resolve()
+        if not new_path.is_relative_to(new_root):
+            raise ValueError(f"Path must not escape the workspace root: '{path}'")
+        return Path(new_path)
+    # Otherwise, we treat it as absolute and change the root
+    return new_root / Path(path).relative_to(old_root)
+
+
+def container_to_host_path(path: PurePosixPath, *, ctx: Context | tuple | None = None) -> Path | None:
     """
     Convert a path in a sandbox container to a host path
 
@@ -54,19 +89,41 @@ def container_to_host_path(path: PurePosixPath, *, ctx: Context | tuple = None)
     Returns:
         Path to the file on the host, or None if the conversion fails.
     """
-    old_root = CONTAINER_WORKSPACE_PATH
+    # Try without service name (maybe the LLM forgot to put the SERVICE_NAME in the path)
+    old_root = get_workspace_path_sandbox()
     new_root = get_workspace_path(ctx=ctx)
     try:
-        return new_root / PurePosixPath(path).relative_to(old_root)
+        # Relative paths are treated as relative to the new_root
+        if not PurePosixPath(path).is_absolute():
+            # Resolve paths to prevent escaping the workspace root
+            new_path = Path(new_root / path).resolve()
+            new_root = Path(new_root.resolve())
+            if not new_path.is_relative_to(new_root):
+                raise ValueError(f"Path must not escape the workspace root: '{path}'")
+            return new_path
+        # Otherwise, we treat it as absolute and change the root
+        return new_root / Path(path).relative_to(old_root)
     except ValueError as e:
         raise ValueError(f"Container path must be a subdir of '{old_root}', got '{path}' instead") from e
 
 
-def host_to_container_path(path: Path, *, ctx: Context | tuple = None) -> PurePosixPath:
-    """Convert a host path to a path in a sandbox container."""
+def host_to_container_path(path: Path, *, ctx: Context | tuple | None = None) -> PurePosixPath:
+    """
+    Convert a host path to a path in a sandbox container.
+    Paths inside the sandbox MUST be PurePosixPath (i.e. we use Linux containers).
+    """
     old_root = get_workspace_path(ctx=ctx)
-    new_root = CONTAINER_WORKSPACE_PATH
+    new_root = get_workspace_path_sandbox()
     try:
+        # Relative paths are treated as relative to the new_root
+        if not Path(path).is_absolute():
+            # Normalize paths to prevent escaping the workspace root (we cannot resolve PurePosixPaths)
+            new_path = path_normalize(new_root / path)
+            new_root = path_normalize(new_root)
+            if not new_path.is_relative_to(new_root):
+                raise ValueError(f"Path must not escape the workspace root: '{path}'")
+            return new_path
+        # Otherwise, we treat it as absolute and change the root
         return new_root / Path(path).relative_to(old_root)
-    except ValueError as exc:
-        raise ValueError(f"Host path must be a subdir of '{old_root}', got '{path}' instead") from exc
+    except ValueError as e:
+        raise ValueError(f"Host path must be a subdir of either '{old_root}', got '{path}' instead") from e

aixtools/testing/aix_test_model.py

@@ -108,9 +108,17 @@ class AixTestModel(Model):
         messages: list[ModelMessage],
         model_settings: ModelSettings | None,
         model_request_parameters: ModelRequestParameters,
+        *args,  # pylint: disable=unused-argument # Accept additional arguments for compatibility with pydantic-ai 1.0.9
+        **kwargs,  # pylint: disable=unused-argument
     ) -> AsyncIterator[StreamedResponse]:
         model_response = await self._request(messages, model_settings, model_request_parameters)
-        yield TestStreamedResponse(_model_name=self.model_name, _structured_response=model_response, _messages=messages)
+        yield TestStreamedResponse(
+            _model_name=self.model_name,
+            _structured_response=model_response,
+            _messages=messages,
+            model_request_parameters=model_request_parameters,
+            _provider_name="",
+        )
 
     @property
     def model_name(self) -> str:

aixtools/tools/doctor/mcp_tool_doctor.py

@@ -0,0 +1,79 @@
+import argparse
+import asyncio
+
+from pydantic_ai.mcp import MCPServerStdio, MCPServerStreamableHTTP
+
+from aixtools.agents import get_agent, run_agent
+from aixtools.tools.doctor.tool_doctor import TOOL_DOCTOR_PROMPT
+from aixtools.tools.doctor.tool_recommendation import ToolRecommendation
+
+
+async def tool_doctor_mcp(
+    mcp_url: str = "http://127.0.0.1:8000/mcp",
+    mcp_server: MCPServerStreamableHTTP | MCPServerStdio | None = None,
+    verbose: bool = False,
+    debug: bool = False,
+) -> list[ToolRecommendation]:
+    """
+    Run the tool doctor agent to analyze tools from an MCP server and give recommendations.
+
+    Usage examples:
+        # Using an http MCP server
+        ret = await tool_doctor_mcp(mcp_url='http://127.0.0.1:8000/mcp')
+        print(ret)
+
+        # Using a stdio MCP server
+        server = MCPServerStdio(command='fastmcp', args=['run', 'my_mcp_server.py'])
+        ret = await tool_doctor_mcp(mcp_server=server)
+        print(ret)
+    """
+    if mcp_server is None:
+        mcp_server = MCPServerStreamableHTTP(url=mcp_url)
+    agent = get_agent(toolsets=[mcp_server], output_type=list[ToolRecommendation])
+    async with agent:
+        ret, nodes = await run_agent(agent, TOOL_DOCTOR_PROMPT, verbose=verbose, debug=debug)
+    return ret  # type: ignore
+
+
+def main_cli():
+    """Command line interface for tool doctor MCP."""
+    parser = argparse.ArgumentParser(description="Analyze tools from an MCP server and provide recommendations")
+
+    # MCP server connection options
+    server_group = parser.add_mutually_exclusive_group()
+    server_group.add_argument(
+        "--mcp-url",
+        default="http://127.0.0.1:8000/mcp",
+        help="URL of the HTTP MCP server (default: http://127.0.0.1:8000/mcp)",
+    )
+    server_group.add_argument("--stdio-command", help="Command to run STDIO MCP server (e.g., 'fastmcp')")
+
+    parser.add_argument(
+        "--stdio-args",
+        nargs="*",
+        default=[],
+        help="Arguments for STDIO MCP server command (e.g., 'run', 'my_server.py')",
+    )
+    parser.add_argument("--verbose", action="store_true", help="Enable verbose output")
+    parser.add_argument("--debug", action="store_true", help="Enable debug output")
+
+    args = parser.parse_args()
+
+    async def run():
+        mcp_server = None
+        if args.stdio_command:
+            mcp_server = MCPServerStdio(command=args.stdio_command, args=args.stdio_args)
+            recommendations = await tool_doctor_mcp(mcp_server=mcp_server, verbose=args.verbose, debug=args.debug)
+        else:
+            recommendations = await tool_doctor_mcp(mcp_url=args.mcp_url, verbose=args.verbose, debug=args.debug)
+
+        print("Tool Doctor Recommendations:")
+        print("=" * 50)
+        for i, rec in enumerate(recommendations, 1):
+            print(f"\n{i}. {rec}")
+
+    asyncio.run(run())
+
+
+if __name__ == "__main__":
+    main_cli()

aixtools/tools/doctor/tool_doctor.py

@@ -1,3 +1,7 @@
+"""
+Tool doctor: Analyze tools and give recommendations for improvement.
+"""
+
 from aixtools.agents import get_agent, run_agent
 from aixtools.tools.doctor.tool_recommendation import ToolRecommendation
 

aixtools/tools/doctor/tool_recommendation.py

@@ -1,5 +1,10 @@
 from pydantic import BaseModel
 
+"""
+These are classes that represent recommendations for improving tools.
+They are ued by ToolDoctor
+"""
+
 
 class ArgumentRecommendation(BaseModel):
     """A recommendation for an argument"""

aixtools/utils/config.py

@@ -71,7 +71,6 @@ logging.warning("Using         DATA_DIR='%s'", DATA_DIR)
 VDB_CHROMA_PATH = DATA_DB_DIR / "chroma.db"
 VDB_DEFAULT_SIMILARITY_THRESHOLD = 0.85
 
-
 # ---
 # Variables in '.env' file
 # Explicitly load specific variables