posh-mcp.psm1

<#
.SYNOPSIS
    PowerShell MCP (Model Context Protocol) Server Module.

.DESCRIPTION
    This module implements an MCP server that dynamically loads tools from a JSON configuration file.
    It uses PowerShell's reflection capabilities to generate tool schemas from cmdlet documentation.
#>


#region Private Functions

function ConvertTo-JsonSchemaType {
    <#
    .SYNOPSIS
        Converts PowerShell parameter type to JSON schema type.
    #>

    param(
        [Type]$Type
    )
    
    switch ($Type.FullName) {
        "System.String" { return "string" }
        "System.Int32" { return "integer" }
        "System.Int64" { return "integer" }
        "System.Double" { return "number" }
        "System.Single" { return "number" }
        "System.Decimal" { return "number" }
        "System.Boolean" { return "boolean" }
        "System.DateTime" { return "string" }
        "System.String[]" { return "array" }
        "System.Int32[]" { return "array" }
        default {
            if ($Type.IsArray) { return "array" }
            if ($Type.IsEnum) { return "string" }
            return "string"
        }
    }
}

function Get-CmdletToolSchema {
    <#
    .SYNOPSIS
        Generates MCP tool schema from a cmdlet using reflection.
    #>

    param(
        [string]$CmdletName
    )
    
    # Get command info
    $cmdInfo = Get-Command $CmdletName -ErrorAction Stop
    
    # Get help for the cmdlet
    $help = Get-Help $CmdletName -Full -ErrorAction SilentlyContinue
    
    # Get description from help
    $toolDescription = if ($help.Synopsis) {
        $help.Synopsis.Trim()
    } else {
        "Executes $CmdletName"
    }
    
    # Generate tool name from cmdlet name (Get-Something -> getSomething)
    $toolName = $CmdletName -replace '-', ''
    $toolName = $toolName.Substring(0, 1).ToLower() + $toolName.Substring(1)
    
    # Build parameter schema
    $properties = @{}
    $required = @()
    
    foreach ($param in $cmdInfo.Parameters.GetEnumerator()) {
        $paramName = $param.Key
        $paramInfo = $param.Value
        
        # Skip common parameters
        $commonParams = @('Verbose', 'Debug', 'ErrorAction', 'WarningAction', 'InformationAction', 
                         'ErrorVariable', 'WarningVariable', 'InformationVariable', 'OutVariable', 
                         'OutBuffer', 'PipelineVariable', 'WhatIf', 'Confirm')
        if ($paramName -in $commonParams) { continue }
        
        # Convert to camelCase for JSON
        $toolParamName = $paramName.Substring(0, 1).ToLower() + $paramName.Substring(1)
        
        # Get parameter description from help
        $paramDescription = ""
        if ($help.parameters.parameter) {
            $helpParam = $help.parameters.parameter | Where-Object { $_.name -eq $paramName }
            if ($helpParam -and $helpParam.description) {
                $paramDescription = ($helpParam.description | Out-String).Trim()
            }
        }
        
        # Build property schema
        $propSchema = @{
            type = ConvertTo-JsonSchemaType -Type $paramInfo.ParameterType
        }
        
        if ($paramDescription) {
            $propSchema.description = $paramDescription
        }
        
        $properties[$toolParamName] = $propSchema
        
        # Check if parameter is mandatory
        $mandatoryAttr = $paramInfo.Attributes | Where-Object { 
            $_ -is [System.Management.Automation.ParameterAttribute] -and $_.Mandatory 
        }
        if ($mandatoryAttr) {
            $required += $toolParamName
        }
    }
    
    return @{
        name = $toolName
        description = $toolDescription
        inputSchema = @{
            type = "object"
            properties = $properties
            required = $required
        }
    }
}

function Build-ToolRegistry {
    <#
    .SYNOPSIS
        Builds the complete tool registry from configuration.
    #>

    param(
        [object]$Config,
        [string]$BasePath
    )
    
    $registry = @{
        tools = @()
        cmdletMap = @{}
    }
    
    # Process modules
    foreach ($module in $Config.modules) {
        foreach ($cmdletName in $module.cmdlets) {
            try {
                $schema = Get-CmdletToolSchema -CmdletName $cmdletName
                
                $registry.tools += $schema
                $registry.cmdletMap[$schema.name] = $cmdletName
            }
            catch {
                Write-Warning "Failed to register tool ${cmdletName}: $_"
            }
        }
    }
    
    # Process scripts
    foreach ($script in $Config.scripts) {
        foreach ($cmdletName in $script.cmdlets) {
            try {
                $schema = Get-CmdletToolSchema -CmdletName $cmdletName
                
                $registry.tools += $schema
                $registry.cmdletMap[$schema.name] = $cmdletName
            }
            catch {
                Write-Warning "Failed to register tool ${cmdletName}: $_"
            }
        }
    }
    
    return $registry
}

function Invoke-McpTool {
    <#
    .SYNOPSIS
        Dynamically invokes a tool based on the registry.
    #>

    param(
        [string]$ToolName,
        [object]$Arguments,
        [hashtable]$Registry
    )
    
    $cmdletName = $Registry.cmdletMap[$ToolName]
    if (-not $cmdletName) {
        throw "Unknown tool: $ToolName"
    }
    
    # Build parameter hashtable
    $params = @{}
    
    # Map incoming arguments to cmdlet parameters
    if ($Arguments) {
        $cmdInfo = Get-Command $cmdletName
        
        foreach ($prop in $Arguments.PSObject.Properties) {
            $argName = $prop.Name
            $argValue = $prop.Value
            
            # Convert from camelCase to PascalCase
            $cmdletParamName = $argName.Substring(0, 1).ToUpper() + $argName.Substring(1)
            
            # Handle special type conversions
            $paramInfo = $cmdInfo.Parameters[$cmdletParamName]
            
            if ($paramInfo) {
                if ($paramInfo.ParameterType -eq [DateTime]) {
                    $argValue = [DateTime]::Parse($argValue)
                }
                elseif ($paramInfo.ParameterType -eq [switch]) {
                    $argValue = [bool]$argValue
                }
            }
            
            $params[$cmdletParamName] = $argValue
        }
    }
    
    # Invoke the cmdlet
    $result = & $cmdletName @params
    
    # Ensure result is array for consistent output
    if ($null -eq $result) {
        $result = @()
    } elseif ($result -isnot [System.Array]) {
        $result = @($result)
    }
    
    return $result
}

function Send-JsonRpcResponse {
    param($id, $result)
    $response = @{
        jsonrpc = "2.0"
        id = $id
        result = $result
    }
    $json = ($response | ConvertTo-Json -Depth 20 -Compress)
    [Console]::WriteLine($json)
    [Console]::Out.Flush()
}

function Send-JsonRpcError {
    param($id, $code, $message)
    $response = @{
        jsonrpc = "2.0"
        id = $id
        error = @{
            code = $code
            message = $message
        }
    }
    $json = ($response | ConvertTo-Json -Depth 20 -Compress)
    [Console]::WriteLine($json)
    [Console]::Out.Flush()
}

#endregion

#region Public Functions

function Start-PoshMcp {
    <#
    .SYNOPSIS
        Starts the PowerShell MCP server.
    
    .DESCRIPTION
        Starts the Model Context Protocol (MCP) server that exposes PowerShell cmdlets as tools.
        The server reads its configuration from a JSON file and dynamically loads modules and scripts.
    
    .PARAMETER ConfigPath
        Path to the MCP configuration JSON file. Defaults to mcp-config.json in the module directory.
    
    .EXAMPLE
        Start-PoshMcp
        Starts the MCP server with the default configuration.
    
    .EXAMPLE
        Start-PoshMcp -ConfigPath "C:\myconfig\custom-mcp-config.json"
        Starts the MCP server with a custom configuration file.
    #>

    [CmdletBinding()]
    param(
        [Parameter(Mandatory = $false)]
        [string]$ConfigPath
    )
    
    # Ensure output is UTF-8 without BOM
    $OutputEncoding = [Console]::OutputEncoding = [System.Text.UTF8Encoding]::new($false)
    
    # Determine config path
    if (-not $ConfigPath) {
        $ConfigPath = Join-Path $PSScriptRoot "mcp-config.json"
    }
    
    # Determine base path for relative paths in config
    $basePath = Split-Path $ConfigPath -Parent
    
    # Load configuration
    if (-not (Test-Path $ConfigPath)) {
        throw "Configuration file not found: $ConfigPath"
    }
    $config = Get-Content $ConfigPath -Raw | ConvertFrom-Json
    
    # Import modules
    foreach ($module in $config.modules) {
        $modulePath = if ([System.IO.Path]::IsPathRooted($module.path)) {
            $module.path
        } else {
            Join-Path $basePath $module.path
        }
        Import-Module $modulePath -Force -Global
    }
    
    # Dot-source scripts in global scope
    foreach ($script in $config.scripts) {
        $scriptPath = if ([System.IO.Path]::IsPathRooted($script.path)) {
            $script.path
        } else {
            Join-Path $basePath $script.path
        }
        . $scriptPath
    }
    
    # Build tool registry
    $toolRegistry = Build-ToolRegistry -Config $config -BasePath $basePath
    
    # Main STDIO loop - MCP uses JSON-RPC 2.0
    while ($true) {
        $line = [Console]::In.ReadLine()
        if ($null -eq $line) { break }
        if ([string]::IsNullOrWhiteSpace($line)) { continue }

        try {
            $req = $line | ConvertFrom-Json
            $method = $req.method
            $id = $req.id

            switch ($method) {
                "initialize" {
                    Send-JsonRpcResponse -id $id -result @{
                        protocolVersion = "2024-11-05"
                        capabilities = @{
                            tools = @{}
                        }
                        serverInfo = @{
                            name = $config.serverInfo.name
                            version = $config.serverInfo.version
                        }
                    }
                }

                "notifications/initialized" {
                    # This is a notification, no response needed
                }

                "tools/list" {
                    Send-JsonRpcResponse -id $id -result @{
                        tools = $toolRegistry.tools
                    }
                }

                "tools/call" {
                    $toolName = $req.params.name
                    $arguments = $req.params.arguments

                    try {
                        $result = Invoke-McpTool -ToolName $toolName -Arguments $arguments -Registry $toolRegistry
                        
                        Send-JsonRpcResponse -id $id -result @{
                            content = @(
                                @{
                                    type = "text"
                                    text = ($result | ConvertTo-Json -Depth 10)
                                }
                            )
                        }
                    }
                    catch {
                        Send-JsonRpcError -id $id -code -32601 -message "Tool error: $($_.Exception.Message)"
                    }
                }

                default {
                    if ($null -ne $id) {
                        Send-JsonRpcError -id $id -code -32601 -message "Method not found: $method"
                    }
                }
            }
        }
        catch {
            if ($null -ne $id) {
                Send-JsonRpcError -id $id -code -32603 -message "Internal error: $($_.Exception.Message)"
            }
        }
    }
}

#endregion

Export-ModuleMember -Function Start-PoshMcp