aixtools 0.1.0__py3-none-any.whl

aixtools/mcp/faulty_mcp.py

@@ -0,0 +1,320 @@
+#!/usr/bin/env python3
+"""
+Faulty MCP Server for Testing MCP Errors
+- Simulates 404 errors for specific MCP requests
+"""
+
+import argparse
+import asyncio
+import json
+import logging.config
+import os
+from dataclasses import dataclass
+from random import choice, random
+
+from fastapi import HTTPException, status
+from fastmcp import FastMCP
+from fastmcp.exceptions import (
+    ClientError,
+    DisabledError,
+    InvalidSignature,
+    NotFoundError,
+    PromptError,
+    ResourceError,
+    ToolError,
+    ValidationError,
+)
+from fastmcp.server.middleware import Middleware as McpMiddleware
+from fastmcp.server.middleware import MiddlewareContext
+from fastmcp.server.middleware.logging import LoggingMiddleware
+from starlette.middleware import Middleware as StarletteMiddleware
+from starlette.types import Receive, Scope, Send
+
+from aixtools.logging.logging_config import DEFAULT_LOGGING_CONFIG
+from aixtools.utils import get_logger
+
+# Remove the user/session ID from logger line to shorten it
+DEFAULT_LOGGING_CONFIG["formatters"]["color"]["format"] = (
+    "%(log_color)s%(asctime)s.%(msecs)03d %(levelname)-8s%(reset)s %(message)s"
+)
+logging.config.dictConfig(DEFAULT_LOGGING_CONFIG)
+
+# Get the logger
+logger = get_logger(__name__)
+
+
+@dataclass
+class Config:
+    """Global configuration for the faulty MCP server."""
+
+    port: int = 9999
+    prob_throw_in_list_handle: float = 0.5  # Probability of throwing an exception in list tools handling
+    prob_delete_404: float = 0.5  # Probability of injecting a 404 error for DELETE requests
+    prob_general_404: float = 0.5  # Probability of injecting a 404 error for other requests
+    prob_terminate_on_empty_request: float = 0.3  # Probability of terminating the process on empty request
+    prob_terminate_in_list_handle: float = 0.3  # Probability of terminating the process in list tools handling
+
+
+# Global configuration
+config = Config()
+
+
+class McpErrorMiddleware(McpMiddleware):
+    """Custom middleware to simulate errors in MCP requests."""
+
+    async def __call__(self, context: MiddlewareContext, call_next):
+        # This method receives ALL messages regardless of type
+        logger.info("[McpErrorMiddleware] processing: %s", context.method)
+
+        if context.method == "tools/list":
+            random_number = random()
+            logger.info("[McpErrorMiddleware] random number: %f", random_number)
+            if random_number < config.prob_terminate_in_list_handle:
+                logger.warning("[McpErrorMiddleware] Simulating server crash!")
+                os.kill(os.getpid(), 9)
+
+            if random_number < config.prob_throw_in_list_handle:
+                exception_class = choice(
+                    [
+                        ValidationError,
+                        ResourceError,
+                        ToolError,
+                        PromptError,
+                        InvalidSignature,
+                        ClientError,
+                        NotFoundError,
+                        DisabledError,
+                    ]
+                )
+                logger.warning("[McpErrorMiddleware] throwing %s for: %s", exception_class.__name__, context.method)
+                raise exception_class(f"[McpErrorMiddleware] {exception_class.__name__}.")
+
+        result = await call_next(context)
+        logger.info("[McpErrorMiddleware] completed: %s", context.method)
+        return result
+
+
+class StarletteErrorMiddleware:  # pylint: disable=too-few-public-methods
+    """Custom Starlette middleware to log and inject errors."""
+
+    def __init__(self, app):
+        """Initialize middleware."""
+        self.app = app
+        logger.info("[StarletteErrorMiddleware] Middleware initialized!")
+        logger.info("Current configuration:")
+        logger.info("Exception in list tools handling probability: %f", config.prob_throw_in_list_handle)
+        logger.info("DELETE 404 probability: %f", config.prob_delete_404)
+        logger.info("General 404 probability: %f", config.prob_general_404)
+        logger.info("Terminate on empty request probability: %f", config.prob_terminate_on_empty_request)
+        logger.info("Terminate in list handle probability: %f", config.prob_terminate_in_list_handle)
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send):  # noqa: PLR0915 # pylint: disable=too-many-statements
+        # Log all the condition variables for debugging
+
+        logger.info("[StarletteErrorMiddleware] >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>")
+        logger.info("[StarletteErrorMiddleware] scope['type']: %s", scope.get("type", "unknown"))
+        logger.info("[StarletteErrorMiddleware] scope['path']: %s", scope.get("path", "unknown"))
+        logger.info("[StarletteErrorMiddleware] HTTP method: %s", http_method := scope.get("method", "unknown"))
+        logger.info("[StarletteErrorMiddleware] Headers: %s", str(dict(scope.get("headers", []))))
+
+        # Wrap receive to log body content without breaking the flow
+        body_parts = []
+        should_inject_404 = False
+
+        if http_method == "DELETE":
+            random_number = random()
+            logger.info("[StarletteErrorMiddleware] random number: %f", random_number)
+            if random_number < config.prob_delete_404:
+                logger.info("[StarletteErrorMiddleware] Simulating 404 error for DELETE request")
+                should_inject_404 = True
+
+        async def logging_receive():
+            nonlocal should_inject_404
+            message = await receive()
+            logger.info("[StarletteErrorMiddleware] Received message: %s", str(message))
+
+            # Check for empty request and possibly terminate the process
+            if message["type"] == "http.request" and message["body"] == b"" and not message.get("more_body", False):
+                random_number = random()
+                logger.info("[StarletteErrorMiddleware] Empty request received, random number: %f", random_number)
+                if random_number < config.prob_terminate_on_empty_request:
+                    logger.warning("[StarletteErrorMiddleware] Simulating server crash on empty request!")
+                    os.kill(os.getpid(), 9)
+
+            if message["type"] == "http.request":  # pylint: disable=too-many-nested-blocks
+                body = message.get("body", b"")
+                if body:
+                    body_parts.append(body)
+
+                # Log when we have the complete body
+                if not message.get("more_body", False):
+                    complete_body = b"".join(body_parts)
+                    if complete_body:
+                        try:
+                            body_str = complete_body.decode("utf-8")
+                            logger.info("[StarletteErrorMiddleware] Request body: %s", body_str)
+
+                            json_data = json.loads(body_str)
+                            if isinstance(json_data, dict):
+                                method_name = json_data.get("method", "unknown")
+                                if method_name == "initialize" and not json_data.get("params", {}).get("capabilities"):
+                                    logger.info("Detected initial health check from Navari, skipping 404 injection.")
+                                    return message
+                                logger.info("[StarletteErrorMiddleware] MCP method: %s", method_name)
+
+                                # Check if we should inject 404
+                                random_number = random()
+                                logger.info("[StarletteErrorMiddleware] random number: %f", random_number)
+                                if random_number < config.prob_general_404:
+                                    should_inject_404 = True
+                                    logger.info("[StarletteErrorMiddleware] %s - will inject 404!", method_name)
+                        except (UnicodeDecodeError, json.JSONDecodeError) as e:
+                            logger.exception("[StarletteErrorMiddleware] Error parsing body: %s", e)
+                    else:
+                        logger.info("[StarletteErrorMiddleware] Request body: (empty)")
+
+            return message
+
+        async def intercepting_send(message):
+            logger.info("[StarletteErrorMiddleware] Sending message: %s", str(message))
+            if message["type"] == "http.response.start" and should_inject_404:
+                logger.warning("[StarletteErrorMiddleware] Injecting 404!")
+                # Replace the response with 404
+                message = {
+                    "type": "http.response.start",
+                    "status": 404,
+                    "headers": [[b"content-type", b"text/plain"]],
+                }
+            elif message["type"] == "http.response.body" and should_inject_404:
+                # Replace body with 404 message
+                message = {
+                    "type": "http.response.body",
+                    "body": b"Simulated 404 error",
+                }
+
+            await send(message)
+
+        logger.info("[StarletteErrorMiddleware] <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<")
+        await self.app(scope, logging_receive, intercepting_send)
+
+
+# Create the MCP server
+mcp = FastMCP(
+    name="Faulty MCP Server",
+    instructions="A simple test server for reproducing MCP errors.",
+    middleware=[LoggingMiddleware(include_payloads=True), McpErrorMiddleware()],
+)
+
+
+@mcp.tool
+def add(a: float, b: float) -> float:
+    """Add two numbers together."""
+    return a + b
+
+
+@mcp.tool
+def always_error() -> None:
+    """Always throw an exception to simulate errors."""
+    raise ValueError("Simulated error")
+
+
+@mcp.tool
+async def freeze_server(seconds: int = 60) -> str:
+    """Simulate a server freeze for testing purposes."""
+    await asyncio.sleep(seconds)
+    return f"Server frozen for {seconds} seconds"
+
+
+@mcp.tool
+def multiply(a: float, b: float) -> float:
+    """Multiply two numbers together."""
+    return a * b
+
+
+@mcp.tool
+def random_throw_exception(a: float, b: float, prob: float = 0.5) -> float:
+    """Randomly throw an exception to simulate errors."""
+    if random() < prob:
+        raise ValueError("Simulated error")
+    return a * b
+
+
+@mcp.tool
+def throw_404_exception() -> None:
+    """Randomly throw an exception to simulate errors."""
+    raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Throwing a 404 error for testing purposes.")
+
+
+def run_server_on_port():
+    """Run a single MCP server using the global configuration."""
+
+    async def run_async():
+        print(f"[Port {config.port}] Starting MCP server on http://localhost:{config.port}/mcp/")
+        await mcp.run_http_async(
+            transport="streamable-http",
+            host="localhost",
+            port=config.port,
+            path="/mcp/",
+            middleware=[StarletteMiddleware(StarletteErrorMiddleware)],
+        )
+
+    asyncio.run(run_async())
+
+
+if __name__ == "__main__":
+    # Parse command line arguments
+    parser = argparse.ArgumentParser(description="Run a faulty MCP server for testing error handling")
+    parser.add_argument(
+        "--safe-mode",
+        action="store_true",
+        help="Set all error probabilities to 0 by default, only use explicitly provided values",
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        help=f"Port to run the server on (default: {config.port})",
+    )
+    parser.add_argument(
+        "--prob-throw-in-list-handle",
+        type=float,
+        help=f"Probability of exception in list tools handling (default: {config.prob_throw_in_list_handle})",
+    )
+    parser.add_argument(
+        "--prob-delete-404",
+        type=float,
+        help=f"Probability of injecting a 404 error for DELETE requests (default: {config.prob_delete_404})",
+    )
+    parser.add_argument(
+        "--prob-general-404",
+        type=float,
+        help=f"Probability of injecting a 404 error for other requests (default: {config.prob_general_404})",
+    )
+    parser.add_argument(
+        "--prob-terminate-on-empty-request",
+        type=float,
+        help=f"Probability of terminating on empty request (default: {config.prob_terminate_on_empty_request})",
+    )
+    parser.add_argument(
+        "--prob-terminate-in-list-handle",
+        type=float,
+        help=f"Probability of terminating in list tools handling (default: {config.prob_terminate_in_list_handle})",
+    )
+
+    args = parser.parse_args()
+
+    def _update_config_value(attr_name: str):
+        if args.safe_mode:
+            setattr(config, attr_name, 0)
+        if (given_value := getattr(args, attr_name)) is not None:
+            setattr(config, attr_name, given_value)
+
+    # Update the global configuration with command line arguments
+    config.port = args.port or config.port
+    _update_config_value("prob_throw_in_list_handle")
+    _update_config_value("prob_delete_404")
+    _update_config_value("prob_general_404")
+    _update_config_value("prob_terminate_on_empty_request")
+    _update_config_value("prob_terminate_in_list_handle")
+
+    # Run the server
+    run_server_on_port()

aixtools/model_patch/model_patch.py

@@ -0,0 +1,65 @@
+from typing import Any
+
+from pydantic import BaseModel
+
+from aixtools.logging.logging_config import get_logger
+
+logger = get_logger(__name__)
+
+
+class ModelRawRequest(BaseModel):
+    method_name: str  # Model method name
+    request_id: str  # Unique request ID
+    args: tuple  # Method arguments
+    kwargs: dict  # Method keyword arguments
+
+
+class ModelRawRequestResult(BaseModel):
+    method_name: str  # Model method name
+    request_id: str  # Unique request ID
+    result: Any  # Method's result
+
+    class Config:
+        arbitrary_types_allowed = True
+
+
+class ModelRawRequestYieldItem(BaseModel):
+    method_name: str  # Model method name
+    request_id: str  # Unique request ID
+    item_num: int  # Item number in the stream
+    item: Any  # Yielded item
+
+    class Config:
+        arbitrary_types_allowed = True
+
+
+def get_request_fn(model):
+    """Get the original request method"""
+    if is_patched(model):
+        return model._request_ori
+    return model.request
+
+
+def get_request_stream_fn(model):
+    """Get the original request method"""
+    if is_patched(model):
+        return model._request_stream_ori
+    return model.request_stream
+
+
+def is_patched(model):
+    return hasattr(model, "_request_ori") or hasattr(model, "_request_stream_ori")
+
+
+def model_patch(model, request_method, request_stream_method):
+    """Replace model.request and model.request_stream with logging versions"""
+    if is_patched(model):
+        logger.warning(f"Model {model.__class__.__name__} is already patched. Skipping patching.")
+        return model
+    # Save original methods
+    model._request_ori = model.request
+    model._request_stream_ori = model.request_stream
+    # Patch methods
+    model.request = request_method
+    model.request_stream = request_stream_method
+    return model

aixtools/server/__init__.py

@@ -0,0 +1,23 @@
+"""
+FastMCP utilities for:
+    - extracting user metadata from context
+    - running mcp tools tasks in a separate thread.
+"""
+
+from .path import (
+    container_to_host_path,
+    get_workspace_path,
+    host_to_container_path,
+)
+from .utils import (
+    get_session_id_tuple,
+    run_in_thread,
+)
+
+__all__ = [
+    "get_workspace_path",
+    "get_session_id_tuple",
+    "container_to_host_path",
+    "host_to_container_path",
+    "run_in_thread",
+]

aixtools/server/app_mounter.py

@@ -0,0 +1,90 @@
+"""Utility for mounting sub-applications with lifespan management in Starlette/FastAPI."""
+
+from contextlib import asynccontextmanager
+
+from starlette.applications import Starlette
+from starlette.types import ASGIApp
+
+
+class AppMounter:  # pylint: disable=too-few-public-methods
+    """
+    A utility class for mounting sub-applications with their lifespans.
+
+    This class handles the complexity of ensuring that mounted sub-applications
+    have their lifespans properly managed alongside the parent application.
+    """
+
+    def __init__(self, parent_app: Starlette):
+        """
+        Initialize the SubAppMounter with a parent application.
+
+        Args:
+            parent_app: The parent Starlette/FastAPI application
+        """
+        self.parent_app = parent_app
+        self.mounted_apps: list[tuple[str, ASGIApp]] = []
+        self._original_lifespan = parent_app.router.lifespan_context
+        self._setup_combined_lifespan()
+
+    def mount_with_lifespan(self, path: str, app: ASGIApp, name: str = None) -> None:
+        """
+        Mount a sub-application and ensure its lifespan is managed.
+
+        Args:
+            path: The path to mount the application at
+            app: The ASGI application to mount
+            name: Optional name for the mounted application
+        """
+        # Mount the app using the parent app's mount method
+        self.parent_app.mount(path, app=app, name=name)
+
+        # Store the mounted app for lifespan management
+        self.mounted_apps.append((path, app))
+
+    def _setup_combined_lifespan(self) -> None:
+        """
+        Set up a combined lifespan that manages both the parent app and all mounted sub-apps.
+        """
+
+        @asynccontextmanager
+        async def combined_lifespan(app):
+            # First enter the parent app's lifespan
+            async with self._original_lifespan(app):
+                # Then enter each mounted app's lifespan in order
+                # We use nested context managers to ensure proper cleanup
+                async with self._create_nested_lifespans():
+                    yield
+
+        # Replace the parent app's lifespan with our combined one
+        self.parent_app.router.lifespan_context = combined_lifespan
+
+    @asynccontextmanager
+    async def _create_nested_lifespans(self):
+        """
+        Create nested async context managers for all mounted apps.
+
+        This ensures that all sub-app lifespans are entered and exited in the correct order.
+        """
+        # If no apps are mounted, just yield
+        if not self.mounted_apps:
+            yield
+            return
+
+        # Otherwise, create nested context managers for each mounted app
+        async def enter_lifespans(index=0):
+            if index >= len(self.mounted_apps):
+                yield
+                return
+
+            _, app = self.mounted_apps[index]
+            if hasattr(app, "router") and hasattr(app.router, "lifespan_context"):
+                async with app.router.lifespan_context(app):
+                    async for _ in enter_lifespans(index + 1):
+                        yield
+            else:
+                # If the app doesn't have a lifespan, just move to the next one
+                async for _ in enter_lifespans(index + 1):
+                    yield
+
+        async for _ in enter_lifespans():
+            yield

aixtools/server/path.py

@@ -0,0 +1,72 @@
+"""
+Workspace path handling for user sessions.
+"""
+
+from pathlib import Path, PurePath, 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
+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:
+    """
+    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.
+
+    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.
+    """
+    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
+
+
+def container_to_host_path(path: PurePosixPath, *, ctx: Context | tuple = None) -> Path | None:
+    """
+    Convert a path in a sandbox container to a host path
+
+    Args:
+        container_path: Path inside the container (must be a subdir of CONTAINER_WORKSPACE_PATH).
+        user_id: ID of the user.
+        session_id: ID of the session.
+
+    Returns:
+        Path to the file on the host, or None if the conversion fails.
+    """
+    old_root = CONTAINER_WORKSPACE_PATH
+    new_root = get_workspace_path(ctx=ctx)
+    try:
+        return new_root / PurePosixPath(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."""
+    old_root = get_workspace_path(ctx=ctx)
+    new_root = CONTAINER_WORKSPACE_PATH
+    try:
+        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

aixtools/server/utils.py

@@ -0,0 +1,70 @@
+"""
+FastMCP server utilities for handling user context and threading.
+"""
+
+import asyncio
+from functools import wraps
+
+from fastmcp import Context
+from fastmcp.server import dependencies
+
+from ..context import session_id_var, user_id_var
+
+
+def get_session_id_tuple(ctx: Context | None = None) -> tuple[str, str]:
+    """
+    Get the user and session IDs from the user session.
+    If `ctx` is None, the current FastMCP request HTTP headers are used.
+    Returns: Tuple of (user_id, session_id).
+    """
+    user_id = get_user_id_from_request(ctx)
+    user_id = user_id or user_id_var.get("default_user")
+    session_id = get_session_id_from_request(ctx)
+    session_id = session_id or session_id_var.get("default_session")
+    return user_id, session_id
+
+
+def get_session_id_from_request(ctx: Context | None = None) -> str | None:
+    """
+    Get the session ID from the HTTP request headers.
+    If `ctx` is None, the current FastMCP request HTTP headers are used.
+    """
+    try:
+        return (ctx or dependencies).get_http_request().headers.get("session-id")
+    except (ValueError, RuntimeError):
+        return None
+
+
+def get_user_id_from_request(ctx: Context | None = None) -> str | None:
+    """
+    Get the user ID from the HTTP request headers.
+    If `ctx` is None, the current FastMCP request HTTP headers are used.
+    The user_id is always returned as lowercase.
+
+    Returns:
+        str | None: The lowercase user ID, or None if not found or an error occurs.
+    """
+    try:
+        user_id = (ctx or dependencies).get_http_request().headers.get("user-id")
+        return user_id.lower() if user_id else None
+    except (ValueError, RuntimeError, AttributeError):
+        return None
+
+
+def get_session_id_str(ctx: Context | None = None) -> str:
+    """
+    Combined session ID for the current user and session.
+    If `ctx` is None, the current FastMCP request HTTP headers are used.
+    """
+    user_id, session_id = get_session_id_tuple(ctx)
+    return f"{user_id}:{session_id}"
+
+
+def run_in_thread(func):
+    """decorator to run blocking function with `asyncio.to_thread`"""
+
+    @wraps(func)
+    async def wrapper(*args, **kwargs):
+        return await asyncio.to_thread(func, *args, **kwargs)
+
+    return wrapper

aixtools/testing/__init__.py

@@ -0,0 +1,9 @@
+"""
+Testing utilities for AI agents and model patching.
+"""
+
+from aixtools.testing.model_patch_cache import model_patch_cache
+
+__all__ = [
+    "model_patch_cache",
+]