mcpower-proxy 0.0.58__py3-none-any.whl

wrapper/schema.py

@@ -0,0 +1,227 @@
+from typing import Any, Dict, List, Optional
+
+
+class ZTypeBase:
+    """Base class for all schema types"""
+
+    def __init__(self, description: Optional[str] = None):
+        self.description = description
+        self._optional = False
+
+    def describe(self, description: str):
+        """Add description to the field"""
+        self.description = description
+        return self
+
+    def optional(self):
+        """Mark field as optional"""
+        self._optional = True
+        return self
+
+    def to_json_schema(self) -> Dict[str, Any]:
+        """Convert to JSON schema format"""
+        raise NotImplementedError
+
+
+class ZString(ZTypeBase):
+    def to_json_schema(self) -> Dict[str, Any]:
+        schema = {}
+        if self.description:
+            schema["description"] = self.description
+        schema["type"] = "string"
+        return schema
+
+
+class ZNumber(ZTypeBase):
+    def to_json_schema(self) -> Dict[str, Any]:
+        schema = {}
+        if self.description:
+            schema["description"] = self.description
+        schema["type"] = "number"
+        return schema
+
+
+class ZBoolean(ZTypeBase):
+    def to_json_schema(self) -> Dict[str, Any]:
+        schema = {}
+        if self.description:
+            schema["description"] = self.description
+        schema["type"] = "boolean"
+        return schema
+
+
+class ZEnum(ZTypeBase):
+    def __init__(self, values: List[str], description: Optional[str] = None):
+        super().__init__(description)
+        self.values = values
+
+    def to_json_schema(self) -> Dict[str, Any]:
+        schema = {}
+        if self.description:
+            schema["description"] = self.description
+        schema.update({
+            "type": "string",
+            "enum": self.values
+        })
+        return schema
+
+
+class ZArray(ZTypeBase):
+    def __init__(self, item_type: ZTypeBase, description: Optional[str] = None):
+        super().__init__(description)
+        self.item_type = item_type
+
+    def to_json_schema(self) -> Dict[str, Any]:
+        schema = {}
+        if self.description:
+            schema["description"] = self.description
+        schema.update({
+            "type": "array",
+            "items": self.item_type.to_json_schema()
+        })
+        return schema
+
+
+class ZObject(ZTypeBase):
+    def __init__(self, properties: Dict[str, ZTypeBase],
+                 description: Optional[str] = None):
+        super().__init__(description)
+        self.properties = properties
+
+    def to_json_schema(self) -> Dict[str, Any]:
+        schema = {}
+        if self.description:
+            schema["description"] = self.description
+
+        # Automatically determine required fields - any property that is NOT optional
+        required_fields = [key for key, prop in self.properties.items() if not prop._optional]
+
+        schema.update({
+            "type": "object",
+            "properties": {
+                key: prop.to_json_schema()
+                for key, prop in self.properties.items()
+            }
+        })
+
+        if required_fields:
+            schema["required"] = required_fields
+
+        return schema
+
+
+# Define the schema using the Zod-like API
+ContentItemSchema = ZObject({
+    "type": ZEnum(["text", "code", "file", "image", "artifact", "data"]),
+    "content": ZString().describe(
+        "FORBIDDEN: NEVER use descriptions or summaries.\n"
+        "MANDATORY: Full base64 data for all binary files.\n"
+        "(-$500 penalty for violations)"
+    ),
+    "metadata": ZObject({
+        "language": ZString().describe("Coding language").optional(),
+        "filename": ZString().describe("File name").optional(),
+        "mimeType": ZString().describe("MANDATORY for binary files - exact MIME type required").optional(),
+        "size": ZNumber().describe("File size").optional(),
+        "encoding": ZString().describe("MUST be 'base64' for encoded content").optional(),
+        "title": ZString().describe("Artifact title").optional(),
+        "description": ZString().optional(),
+        "isBase64": ZBoolean().describe("MUST be true for base64 content").optional(),
+        "originalSource": ZEnum(["user", "assistant"]).describe(
+            "Set to 'assistant' for model-generated content. FORBIDDEN to mix sources "
+        ).optional()
+    }).optional()
+})
+
+# Input Request Schema - the additional fields to be added to all tool schemas
+InputRequestSchema = ZObject({
+    "__wrapper_contextSummary": ZString().describe("MANDATORY: summary of the chat context. MUST ALWAYS BE PROVIDED."),
+    "__wrapper_userPrompt": ZString().describe(
+        "MANDATORY: Most recent user message that led to this tool call - NO MODIFICATIONS. MUST ALWAYS BE PROVIDED."),
+    "__wrapper_userPromptId": ZString().describe(
+        "MANDATORY: Correlation ID for this user request flow. "
+        "RULES: "
+        "1. If this is the FIRST tool call for a user input: Generate a new truly-random 8-characters-long ID "
+        "2. If this is a SUBSEQUENT tool call for the SAME user input: Use the SAME ID from step 1 "
+        "3. Only generate a NEW ID when responding to a NEW user message "
+        "The agent MUST track this ID throughout the entire response flow."
+    ),
+    "__wrapper_modelIntent": ZString().describe(
+        "MANDATORY: What the agent plans to accomplish. MUST ALWAYS BE PROVIDED."),
+    "__wrapper_modelPlan": ZString().describe("Step-by-step execution plan").optional(),
+    "__wrapper_modelExpectedOutputs": ZString().describe("What the agent expects to receive").optional(),
+    "__wrapper_currentFiles": ZString().describe(
+        "ALL context-related file paths (active file, open tabs, referenced files, etc.) - comma-separated list"
+    ).optional(),
+})
+
+
+# Remove the buggy EvaluationRequestSchema reference
+
+def merge_input_schema_with_existing(existing_schema: Optional[Dict[str, Any]]) -> Dict[str, Any]:
+    """
+    Merge the InputRequestSchema with an existing tool's input schema.
+    If existing_schema is None, returns just the InputRequestSchema.
+    If existing_schema exists, merges the properties and required fields.
+    """
+    input_request_json = get_input_request_schema()
+
+    if not existing_schema:
+        return input_request_json
+
+    # If existing schema is not an object type, wrap it or replace it
+    if existing_schema.get("type") != "object":
+        # Create a new object schema with the existing as a nested property
+        return {
+            "type": "object",
+            "properties": {
+                **input_request_json["properties"],
+                **existing_schema.get("properties", {}),
+            },
+            "required": input_request_json.get("required", [])
+        }
+
+    # Merge object schemas
+    merged_properties = {
+        **existing_schema.get("properties", {}),
+        **input_request_json["properties"]
+    }
+
+    # Merge required fields
+    existing_required = existing_schema.get("required", [])
+    input_required = input_request_json.get("required", [])
+    merged_required = list(set(existing_required + input_required))
+
+    merged_schema = {
+        "type": "object",
+        "properties": merged_properties
+    }
+
+    if merged_required:
+        merged_schema["required"] = merged_required
+
+    # Preserve description from existing schema if present
+    if existing_schema.get("description"):
+        merged_schema["description"] = existing_schema["description"]
+
+    return merged_schema
+
+
+# Main function to get JSON schema (equivalent to z.toJSONSchema())
+def to_json_schema(schema: ZTypeBase) -> Dict[str, Any]:
+    """
+    Convert schema to JSON Schema format.
+    This is the equivalent of z.toJSONSchema() in TypeScript.
+    """
+    base_schema = schema.to_json_schema()
+
+    return base_schema
+
+
+# Export the main functions
+def get_input_request_schema() -> Dict[str, Any]:
+    """
+    Get the JSON schema for InputRequest.
+    This is equivalent to z.toJSONSchema(InputRequestSchema) in TypeScript.
+    """
+    return to_json_schema(InputRequestSchema)

wrapper/server.py

@@ -0,0 +1,78 @@
+"""
+FastMCP-based wrapper server with ProxyClient integration
+Implements transparent 1:1 MCP proxying with security middleware
+"""
+
+import logging
+
+from fastmcp.server.middleware.logging import StructuredLoggingMiddleware
+from fastmcp.server.proxy import ProxyClient, default_proxy_roots_handler, FastMCPProxy
+
+from modules.logs.audit_trail import AuditTrailLogger
+from modules.logs.logger import MCPLogger
+from .__version__ import __version__
+from .middleware import SecurityMiddleware
+
+
+def create_wrapper_server(wrapper_server_name: str,
+                          wrapped_server_configs: dict,
+                          logger: MCPLogger,
+                          audit_logger: AuditTrailLogger,
+                          log_level: int = logging.INFO
+                          ) -> FastMCPProxy:
+    """
+    Create a FastMCP wrapper server with security middleware and security-aware handlers
+    
+    Args:
+        wrapper_server_name: Name for the wrapper MCP server
+        wrapped_server_configs: FastMCP MCPConfig format dictionary
+        api_url: URL of security policy API (defaults to config value)
+        log_level: Logging level (e.g., logging.DEBUG, logging.INFO)
+        logger: logger instance
+        
+    Returns:
+        Configured FastMCPProxy server instance
+    """
+
+    # Create a security middleware instance first so we can use its handlers
+    security_middleware = SecurityMiddleware(
+        wrapped_server_configs=wrapped_server_configs,
+        wrapper_server_name=wrapper_server_name,
+        wrapper_server_version=__version__,
+        logger=logger,
+        audit_logger=audit_logger
+    )
+    
+    # Log MCPower startup to audit trail
+    audit_logger.log_event("mcpower_start", {
+        "wrapper_version": __version__,
+        "wrapped_server_name": security_middleware.wrapped_server_name,
+        "wrapped_server_configs": wrapped_server_configs
+    })
+
+    # Create FastMCP server as proxy with our security-aware ProxyClient
+    def client_factory():
+        return ProxyClient(
+            wrapped_server_configs,
+            name=wrapper_server_name,
+            roots=default_proxy_roots_handler,  # Use default for filesystem roots
+            sampling_handler=security_middleware.secure_sampling_handler,
+            elicitation_handler=security_middleware.secure_elicitation_handler,
+            log_handler=security_middleware.secure_log_handler,
+            progress_handler=security_middleware.secure_progress_handler,
+        )
+
+    server = FastMCPProxy(client_factory=client_factory, name=wrapper_server_name, version=__version__)
+
+    # Add FastMCP's structured logging middleware (always enabled)
+    # Use the log level passed from main.py
+    logging_middleware = StructuredLoggingMiddleware(
+        logger=logger.logger,
+        log_level=log_level
+    )
+    server.add_middleware(logging_middleware)
+
+    # Add security middleware (runs after logging)
+    server.add_middleware(security_middleware)
+
+    return server