@macaw-ai/client 0.0.0-stage → 0.9.11

package/payload/secureAI/macaw_adapters/mcp/mcp.py

@@ -0,0 +1,785 @@
+"""
+MACAW SecureMCP - FastMCP-compatible API with MACAW Security.
+
+Provides the same decorator-based API as FastMCP while routing all tool
+invocations through the MACAW security layer for policy enforcement,
+cryptographic signing, and audit logging.
+
+Usage:
+    from macaw_adapters.mcp import SecureMCP, Context
+
+    mcp = SecureMCP("calculator")
+
+    @mcp.tool(description="Add two numbers")
+    def add(a: float, b: float) -> float:
+        return a + b
+
+    @mcp.resource("calc://history")
+    def get_history(ctx: Context) -> list:
+        return ctx.get("calc_history") or []
+
+    if __name__ == "__main__":
+        mcp.run()
+"""
+
+import asyncio
+import inspect
+import logging
+from typing import Any, Callable, Dict, List, Optional
+from functools import wraps
+from dataclasses import dataclass, field
+from datetime import datetime
+
+logger = logging.getLogger(__name__)
+
+
+def _json_schema(parameters: Dict[str, Any]) -> Dict[str, Any]:
+    """Convert _extract_parameters() output into a JSON Schema for MCP's inputSchema."""
+    return {
+        "type": "object",
+        "properties": {
+            name: {"type": info.get("type", "string")}
+            for name, info in parameters.items()
+        },
+        "required": [
+            name for name, info in parameters.items() if info.get("required")
+        ],
+    }
+
+
+@dataclass
+class Context:
+    """
+    Request context passed to tool handlers.
+
+    Provides access to context vault, progress reporting, and cross-tool calls.
+    Compatible with FastMCP's Context object.
+    """
+
+    request_id: str = ""
+    tool_name: str = ""
+    source_agent: str = ""  # Caller's agent ID (for sampling/elicitation callbacks)
+    _client: Any = None  # MACAWClient reference
+    _metadata: Dict[str, Any] = field(default_factory=dict)
+
+    def get(self, key: str) -> Any:
+        """Get value from context vault."""
+        if self._client:
+            return self._client.context_get(key)
+        return None
+
+    def set(self, key: str, value: Any) -> None:
+        """Set value in context vault."""
+        if self._client:
+            self._client.context_set(key, value)
+
+    async def report_progress(self, progress: float, message: str = None) -> None:
+        """
+        Report execution progress (0.0 to 1.0).
+
+        This is async to support streaming progress updates.
+
+        Args:
+            progress: Progress percentage from 0.0 to 1.0
+            message: Optional progress message
+        """
+        if self._client:
+            self._client.log_event(
+                event_type="tool_progress",
+                action="progress",
+                target=self.tool_name,
+                metadata={
+                    "request_id": self.request_id,
+                    "progress": progress,
+                    "message": message
+                }
+            )
+
+    async def read_resource(self, uri: str) -> Any:
+        """
+        Read another resource (for resource dependencies).
+
+        Args:
+            uri: Resource URI to read
+
+        Returns:
+            Resource content
+        """
+        if self._client:
+            return self._client.invoke_tool(f"resource:{uri}", {})
+        return None
+
+    async def sample(
+        self,
+        prompt: str,
+        system_prompt: str = None,
+        max_tokens: int = 1000,
+        temperature: Optional[float] = None,
+        **kwargs
+    ) -> str:
+        """
+        Request an LLM completion from the client (MCP Sampling).
+
+        This allows server tools to leverage the client's LLM capabilities
+        for AI-assisted processing within tool execution.
+
+        Args:
+            prompt: The user prompt to send to the LLM
+            system_prompt: Optional system prompt
+            max_tokens: Maximum tokens in response
+            temperature: Sampling temperature. Left unset by default: this server
+                does not know which model the client will use, and sampling
+                parameters are rejected outright by some models. Set it only if
+                the caller has a reason to, and accept that a client may ignore
+                or reject it.
+            **kwargs: Additional LLM parameters
+
+        Returns:
+            The LLM's text response
+
+        Example:
+            summary = await ctx.sample(
+                prompt=f"Summarize this text: {long_text}",
+                system_prompt="You are a helpful summarization assistant.",
+                max_tokens=500
+            )
+        """
+        if not self._client:
+            raise RuntimeError("No client available for sampling")
+
+        if not self.source_agent:
+            raise RuntimeError("No source_agent available - sampling requires caller context")
+
+        params = {
+            "prompt": prompt,
+            "system_prompt": system_prompt,
+            "max_tokens": max_tokens,
+            "request_id": self.request_id,
+            "source_tool": self.tool_name,
+            **kwargs
+        }
+        # Only send what the caller actually asked for. Inventing a temperature
+        # here means shipping a model-specific parameter to a model we know
+        # nothing about - which is the whole point of sampling.
+        if temperature is not None:
+            params["temperature"] = temperature
+
+        # Route sampling request back to the calling client's sampling handler
+        logger.info(f"[ctx.sample] Calling back to source_agent: {self.source_agent}")
+        result = self._client.invoke_tool(
+            "_mcp_sample",
+            params,
+            target_agent=self.source_agent  # Call back to the original caller
+        )
+
+        if isinstance(result, dict):
+            return result.get("response", result.get("result", str(result)))
+        return str(result)
+
+    async def elicit(
+        self,
+        prompt: str,
+        options: List[str] = None,
+        input_type: str = "text",
+        default: str = None,
+        required: bool = True,
+        **kwargs
+    ) -> Any:
+        """
+        Request user input from the client (MCP Elicitation).
+
+        This allows server tools to interactively request information
+        from the user during tool execution.
+
+        Args:
+            prompt: The question or prompt to show the user
+            options: List of choices for selection input
+            input_type: Type of input ("text", "select", "confirm", "number")
+            default: Default value if user doesn't provide input
+            required: Whether input is required
+            **kwargs: Additional elicitation parameters
+
+        Returns:
+            User's input (string, bool for confirm, or selected option)
+
+        Example:
+            # Simple text input
+            name = await ctx.elicit("What is your name?")
+
+            # Confirmation
+            proceed = await ctx.elicit(
+                "Delete this file?",
+                input_type="confirm",
+                default="no"
+            )
+
+            # Selection from options
+            color = await ctx.elicit(
+                "Choose a color:",
+                options=["red", "green", "blue"],
+                input_type="select"
+            )
+        """
+        if not self._client:
+            raise RuntimeError("No client available for elicitation")
+
+        if not self.source_agent:
+            raise RuntimeError("No source_agent available - elicitation requires caller context")
+
+        # Route elicitation request back to the calling client's handler
+        result = self._client.invoke_tool(
+            "_mcp_elicit",
+            {
+                "prompt": prompt,
+                "options": options,
+                "input_type": input_type,
+                "default": default,
+                "required": required,
+                "request_id": self.request_id,
+                "source_tool": self.tool_name,
+                **kwargs
+            },
+            target_agent=self.source_agent  # Call back to the original caller
+        )
+
+        if isinstance(result, dict):
+            response = result.get("response", result.get("result"))
+
+            # Convert response based on input_type
+            if input_type == "confirm":
+                if isinstance(response, bool):
+                    return response
+                return str(response).lower() in ("yes", "y", "true", "1")
+            elif input_type == "number":
+                try:
+                    return float(response) if "." in str(response) else int(response)
+                except (ValueError, TypeError):
+                    return default
+            return response
+
+        return result
+
+    # ========================================
+    # Logging methods (FastMCP-compatible)
+    # ========================================
+
+    def log(self, level: str, message: str, **kwargs) -> None:
+        """
+        Log a message at the specified level.
+
+        Args:
+            level: Log level (debug, info, warning, error)
+            message: Log message
+            **kwargs: Additional metadata
+        """
+        if self._client:
+            self._client.log_event(
+                event_type="tool_log",
+                source=self.tool_name,
+                action=level,
+                target=self.request_id,
+                outcome="success",  # Log events are informational, always "success"
+                metadata={"level": level, "message": message, **kwargs}
+            )
+
+    def debug(self, message: str, **kwargs) -> None:
+        """Log a debug message."""
+        self.log("debug", message, **kwargs)
+
+    def info(self, message: str, **kwargs) -> None:
+        """Log an info message."""
+        self.log("info", message, **kwargs)
+
+    def warning(self, message: str, **kwargs) -> None:
+        """Log a warning message."""
+        self.log("warning", message, **kwargs)
+
+    def error(self, message: str, **kwargs) -> None:
+        """Log an error message."""
+        self.log("error", message, **kwargs)
+
+    def audit(
+        self,
+        action: str,
+        target: str = None,
+        outcome: str = "success",
+        **metadata
+    ) -> None:
+        """
+        Create a cryptographically signed audit entry.
+
+        This is a MACAW-specific feature that creates tamper-evident audit logs
+        with cryptographic signatures for compliance and forensics.
+
+        Args:
+            action: The action being audited (e.g., "data_access", "approval", "modification")
+            target: What the action affects (e.g., resource name, record ID)
+            outcome: Result of the action ("success", "failure", "denied")
+            **metadata: Additional audit metadata
+
+        Example:
+            ctx.audit("data_access", target="customer_records", outcome="success",
+                      records_accessed=42, query="SELECT * FROM customers")
+        """
+        if self._client:
+            self._client.log_event(
+                event_type="audit",
+                source=self.tool_name,
+                action=action,
+                target=target or self.tool_name,
+                outcome=outcome,
+                signed=True,  # Top-level param for cryptographic signing
+                metadata={
+                    "request_id": self.request_id,
+                    **metadata
+                }
+            )
+
+    @property
+    def metadata(self) -> Dict[str, Any]:
+        """Get request metadata."""
+        return self._metadata
+
+    def get_roots(self) -> List[str]:
+        """
+        Get the list of filesystem roots this server can access.
+
+        Returns:
+            List of root paths declared for this server
+        """
+        return self._metadata.get('roots', [])
+
+
+class SecureMCP:
+    """
+    FastMCP-compatible server with MACAW security.
+
+    All tool invocations are routed through MACAWClient for:
+    - Cryptographic signing
+    - Policy enforcement
+    - Audit logging
+    - Prompt lineage tracking
+    """
+
+    def __init__(
+        self,
+        name: str,
+        version: str = "1.0.0",
+        intent_policy: Optional[Dict[str, Any]] = None,
+        roots: Optional[List[str]] = None,
+        **kwargs
+    ):
+        """
+        Initialize SecureMCP server.
+
+        Args:
+            name: Server name (used in agent_id)
+            version: Server version
+            intent_policy: MAPL policy declaration (resources, denied_resources, etc.)
+            roots: List of filesystem paths this server can access (MCP roots)
+            **kwargs: Additional MACAWClient options (iam_token, provider, etc.)
+
+        Example:
+            # Server that can only access specific directories
+            mcp = SecureMCP(
+                "file-manager",
+                roots=["/home/user/documents", "/tmp/workspace"]
+            )
+        """
+        self.name = name
+        self.version = version
+        self.extra_config = kwargs
+
+        # Store custom intent_policy or build default
+        self._custom_intent_policy = intent_policy
+
+        # MCP Roots - directories this server can access
+        self._roots: List[str] = roots or []
+
+        # Registries for MCP components
+        self._tools: Dict[str, Dict[str, Any]] = {}
+        self._resources: Dict[str, Dict[str, Any]] = {}
+        self._prompts: Dict[str, Dict[str, Any]] = {}
+
+        # MACAWClient (initialized in run())
+        self._client = None
+        self.agent_id = None
+
+    @property
+    def roots(self) -> List[str]:
+        """Get the list of filesystem roots this server can access."""
+        return self._roots
+
+    def tool(
+        self,
+        name: str = None,
+        description: str = None,
+        prompts: List[str] = None
+    ) -> Callable:
+        """
+        Decorator to register a tool.
+
+        Args:
+            name: Tool name (defaults to function name)
+            description: Tool description (defaults to docstring)
+            prompts: List of parameter names that should be treated as prompts
+
+        Example:
+            @mcp.tool(description="Add two numbers")
+            def add(a: float, b: float) -> float:
+                return a + b
+
+            @mcp.tool(prompts=["query"])
+            def search(query: str) -> list:
+                # query will be wrapped as AuthenticatedPrompt
+                return do_search(query)
+        """
+        def decorator(func: Callable) -> Callable:
+            tool_name = name or func.__name__
+            tool_desc = description or func.__doc__ or ""
+
+            # Check if function wants Context parameter
+            sig = inspect.signature(func)
+            wants_context = "ctx" in sig.parameters or "context" in sig.parameters
+
+            # Store tool info
+            self._tools[tool_name] = {
+                "handler": func,
+                "description": tool_desc,
+                "wants_context": wants_context,
+                "prompts": prompts or [],
+                "parameters": self._extract_parameters(func)
+            }
+
+            logger.debug(f"Registered tool: {tool_name}")
+            return func
+
+        return decorator
+
+    def resource(
+        self,
+        uri_pattern: str,
+        description: str = None
+    ) -> Callable:
+        """
+        Decorator to register a resource (read-only).
+
+        Args:
+            uri_pattern: URI pattern (e.g., "config://settings", "file://{path}")
+            description: Resource description
+
+        Example:
+            @mcp.resource("config://app")
+            def get_config(ctx: Context) -> dict:
+                return {"theme": "dark"}
+        """
+        def decorator(func: Callable) -> Callable:
+            self._resources[uri_pattern] = {
+                "handler": func,
+                "description": description or func.__doc__ or "",
+                "uri_pattern": uri_pattern
+            }
+
+            logger.debug(f"Registered resource: {uri_pattern}")
+            return func
+
+        return decorator
+
+    def prompt(
+        self,
+        name: str = None,
+        description: str = None
+    ) -> Callable:
+        """
+        Decorator to register a prompt template.
+
+        Args:
+            name: Prompt name (defaults to function name)
+            description: Prompt description
+
+        Example:
+            @mcp.prompt(description="Generate a greeting")
+            def greeting(name: str) -> str:
+                return f"Hello, {name}! How can I help?"
+        """
+        def decorator(func: Callable) -> Callable:
+            prompt_name = name or func.__name__
+
+            self._prompts[prompt_name] = {
+                "handler": func,
+                "description": description or func.__doc__ or "",
+                "parameters": self._extract_parameters(func)
+            }
+
+            logger.debug(f"Registered prompt: {prompt_name}")
+            return func
+
+        return decorator
+
+    def run(self, transport: str = None, port: int = 8080) -> None:
+        """
+        Start the MCP server.
+
+        Args:
+            transport: How external (non-MACAW) MCP clients may reach this server.
+                None    - MACAW mesh only. Tools are reachable via invoke_tool()
+                          from MACAW-aware clients. This is the default.
+                "stdio" - also serve native MCP clients over stdin/stdout.
+                          NOTE: stdout becomes JSON-RPC; log to stderr instead.
+                "http"  - also serve native MCP clients at http://host:port/mcp
+            port: Bind port for transport="http"
+
+        Either way every invocation goes through the MACAW PEP: policy, signing,
+        and audit are identical on both paths.
+        """
+        asyncio.run(self._run_async(transport, port))
+
+    async def _run_async(self, transport: str = None, port: int = 8080) -> None:
+        """Async server run."""
+        try:
+            # Initialize MACAWClient with registered tools
+            self._init_macaw_client()
+
+            logger.info(f"SecureMCP server '{self.name}' started")
+            logger.info(f"  Agent ID: {self.agent_id}")
+            logger.info(f"  Tools: {list(self._tools.keys())}")
+            logger.info(f"  Resources: {list(self._resources.keys())}")
+            logger.info(f"  Prompts: {list(self._prompts.keys())}")
+
+            if transport is None:
+                # MACAW mesh only. The listener lives inside MACAWClient (started by
+                # register()); this loop just keeps the process alive.
+                while True:
+                    await asyncio.sleep(1)
+            else:
+                from ._serve import serve
+                await serve(
+                    self.name,
+                    self.version,
+                    registry=self._client,
+                    target_agent=self.agent_id,
+                    prefix="",
+                    transport=transport,
+                    port=port,
+                )
+
+        except KeyboardInterrupt:
+            logger.info("Shutting down...")
+        finally:
+            if self._client:
+                self._client.unregister()
+                logger.info(f"Unregistered agent: {self.agent_id}")
+
+    def _init_macaw_client(self) -> None:
+        """Initialize MACAWClient with all registered tools."""
+        try:
+            from macaw_client import MACAWClient
+
+            # Collect unified tools config (handler + metadata together)
+            tools_config = {}
+            tool_names = []
+
+            for tool_name, tool_info in self._tools.items():
+                tool_names.append(tool_name)
+                # Unified format: handler + metadata together
+                tools_config[tool_name] = {
+                    "handler": self._create_handler(
+                        tool_name,
+                        tool_info["handler"],
+                        tool_info["wants_context"]
+                    ),
+                    "description": tool_info["description"],
+                    "prompts": tool_info.get("prompts", []),
+                    # Schema lives on the MACAW registry so the MCP endpoint can
+                    # answer tools/list from the same place policy is enforced.
+                    "metadata": {"schema": _json_schema(tool_info["parameters"])}
+                }
+
+            # Register resources as tools with resource: prefix
+            for uri_pattern, res_info in self._resources.items():
+                res_tool_name = f"resource:{uri_pattern}"
+                tool_names.append(res_tool_name)
+                tools_config[res_tool_name] = {
+                    "handler": self._create_handler(
+                        uri_pattern,
+                        res_info["handler"],
+                        True  # Resources always get context
+                    ),
+                    "description": res_info["description"],
+                    "read_only": True
+                }
+
+            # Register prompts as tools with prompt: prefix
+            for prompt_name, prompt_info in self._prompts.items():
+                prompt_tool_name = f"prompt:{prompt_name}"
+                tool_names.append(prompt_tool_name)
+                tools_config[prompt_tool_name] = {
+                    "handler": self._create_handler(
+                        prompt_name,
+                        prompt_info["handler"],
+                        False
+                    ),
+                    "description": prompt_info["description"]
+                }
+
+            # Build intent policy
+            if self._custom_intent_policy:
+                intent_policy = self._custom_intent_policy
+            else:
+                intent_policy = {
+                    "provided_capabilities": [{
+                        "type": "tools",
+                        "tools": tool_names,
+                        "description": f"Tools provided by SecureMCP server {self.name}"
+                    }],
+                    "required_resources": [],
+                    "data_access": ["none"],
+                    "description": f"SecureMCP server: {self.name}"
+                }
+
+                # Add roots as MAPL resource declarations
+                if self._roots:
+                    intent_policy["roots"] = self._roots
+                    intent_policy["required_resources"] = [
+                        {
+                            "resource": f"file://{root}/*",
+                            "operations": ["read", "write", "list"],
+                            "justification": f"MCP root: {root}"
+                        }
+                        for root in self._roots
+                    ]
+                    intent_policy["data_access"] = ["filesystem"]
+
+            # Create MACAWClient with unified tools config
+            self._client = MACAWClient(
+                app_name=f"securemcp-{self.name}",
+                app_version=self.version,
+                intent_policy=intent_policy,
+                tools=tools_config,  # Unified: {name: {handler, description, ...}}
+                **self.extra_config
+            )
+
+            # Register with Local Agent
+            if not self._client.register():
+                raise RuntimeError("Failed to register with MACAW Local Agent")
+
+            self.agent_id = self._client.agent_id
+            logger.info(f"Registered as agent: {self.agent_id}")
+
+        except ImportError as e:
+            raise RuntimeError(f"MACAWClient not available: {e}")
+        except Exception as e:
+            raise RuntimeError(f"MACAW initialization failed: {e}")
+
+    def _create_handler(
+        self,
+        tool_name: str,
+        func: Callable,
+        wants_context: bool
+    ) -> Callable:
+        """Create a handler wrapper for MACAWClient."""
+        def handler(params: Dict[str, Any]) -> Any:
+            try:
+                if wants_context:
+                    # Get source_agent from MACAWClient's current request context
+                    # This is set during tool execution handling
+                    source_agent = ""
+                    if self._client and hasattr(self._client, 'current_request_context'):
+                        source_agent = self._client.current_request_context.get('source_agent', '')
+                        logger.debug(f"[SecureMCP] Tool '{tool_name}' invoked by source_agent: {source_agent}")
+
+                    # Create context for this invocation
+                    # Include roots in metadata for get_roots() access
+                    metadata = params.pop("_metadata", {})
+                    metadata["roots"] = self._roots
+                    ctx = Context(
+                        request_id=params.pop("_request_id", ""),
+                        tool_name=tool_name,
+                        source_agent=source_agent,  # For sampling/elicitation callbacks
+                        _client=self._client,
+                        _metadata=metadata
+                    )
+                    # Check parameter name (ctx or context)
+                    sig = inspect.signature(func)
+                    if "ctx" in sig.parameters:
+                        params["ctx"] = ctx
+                    elif "context" in sig.parameters:
+                        params["context"] = ctx
+
+                # Execute handler
+                result = func(**params)
+
+                # Handle async functions
+                if asyncio.iscoroutine(result):
+                    # Use asyncio.run() which creates a new event loop
+                    # This works in threads without an existing event loop
+                    result = asyncio.run(result)
+
+                # Wrap result in dict format expected by MACAWClient
+                if not isinstance(result, dict):
+                    result = {"result": result}
+
+                return result
+
+            except Exception as e:
+                logger.error(f"Handler error for {tool_name}: {e}")
+                raise
+
+        return handler
+
+    def _extract_parameters(self, func: Callable) -> Dict[str, Any]:
+        """Extract parameter info from function signature."""
+        sig = inspect.signature(func)
+        params = {}
+
+        for param_name, param in sig.parameters.items():
+            if param_name in ("self", "ctx", "context"):
+                continue
+
+            param_info = {
+                "required": param.default == inspect.Parameter.empty
+            }
+
+            # Map type annotations to JSON schema types
+            if param.annotation != inspect.Parameter.empty:
+                type_map = {
+                    str: "string",
+                    int: "integer",
+                    float: "number",
+                    bool: "boolean",
+                    dict: "object",
+                    list: "array"
+                }
+                param_info["type"] = type_map.get(param.annotation, "string")
+            else:
+                param_info["type"] = "string"
+
+            if param.default != inspect.Parameter.empty:
+                param_info["default"] = param.default
+
+            params[param_name] = param_info
+
+        return params
+
+    # Convenience methods for testing
+    async def call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
+        """
+        Call a tool programmatically (for testing).
+
+        Args:
+            tool_name: Name of tool to call
+            arguments: Tool arguments
+
+        Returns:
+            Tool result
+        """
+        if tool_name not in self._tools:
+            raise ValueError(f"Unknown tool: {tool_name}")
+
+        handler = self._tools[tool_name]["handler"]
+        result = handler(**arguments)
+
+        if asyncio.iscoroutine(result):
+            result = await result
+
+        return result