mcp-server-cli/src/mcp_server_cli/models.py

"""Pydantic models for MCP protocol messages and tool definitions."""

from typing import Any, Dict, List, Optional, Union
from enum import Enum
from pydantic import BaseModel, Field, field_validator
from datetime import datetime


class MCPMessageType(str, Enum):
    """MCP protocol message types."""

    REQUEST = "request"
    RESPONSE = "response"
    NOTIFICATION = "notification"
    RESULT = "result"


class MCPMethod(str, Enum):
    """MCP protocol methods."""

    INITIALIZE = "initialize"
    INITIALIZED = "initialized"
    TOOLS_LIST = "tools/list"
    TOOLS_CALL = "tools/call"
    RESOURCES_LIST = "resources/list"
    RESOURCES_READ = "resources/read"
    PROMPTS_LIST = "prompts/list"
    PROMPTS_GET = "prompts/get"


class MCPRequest(BaseModel):
    """MCP protocol request message."""

    jsonrpc: str = "2.0"
    id: Optional[Union[int, str]] = None
    method: MCPMethod
    params: Optional[Dict[str, Any]] = None


class MCPResponse(BaseModel):
    """MCP protocol response message."""

    jsonrpc: str = "2.0"
    id: Optional[Union[int, str]] = None
    result: Optional[Dict[str, Any]] = None
    error: Optional[Dict[str, Any]] = None


class MCPNotification(BaseModel):
    """MCP protocol notification message."""

    jsonrpc: str = "2.0"
    method: str
    params: Optional[Dict[str, Any]] = None


class InitializeParams(BaseModel):
    """Parameters for MCP initialize request."""

    protocol_version: str = "2024-11-05"
    capabilities: Dict[str, Any] = Field(default_factory=dict)
    client_info: Optional[Dict[str, str]] = None


class ToolParameter(BaseModel):
    """Schema for a tool parameter."""

    name: str
    type: str = "string"
    description: Optional[str] = None
    required: bool = False
    enum: Optional[List[str]] = None
    default: Optional[Any] = None
    properties: Optional[Dict[str, Any]] = None


class ToolSchema(BaseModel):
    """Schema for tool input validation."""

    type: str = "object"
    properties: Dict[str, ToolParameter] = Field(default_factory=dict)
    required: List[str] = Field(default_factory=list)


class ToolDefinition(BaseModel):
    """Definition of a tool that can be called via MCP."""

    name: str
    description: str
    input_schema: ToolSchema = Field(default_factory=ToolSchema)
    annotations: Optional[Dict[str, Any]] = None

    @field_validator("name")
    @classmethod
    def validate_name(cls, v: str) -> str:
        if not v.isidentifier():
            raise ValueError(f"Tool name must be a valid identifier: {v}")
        return v


class ToolsListResult(BaseModel):
    """Result of tools/list request."""

    tools: List[ToolDefinition] = Field(default_factory=list)


class ToolCallParams(BaseModel):
    """Parameters for tools/call request."""

    name: str
    arguments: Optional[Dict[str, Any]] = None


class ToolCallResult(BaseModel):
    """Result of a tool call."""

    content: List[Dict[str, Any]]
    is_error: bool = False
    error_message: Optional[str] = None


class ServerInfo(BaseModel):
    """Information about the MCP server."""

    name: str = "mcp-server-cli"
    version: str = "0.1.0"


class ServerCapabilities(BaseModel):
    """Server capabilities announcement."""

    tools: Dict[str, Any] = Field(default_factory=dict)
    resources: Dict[str, Any] = Field(default_factory=dict)
    prompts: Dict[str, Any] = Field(default_factory=dict)


class InitializeResult(BaseModel):
    """Result of initialize request."""

    protocol_version: str
    server_info: ServerInfo
    capabilities: ServerCapabilities


class ConfigVersion(BaseModel):
    """Configuration version info."""

    version: str = "1.0"
    last_updated: Optional[str] = None


class ToolConfig(BaseModel):
    """Configuration for a registered tool."""

    name: str
    source: str
    enabled: bool = True
    config: Dict[str, Any] = Field(default_factory=dict)


class ServerConfig(BaseModel):
    """Main server configuration."""

    host: str = "127.0.0.1"
    port: int = 3000
    log_level: str = "INFO"
    config_version: str = "1.0"

    class Config:
        protected_namespaces = ()


class LocalLLMConfig(BaseModel):
    """Configuration for local LLM provider."""

    enabled: bool = False
    base_url: str = "http://localhost:11434"
    model: str = "llama2"
    temperature: float = 0.7
    max_tokens: int = 2048
    timeout: int = 60


class SecurityConfig(BaseModel):
    """Security configuration for the server."""

    allowed_commands: List[str] = Field(default_factory=lambda: ["ls", "cat", "echo", "pwd", "git"])
    blocked_paths: List[str] = Field(default_factory=lambda: ["/etc", "/root", "/home/*/.ssh"])
    max_shell_timeout: int = 30
    require_confirmation: bool = False


class AppConfig(BaseModel):
    """Application configuration combining all settings."""

    server: ServerConfig = Field(default_factory=ServerConfig)
    llm: LocalLLMConfig = Field(default_factory=LocalLLMConfig)
    security: SecurityConfig = Field(default_factory=SecurityConfig)
    tools: List[ToolConfig] = Field(default_factory=list)