Imports System.Collections.Concurrent
Imports System.Globalization
Imports System.IO
Imports System.Linq
Imports System.Net
Imports System.Reflection
Imports System.Security.Cryptography
Imports System.Threading.Tasks
Imports System.Web.Script.Serialization
' Required for IRequiresSessionState if you uncomment it below
' Imports System.Web.SessionState
''' <summary>
''' Attribute for automatically registering commands in x64Dbg or MCP.
''' </summary>
<AttributeUsage(AttributeTargets.Method, AllowMultiple:=True, Inherited:=False)>
Public Class CommandAttribute
Inherits Attribute
' Backing field for read-only property assigned in constructor
Private ReadOnly _name As String
Public ReadOnly Property Name As String
Get
Return _name
End Get
End Property
Public Property MCPCmdDescription As String ' Description specifically for MCP clients.
' Default constructor might be used if Name is inferred from method name later
Public Sub New()
End Sub
Public Sub New(name As String)
_name = name
End Sub
End Class
Public Module CommandImplementations
<Command("Test.Echo", MCPCmdDescription:="Simple echo command for testing.")>
Public Function EchoCommand(message As String) As String
Debug.WriteLine($"EchoCommand called with: {message}")
Return $"You sent: {message}"
End Function
<Command("Test.Echo2", MCPCmdDescription:="Simple echo command for testing.")>
Public Function Echo2Command(message As String) As String
Debug.WriteLine($"Echo2Command called with: {message}")
Return $"You sent: {message}"
End Function
' Add other command methods here following the pattern
' Example structure for other methods:
' <Command(...)>
' Public Function SomeOtherCommand(param1 As String, Optional param2 As Integer = 0) As String
' Debug.WriteLine(...)
' ' Logic
' Return "Result"
' End Function
End Module
''' <summary>
''' Handles MCP protocol requests via Server-Sent Events.
''' Implements IHttpAsyncHandler for non-blocking SSE connections.
''' </summary>
Public Class McpHandler
Implements IHttpAsyncHandler ', IRequiresSessionState ' Add if ASP.NET Session state is needed
' --- Shared State ---
Private Shared ReadOnly _sseSessions As New ConcurrentDictionary(Of String, StreamWriter)()
Private Shared ReadOnly _commands As New Dictionary(Of String, MethodInfo)(StringComparer.OrdinalIgnoreCase)
' *** POINT THIS TO YOUR ACTUAL MODULE/CLASS CONTAINING [Command] METHODS ***
Private Shared ReadOnly _targetType As Type = GetType(CommandImplementations)
Private Shared ReadOnly _jsonSerializer As New JavaScriptSerializer()
' Static constructor for one-time initialization
Shared Sub New()
RegisterCommands()
End Sub
''' <summary>
''' Discovers and registers methods marked with [CommandAttribute] suitable for MCP.
''' </summary>
Private Shared Sub RegisterCommands()
SyncLock _commands ' Lock during registration
Debug.WriteLine("Registering MCP commands...")
_commands.Clear() ' Ensure clean slate
' Reflect and register [Command] methods from the target type
For Each method In _targetType.GetMethods(BindingFlags.Static Or BindingFlags.Public Or BindingFlags.NonPublic)
' Process all Command attributes on a single method
For Each attr In method.GetCustomAttributes(Of CommandAttribute)()
' Determine the command name (use attribute Name if provided, else method name)
Dim commandName = If(attr?.Name, method.Name) ' VB equivalent of ??
' Basic validation
If String.IsNullOrWhiteSpace(commandName) Then
Debug.WriteLine($"Skipping registration for method '{method.Name}' due to missing command name in attribute.")
Continue For ' Skip this attribute
End If
If _commands.ContainsKey(commandName) Then
' Prevent duplicate *names*.
Debug.WriteLine($"Warning: Duplicate MCP command name '{commandName}' detected for method '{method.Name}'. Skipping duplicate registration.")
Continue For ' Skip this attribute
End If
Debug.WriteLine($"Registering MCP command: '{commandName}' for method '{method.Name}'")
_commands.Add(commandName, method) ' Use Add for clarity after ContainsKey check
Next ' attr
Next ' method
Debug.WriteLine($"Registered {_commands.Count} MCP commands.")
End SyncLock ' End lock
End Sub
' --- IHttpHandler Implementation ---
''' <summary>
''' Indicates if a handler instance can be reused. True for better performance
''' as state is managed safely in static fields.
''' </summary>
Public ReadOnly Property IsReusable As Boolean Implements IHttpHandler.IsReusable
Get
Return True
End Get
End Property
''' <summary>Synchronous ProcessRequest (Not used by IHttpAsyncHandler)</summary>
Public Sub ProcessRequest(context As HttpContext) Implements IHttpHandler.ProcessRequest
' This should not be called when implementing IHttpAsyncHandler
Throw New InvalidOperationException("ProcessRequest should not be called on an IHttpAsyncHandler.")
End Sub
' --- IHttpAsyncHandler Implementation ---
Public Function BeginProcessRequest(context As HttpContext, cb As AsyncCallback, extraData As Object) As IAsyncResult Implements IHttpAsyncHandler.BeginProcessRequest
' Wrap the asynchronous processing logic in a Task
Dim processingTask As Task = ProcessRequestAsync(context)
' Create a TaskCompletionSource to bridge Task and the APM pattern (Begin/End)
Dim tcs As New TaskCompletionSource(Of Object)(extraData)
' When the task completes, call the callback and signal completion
processingTask.ContinueWith(
Sub(t As Task)
If t.IsFaulted Then
' Propagate the exception (will be caught by EndProcessRequest)
Dim ex = If(t.Exception.InnerExceptions.Count = 1, t.Exception.InnerException, t.Exception)
tcs.TrySetException(ex)
ElseIf t.IsCanceled Then
tcs.TrySetCanceled()
Else
' Signal successful completion
tcs.TrySetResult(Nothing)
End If
' Invoke the ASP.NET callback if provided
cb?.Invoke(tcs.Task)
End Sub, TaskScheduler.Default) ' Use default scheduler is usually fine here
Return tcs.Task ' Return the Task which acts as IAsyncResult
End Function
Public Sub EndProcessRequest(result As IAsyncResult) Implements IHttpAsyncHandler.EndProcessRequest
Dim task = DirectCast(result, Task)
Try
' Await the task to force any exceptions to be re-thrown on the ASP.NET thread.
' This ensures ASP.NET's error handling pipeline is triggered correctly.
task.GetAwaiter().GetResult()
Catch ex As Exception
' Log the exception if needed, but ASP.NET will typically handle it now.
Debug.WriteLine($"Exception caught in EndProcessRequest: {ex}")
' Potentially log to a file or monitoring system here.
' Re-throwing is often implicit because GetResult() does it.
' Throw ' Usually not needed here as GetResult re-throws.
Finally
' Dispose the task associated with the TaskCompletionSource
task.Dispose()
End Try
End Sub
' --- Core Async Request Processing Logic ---
Private Async Function ProcessRequestAsync(context As HttpContext) As Task
Dim request As HttpRequest = context.Request
Dim response As HttpResponse = context.Response
Dim method As String = request.HttpMethod
Dim path As String = request.Url.AbsolutePath ' Path relative to the site root
Dim handlerPath As String = request.FilePath ' Path of the .ashx handler itself
Debug.WriteLine($"MCP Handler: Received {method} request for {path} (Handler: {handlerPath})")
' Determine the part of the path *after* the handler itself
' Example: /MyApp/McpHandler.ashx/sse -> actionPath = "sse"
Dim actionPath As String = String.Empty
If path.StartsWith(handlerPath, StringComparison.OrdinalIgnoreCase) Then
actionPath = path.Substring(handlerPath.Length).TrimStart("/"c).ToLowerInvariant()
End If
Debug.WriteLine($"MCP Handler: Determined Action Path: '{actionPath}'")
' --- Default Headers ---
' Disable caching for all MCP responses by default
response.Cache.SetCacheability(HttpCacheability.NoCache)
response.Cache.SetNoStore()
response.AppendHeader("X-Content-Type-Options", "nosniff") ' Security header
Try
If method = "GET" Then
If actionPath.EndsWith("sse") OrElse actionPath.EndsWith("sse/") Then
' HandleSseHandshake will keep the connection open
Await HandleSseHandshake(context)
' IMPORTANT: Don't do anything else with the 'response' object after this
' for SSE connections, as HandleSseHandshake manages its lifetime.
ElseIf actionPath.EndsWith("discover") OrElse actionPath.EndsWith("discover/") OrElse actionPath.EndsWith("mcp") OrElse actionPath.EndsWith("mcp/") Then
' Simple synchronous GET for basic command list (legacy support?)
HandleDiscoverRequest(context)
ElseIf actionPath.EndsWith("ping") OrElse actionPath.EndsWith("ping/") Then ' Simple health check
response.ContentType = "text/plain"
response.Write("pong")
Else
SendErrorResponse(response, HttpStatusCode.NotFound, "Resource not found.")
End If
ElseIf method = "POST" Then
If actionPath.EndsWith("message") OrElse actionPath.EndsWith("message/") Then
' Check session validity *before* accepting the request
Dim sessionId As String = context.Request.QueryString("sessionId")
If String.IsNullOrWhiteSpace(sessionId) Then
SendErrorResponse(context.Response, HttpStatusCode.BadRequest, "Missing sessionId query parameter.")
Exit Function ' Stop processing
End If
If Not _sseSessions.ContainsKey(sessionId) Then
Debug.WriteLine($"Received POST for unknown/closed session: {sessionId}")
SendErrorResponse(context.Response, HttpStatusCode.BadRequest, "Invalid or expired session ID.")
Exit Function ' Stop processing
End If
' Session looks okay, send 202 Accepted immediately for the POST request
Await SendAcceptedResponse(context.Response)
' Process the message body asynchronously and push results via SSE
Await HandleMessagePostBodyProcessing(context, sessionId)
Else
SendErrorResponse(response, HttpStatusCode.NotFound, "Resource not found.")
End If
Else ' Method Not Allowed (e.g., PUT, DELETE)
response.AppendHeader("Allow", "GET, POST") ' Inform client of allowed methods
SendErrorResponse(response, HttpStatusCode.MethodNotAllowed, "Method not allowed.")
End If
Catch ex As Exception ' Catch unexpected errors during routing/initial processing
Debug.WriteLine($"!!! Unhandled exception in ProcessRequestAsync: {ex}")
' Try to send an error response if possible (headers might already be sent for SSE)
If response.IsClientConnected AndAlso Not response.HeadersWritten Then
Try
SendErrorResponse(response, HttpStatusCode.InternalServerError, $"Internal Server Error: {ex.Message}")
Catch
' Ignore errors trying to send error response '
End Try
End If
' Rethrow the exception so EndProcessRequest can log it or handle it.
Throw
End Try
End Function
' --- GET Handlers ---
''' <summary>
''' Handles legacy GET /discover or /mcp requests (sends basic command list).
''' </summary>
Private Sub HandleDiscoverRequest(context As HttpContext)
Dim response As HttpResponse = context.Response
Dim toolList = New List(Of Object)()
SyncLock _commands
' Filter commands suitable for MCP here as well, respecting DebugOnly
Dim isDebugging As Boolean = System.Diagnostics.Debugger.IsAttached
For Each kvp In _commands
Dim methodInfo As MethodInfo = kvp.Value
Dim attribute As CommandAttribute = methodInfo.GetCustomAttributes(Of CommandAttribute)() _
.FirstOrDefault(Function(attr) If(attr?.Name, methodInfo.Name) = kvp.Key)
' Apply basic filtering (could add DebugOnly check here too if needed for discover)
If attribute Is Nothing Then
Continue For
End If
toolList.Add(New With {.name = kvp.Key, .parameters = {"string()"}})
Next
End SyncLock
Dim jsonResponse = _jsonSerializer.Serialize(New With {
.jsonrpc = "2.0",
.id = CType(Nothing, String),
.result = toolList
})
response.ContentType = "application/json"
response.Write(jsonResponse)
' Request processing completes here for this synchronous handler part
End Sub
''' <summary>
''' Establishes the Server-Sent Events connection.
''' </summary>
Private Async Function HandleSseHandshake(context As HttpContext) As Task
Dim response As HttpResponse = context.Response
' --- Set SSE Headers ---
response.ContentType = "text/event-stream"
response.StatusCode = CInt(HttpStatusCode.OK)
response.BufferOutput = False ' Crucial: Send data immediately
response.AppendHeader("Connection", "keep-alive")
' Caching handled by default headers set in ProcessRequestAsync
response.Headers.Remove("Content-Encoding")
response.AppendHeader("Content-Encoding", "identity")
Dim sessionId As String = GenerateSessionId()
' Create a StreamWriter that doesn't close the underlying stream on Dispose
' Use UTF8 without BOM. Ensure it leaves the stream open. Auto-flush.
Dim writer = New StreamWriter(response.OutputStream, New UTF8Encoding(False), 1024, leaveOpen:=True) With {
.AutoFlush = True
}
If _sseSessions.TryAdd(sessionId, writer) Then
Debug.WriteLine($"SSE session started: {sessionId}")
Try
' --- Send SSE Handshake ---
' Send the required 'endpoint' event with the message path
' Ensure the path is correct based on your handler mapping (e.g., relative or absolute)
Dim messagePath As String = VirtualPathUtility.ToAbsolute($"~/McpHandler.ashx/message?sessionId={sessionId}") ' Example using virtual path
' Or construct manually if needed: string messagePath = $"{context.Request.ApplicationPath.TrimEnd('/')}/McpHandler.ashx/message?sessionId={sessionId}";
Debug.WriteLine($"SSE Handshake: Sending endpoint event with data: {messagePath}")
Await writer.WriteAsync($"event: endpoint" & vbCrLf)
Await writer.WriteAsync($"data: {messagePath}" & vbCrLf & vbCrLf)
Await writer.FlushAsync() ' Ensure handshake is sent
' --- Keep Connection Alive ---
' Use TaskCompletionSource signaled by ClientDisconnectedToken
Dim tcs = New TaskCompletionSource(Of Boolean)()
Using registration = context.Response.ClientDisconnectedToken.Register(Sub()
Debug.WriteLine($"SSE client disconnected callback triggered: {sessionId}")
tcs.TrySetResult(True) ' Signal completion
End Sub)
Await tcs.Task ' Wait here until the client disconnects
Debug.WriteLine($"SSE client disconnected task completed: {sessionId}")
End Using
Catch __unusedObjectDisposedException1__ As ObjectDisposedException
' This might happen if the connection closes very quickly or due to an error elsewhere.
Debug.WriteLine($"SSE stream/writer was disposed during handshake/wait for session {sessionId}.")
Catch hex As HttpException When hex.ErrorCode = -2147023667 OrElse TypeOf hex.InnerException Is ObjectDisposedException ' Client disconnected
Debug.WriteLine($"SSE connection aborted by client (HttpException) during handshake/wait: {sessionId}")
Catch ex As Exception
Debug.WriteLine($"Unexpected error during SSE stream wait for {sessionId}: {ex}")
Finally
Debug.WriteLine($"Cleaning up SSE session after wait/disconnect: {sessionId}")
CleanupSseSession(sessionId)
End Try
Else
' Extremely unlikely with GenerateSessionId, but handle defensively
Debug.WriteLine($"FATAL: Failed to add SSE session {sessionId} to dictionary.")
CleanupSseSession(sessionId) ' Ensure writer is disposed if created
' Cannot send error via SSE; the handshake likely failed. Client will timeout.
End If
End Function
' --- POST Handler Logic ---
''' <summary>
''' Helper to send the 202 Accepted response for the initial POST.
''' </summary>
Private Async Function SendAcceptedResponse(response As HttpResponse) As Task
If Not response.IsClientConnected OrElse response.HeadersWritten Then Exit Function
response.StatusCode = CInt(HttpStatusCode.Accepted)
response.ContentType = "text/plain; charset=utf-8"
Try
Await response.Output.WriteAsync("Accepted")
Await response.FlushAsync()
Debug.WriteLine("Sent 202 Accepted response.")
Catch hex As HttpException When hex.ErrorCode = -2147023667 ' Client disconnected
Debug.WriteLine("Client disconnected before 202 Accepted could be fully sent.")
Catch ex As Exception
Debug.WriteLine($"Error sending 202 Accepted: {ex.Message}")
End Try
End Function
' --- Within Class McpHandler ---
''' <summary>
''' Reads and processes the body of the POST /message request after 202 Accepted has been sent.
''' </summary>
Private Async Function HandleMessagePostBodyProcessing(context As HttpContext, sessionId As String) As Task
Dim request As HttpRequest = context.Request
Dim json As Dictionary(Of String, Object) = Nothing
Dim parseErrorOccurred As Boolean = False
Dim invalidRequestStructure As Boolean = False
Dim processingException As Exception = Nothing
Dim idForError As Object = Nothing ' To store ID if request structure is bad
Try ' Outer try for general processing errors
Try ' Inner try specifically for JSON parsing
Dim requestBody As String
Using reader = New StreamReader(request.InputStream, request.ContentEncoding)
requestBody = Await reader.ReadToEndAsync()
End Using
Debug.WriteLine($"POST /message Body for Session {sessionId}: {requestBody}")
If String.IsNullOrWhiteSpace(requestBody) Then
Debug.WriteLine($"Warning: Empty body received for POST /message?sessionId={sessionId}")
Exit Function ' Nothing to process
End If
json = _jsonSerializer.Deserialize(Of Dictionary(Of String, Object))(requestBody)
Catch ex As Exception
Debug.WriteLine($"JSON Deserialization Error for Session {sessionId}: {ex.Message}")
parseErrorOccurred = True ' Flag the error
' Cannot reliably get ID here, will send error without it below
End Try
' --- Exit if parsing failed ---
If parseErrorOccurred Then
Await SendJsonRpcErrorAsync(sessionId, Nothing, -32700, "Parse error: Invalid JSON received.")
Exit Function
End If
' --- Validate JSON structure --- Init JSON message should not contain a key
If json Is Nothing OrElse Not json.ContainsKey("method") Then 'OrElse Not json.ContainsKey("id")
Debug.WriteLine($"Invalid JSON RPC structure received for Session {sessionId}")
invalidRequestStructure = True ' Flag the error
If json IsNot Nothing Then
json.TryGetValue("id", idForError) ' Try to get ID for error message
End If
Else
' --- Structure is valid, proceed to process ---
Dim methodName As String = Convert.ToString(json("method"))
Dim id As Object ' = json("id")
json.TryGetValue("id", id)
Await ProcessJsonRpcRequest(sessionId, id, methodName, json)
End If
Catch ex As Exception
' Catch errors during ProcessJsonRpcRequest or other steps
Debug.WriteLine($"Error processing POST /message body for session {sessionId}: {ex}")
processingException = ex ' Store exception
If json IsNot Nothing Then json.TryGetValue("id", idForError) ' Try get ID for error
End Try
' --- Send errors *after* Try blocks ---
If invalidRequestStructure Then
Await SendJsonRpcErrorAsync(sessionId, idForError, -32600, "Invalid Request: Missing method or id.")
ElseIf processingException IsNot Nothing Then
Await SendJsonRpcErrorAsync(sessionId, idForError, -32000, $"Internal Server Error processing message: {processingException.Message}")
End If
' Note: Parse error is handled and awaited earlier
End Function
''' <summary>Handles the 'initialize' request.</summary>
Private Async Function HandleInitialize(sessionId As String, id As Object, requestJson As Dictionary(Of String, Object)) As Task
Dim serverInfo = New With {.name = "ASPMCPHandler-VB", .version = "1.0.1"}
' CORRECTED: Use an empty dictionary for an empty JSON object
Dim capabilities = New With {.tools = New Dictionary(Of String, Object)()}
Dim responsePayload = New With {
.protocolVersion = "2024-11-05",
.capabilities = capabilities,
.serverInfo = serverInfo,
.instructions = "Welcome to the VB.NET MCP Handler!"
}
Await SendJsonRpcResultAsync(sessionId, id, responsePayload)
End Function
''' <summary>
''' Routes incoming JSON RPC requests (received via POST) to appropriate handlers.
''' Results/errors are sent back via the SSE stream associated with the sessionId.
''' </summary>
Private Async Function ProcessJsonRpcRequest(sessionId As String, id As Object, methodName As String, requestJson As Dictionary(Of String, Object)) As Task
' Double-check session is still valid before processing method
If Not _sseSessions.ContainsKey(sessionId) Then
Debug.WriteLine($"Error: Attempted to process RPC method '{methodName}' for unknown/closed session: {sessionId}")
Exit Function
End If
Debug.WriteLine($"Processing RPC Method '{methodName}' for Session {sessionId}, ID: {id}")
Select Case methodName
Case "initialize"
Await HandleInitialize(sessionId, id, requestJson)
Case "notifications/initialized"
Debug.WriteLine($"Received notifications/initialized for Session {sessionId}.")
' No response needed.
Case "tools/list"
Await HandleToolsList(sessionId, id)
Case "tools/call"
Await HandleToolCall(sessionId, id, requestJson)
Case "rpc.discover" ' Legacy support
Debug.WriteLine($"Warning: Received deprecated 'rpc.discover' for Session {sessionId}. Handling like 'tools/list'.")
Await HandleToolsList(sessionId, id)
Case Else
Debug.WriteLine($"Unknown method '{methodName}' received for Session {sessionId}")
Await SendJsonRpcErrorAsync(sessionId, id, -32601, $"Method not found: {methodName}")
End Select
End Function
''' <summary>Handles the 'tools/list' request.</summary>
Private Async Function HandleToolsList(sessionId As String, id As Object) As Task
Dim toolsList = New List(Of Object)()
Dim isDebugging As Boolean = System.Diagnostics.Debugger.IsAttached ' Check if debugger attached to IIS process
Debug.WriteLine($"HandleToolsList called. Debugger Attached: {isDebugging}")
SyncLock _commands ' Lock static commands dictionary while reading
For Each kvp In _commands
Dim commandName As String = kvp.Key
Dim methodInfo As MethodInfo = kvp.Value
' Find the specific attribute associated with this command name registration
Dim attribute As CommandAttribute = methodInfo.GetCustomAttributes(Of CommandAttribute)() _
.FirstOrDefault(Function(attr) If(attr?.Name, methodInfo.Name) = commandName)
' --- Filtering ---
If attribute Is Nothing Then Continue For ' Should be pre-filtered, but double-check
' --- Build Schema ---
Dim parameters = methodInfo.GetParameters()
Dim properties = New Dictionary(Of String, Object)()
Dim required = New List(Of String)()
For Each param In parameters
Dim paramName As String = param.Name
Dim paramType As String = GetJsonSchemaType(param.ParameterType)
Dim paramDescription As String = GetParameterDescription(commandName, paramName)
properties(paramName) = New With {.type = paramType, .description = paramDescription}
If Not param.IsOptional Then required.Add(paramName)
Next
toolsList.Add(New With {
.name = commandName,
.description = If(attribute.MCPCmdDescription, $"Executes the {commandName} command."),
.inputSchema = New With {
.title = commandName,
.description = If(attribute.MCPCmdDescription, $"Input schema for {commandName}."),
.type = "object",
.properties = properties,
.required = required.ToArray()
}
})
Next ' End foreach command
End SyncLock ' End lock
' --- Add Built-in Tools ---
toolsList.Add(New With {
.name = "Echo",
.description = "Echoes the input back.",
.inputSchema = New With {
.type = "object",
.properties = New With {.message = New With {.type = "string", .description = "Message to echo."}},
.required = {"message"}
}
})
' --- Send Response ---
Await SendJsonRpcResultAsync(sessionId, id, New With {.tools = toolsList.ToArray()})
End Function
''' <summary>Handles the 'tools/call' request.</summary>
Private Async Function HandleToolCall(sessionId As String, id As Object, requestJson As Dictionary(Of String, Object)) As Task
' --- Parse Input ---
Dim paramsObj As Object = Nothing
Dim toolCallParams As Dictionary(Of String, Object) = Nothing
If Not requestJson.TryGetValue("params", paramsObj) OrElse Not TypeOf paramsObj Is Dictionary(Of String, Object) Then
Await SendJsonRpcErrorAsync(sessionId, id, -32602, "Invalid params for tools/call")
Exit Function
End If
toolCallParams = DirectCast(paramsObj, Dictionary(Of String, Object))
Dim toolNameObj As Object = Nothing
Dim argumentsObj As Object = Nothing
Dim toolName As String = Nothing
If Not toolCallParams.TryGetValue("name", toolNameObj) OrElse Not TypeOf toolNameObj Is String Then
Await SendJsonRpcErrorAsync(sessionId, id, -32602, "Missing/invalid tool name")
Exit Function
End If
toolName = DirectCast(toolNameObj, String)
If String.IsNullOrEmpty(toolName) Then ' Extra check
Await SendJsonRpcErrorAsync(sessionId, id, -32602, "Missing/invalid tool name")
Exit Function
End If
toolCallParams.TryGetValue("arguments", argumentsObj)
Dim arguments As Dictionary(Of String, Object) = If(TryCast(argumentsObj, Dictionary(Of String, Object)), New Dictionary(Of String, Object)())
Debug.WriteLine($"Attempting tool call '{toolName}' for session {sessionId}")
' --- Handle Built-in Echo ---
If toolName.Equals("Echo", StringComparison.OrdinalIgnoreCase) Then
Dim message As String = "N/A"
Dim msgArg As Object = Nothing
If arguments.TryGetValue("message", msgArg) Then message = Convert.ToString(msgArg)
Await SendToolCallResultAsync(sessionId, id, {New With {.type = "text", .text = $"Echo response: {message}"}}, isError:=False)
Exit Function
End If
' --- Find and Prepare Registered Command ---
Dim methodInfo As MethodInfo = Nothing
Dim isError As Boolean = False
Dim resultText As String
If _commands.TryGetValue(toolName, methodInfo) Then
Dim attribute As CommandAttribute = methodInfo.GetCustomAttributes(Of CommandAttribute)().FirstOrDefault(Function(attr) If(attr?.Name, methodInfo.Name) = toolName)
Dim isDebugging As Boolean = System.Diagnostics.Debugger.IsAttached
' --- Pre-execution Checks ---
If attribute Is Nothing Then
resultText = $"Error: Command '{toolName}' not available or not permitted in this context."
isError = True
Debug.WriteLine($"Denied execution of '{toolName}': Attribute mismatch/missing.")
Else
' --- Parameter Binding and Invocation ---
Try
Dim methodParams = methodInfo.GetParameters()
Dim invokeArgs = New Object(methodParams.Length - 1) {} ' VB arrays are 0-based
For i As Integer = 0 To methodParams.Length - 1
Dim param = methodParams(i)
Dim argValue As Object = Nothing
If arguments.TryGetValue(param.Name, argValue) Then
Try
invokeArgs(i) = ConvertArgumentType(argValue, param.ParameterType, param.Name) ' Use helper
Catch convEx As Exception
Throw New ArgumentException($"Cannot convert argument '{param.Name}'. Error: {convEx.Message}", convEx)
End Try
ElseIf param.IsOptional Then
invokeArgs(i) = param.DefaultValue
Else
Throw New ArgumentException($"Missing required argument: '{param.Name}'")
End If
Next
' *** Invoke the command method ***
Dim result As Object = methodInfo.Invoke(Nothing, invokeArgs) ' Assumes static methods in a Module
resultText = If(result?.ToString(), $"{toolName} executed successfully.")
isError = False
Catch tie As TargetInvocationException ' Exception inside the command
resultText = $"Error executing '{toolName}': {If(tie.InnerException?.Message, tie.Message)}"
isError = True
Debug.WriteLine($"Execution Error in {toolName}: {If(tie.InnerException, tie)}")
Catch ex As Exception ' Error during binding/setup
resultText = $"Error preparing/calling '{toolName}': {ex.Message}"
isError = True
Debug.WriteLine($"Binding/Invocation Error for {toolName}: {ex}")
End Try
End If
Else ' Command name not found
resultText = $"Tool '{toolName}' not found."
isError = True
End If
' --- Send Result via SSE ---
Dim toolContent = New With {.type = "text", .text = resultText}
Await SendToolCallResultAsync(sessionId, id, {toolContent}, isError)
End Function
''' <summary>Sends a structured event over the SSE stream.</summary>
Private Async Function SendSseEventAsync(sessionId As String, eventName As String, data As String) As Task
Dim writer As StreamWriter = Nothing
If _sseSessions.TryGetValue(sessionId, writer) Then
Try
Dim sb As New StringBuilder()
If Not String.IsNullOrEmpty(eventName) Then
sb.Append($"event: {eventName}" & vbLf)
End If
' Handle multi-line data correctly
Using reader = New StringReader(If(data, String.Empty))
Dim line As String = Await reader.ReadLineAsync()
While line IsNot Nothing
sb.Append($"data: {line}" & vbLf)
line = Await reader.ReadLineAsync()
End While
End Using
sb.Append(vbLf) ' End of message marker
Await writer.WriteAsync(sb.ToString())
' Flush is handled by writer's AutoFlush = True
Catch ex As Exception When TypeOf ex Is ObjectDisposedException OrElse TypeOf ex Is IOException OrElse (TypeOf ex Is HttpException AndAlso DirectCast(ex, HttpException).ErrorCode = -2147023667)
Debug.WriteLine($"SSE Write Error/Disconnect for session {sessionId}, cleaning up. Error: {ex.GetType().Name} - {ex.Message}")
CleanupSseSession(sessionId) ' Remove session on write failure
Catch ex As Exception ' Catch other unexpected errors during write
Debug.WriteLine($"Unexpected SSE Write Error for session {sessionId}: {ex}")
CleanupSseSession(sessionId) ' Be aggressive on cleanup
End Try
End If
' else: Session already removed, do nothing.
End Function
''' <summary>Sends a standard JSON RPC result message via SSE.</summary>
Private Function SendJsonRpcResultAsync(sessionId As String, id As Object, result As Object) As Task
Dim response = New With {.jsonrpc = "2.0", .id = id, .result = result}
Dim jsonData As String = _jsonSerializer.Serialize(response)
Return SendSseEventAsync(sessionId, Nothing, jsonData) ' 'message' event type is default
End Function
''' <summary>Sends a standard JSON RPC error message via SSE.</summary>
Private Function SendJsonRpcErrorAsync(sessionId As String, id As Object, code As Integer, message As String, Optional data As Object = Nothing) As Task
Dim errorPayload = New With {.code = code, .message = message, .data = data}
Dim response = New With {.jsonrpc = "2.0", .id = id, .error = errorPayload}
Dim jsonData As String = _jsonSerializer.Serialize(response)
Return SendSseEventAsync(sessionId, Nothing, jsonData) ' 'message' event type is default
End Function
''' <summary>Sends a structured tools/call result message via SSE.</summary>
Private Function SendToolCallResultAsync(sessionId As String, id As Object, content As Object(), isError As Boolean) As Task
Dim resultPayload = New With {.content = content, .isError = isError}
Return SendJsonRpcResultAsync(sessionId, id, resultPayload)
End Function
''' <summary>Removes an SSE session and disposes its writer.</summary>
Private Sub CleanupSseSession(sessionId As String)
Dim writer As StreamWriter = Nothing
If _sseSessions.TryRemove(sessionId, writer) Then
Debug.WriteLine($"Removed SSE session from dictionary: {sessionId}")
Try
writer?.Dispose() ' Dispose writer to release resources
Catch ex As Exception
Debug.WriteLine($"Error disposing SSE writer for {sessionId}: {ex.Message}")
End Try
Else
Debug.WriteLine($"Cleanup requested for session {sessionId}, but it was already removed.")
End If
End Sub
''' <summary>Generates a URL-safe random session ID.</summary>
Private Shared Function GenerateSessionId() As String
Using rng = RandomNumberGenerator.Create()
Dim randomBytes(15) As Byte ' 16 bytes = 128 bits
rng.GetBytes(randomBytes)
Return Convert.ToBase64String(randomBytes).TrimEnd("="c).Replace("+"c, "-"c).Replace("/"c, "_"c)
End Using
End Function
''' <summary>Sends a simple error response if headers haven't been written.</summary>
Private Shared Sub SendErrorResponse(response As HttpResponse, statusCode As HttpStatusCode, message As String)
If Not response.IsClientConnected OrElse response.HeadersWritten Then
Debug.WriteLine($"Cannot send error '{message}' - response already sent/client disconnected.")
Exit Sub
End If
response.StatusCode = CInt(statusCode)
response.ContentType = "text/plain; charset=utf-8"
response.TrySkipIisCustomErrors = True ' Prevent IIS custom error pages
response.Write(message)
' This implicitly ends the response for synchronous error cases.
End Sub
''' <summary>Converts a deserialized argument value to the required parameter type.</summary>
Private Shared Function ConvertArgumentType(argValue As Object, requiredType As Type, paramName As String) As Object
If argValue Is Nothing Then
If requiredType.IsClass OrElse Nullable.GetUnderlyingType(requiredType) IsNot Nothing Then Return Nothing
Throw New ArgumentNullException(paramName, $"Null provided for non-nullable parameter '{paramName}' of type {requiredType.Name}")
End If
' If type already matches (common for strings, bools)
If requiredType.IsInstanceOfType(argValue) Then Return argValue
' Handle common numeric conversions from JavaScriptSerializer (often gives Int32 or Decimal)
Select Case requiredType
Case GetType(Integer) : Return Convert.ToInt32(argValue)
Case GetType(Long) : Return Convert.ToInt64(argValue)
Case GetType(Short) : Return Convert.ToInt16(argValue)
Case GetType(Byte) : Return Convert.ToByte(argValue)
Case GetType(UInteger) : Return Convert.ToUInt32(argValue)
Case GetType(ULong) : Return Convert.ToUInt64(argValue)
Case GetType(UShort) : Return Convert.ToUInt16(argValue)
Case GetType(SByte) : Return Convert.ToSByte(argValue)
Case GetType(Single) : Return Convert.ToSingle(argValue)
Case GetType(Double) : Return Convert.ToDouble(argValue)
Case GetType(Decimal) : Return Convert.ToDecimal(argValue)
Case GetType(Boolean) : Return Convert.ToBoolean(argValue)
Case GetType(Guid) : Return Guid.Parse(argValue.ToString())
End Select
If requiredType.IsEnum Then Return System.Enum.Parse(requiredType, argValue.ToString(), ignoreCase:=True) ' Case-insensitive
' Handle arrays (simple case: assumes string array if target is string[])
If requiredType Is GetType(String()) AndAlso TypeOf argValue Is System.Collections.ArrayList Then
Dim list = DirectCast(argValue, System.Collections.ArrayList)
Return list.Cast(Of Object)().Select(Function(o) o?.ToString()).ToArray()
End If
' TODO: Add more robust array conversion if needed (e.g., arrays of numbers)
' Fallback for other types
Try
Return Convert.ChangeType(argValue, requiredType, CultureInfo.InvariantCulture)
Catch ex As Exception
Throw New InvalidCastException($"Cannot convert value '{argValue}' (type: {argValue.GetType().Name}) to required type '{requiredType.Name}' for parameter '{paramName}'.", ex)
End Try
End Function
''' <summary>Gets a JSON schema type name for a C# type.</summary>
Private Shared Function GetJsonSchemaType(type As Type) As String
If type Is Nothing Then Return "null"
Dim underlyingType As Type = Nullable.GetUnderlyingType(type)
type = If(underlyingType, type) ' Use underlying type if nullable
If type Is GetType(String) OrElse type Is GetType(Guid) OrElse type.IsEnum Then Return "string"
If type Is GetType(Integer) OrElse type Is GetType(Long) OrElse type Is GetType(Short) OrElse type Is GetType(Byte) OrElse
type Is GetType(UInteger) OrElse type Is GetType(ULong) OrElse type Is GetType(UShort) OrElse type Is GetType(SByte) Then Return "integer"
If type Is GetType(Single) OrElse type Is GetType(Double) OrElse type Is GetType(Decimal) Then Return "number"
If type Is GetType(Boolean) Then Return "boolean"
If type Is GetType(DateTime) OrElse type Is GetType(DateTimeOffset) Then Return "string" ' Could add "format": "date-time"
If type.IsArray OrElse GetType(System.Collections.IEnumerable).IsAssignableFrom(type) Then Return "array" ' Basic array/list check
Return "object" ' Default for classes, dictionaries, etc.
End Function
''' <summary>Gets a default description for common parameter names.</summary>
Private Shared Function GetParameterDescription(commandName As String, paramName As String) As String
Select Case paramName.ToLowerInvariant()
'Sample Params and Descriptions to return.
Case "address" : Return "Memory address (e.g., '0x1234ABCD' or module+offset)"
Case "value" : Return "Value to use"
Case "bytecount", "count", "size" : Return "Number of bytes or items"
Case "filepath", "filename" : Return "Path to a file (e.g., 'C:\temp\output.txt')" ' Note VB uses \ for path sep typically
Case "mode" : Return "Operating mode"
Case "bytestring", "bytes" : Return "Hexadecimal byte string (e.g., '90 F3 AA 00')"
Case "message" : Return "Text message or input"
Case Else : Return $"Parameter '{paramName}' for {commandName}"
End Select
End Function
End Class ' McpHandler
