From 40ed102c36eb364bab47aa268886aaad0b28bd83 Mon Sep 17 00:00:00 2001
From: Nhan Nguyen <nguyenxuannhan407@gmail.com>
Date: Fri, 27 Feb 2026 20:22:10 +0700
Subject: [PATCH] Enhance McpDocumentTransformer to support JSON-RPC request
 and response schemas

---
 .../OpenApi/McpDocumentTransformer.cs         | 36 ++++++++++++++++---
 1 file changed, 32 insertions(+), 4 deletions(-)

diff --git a/shared/McpSamples.Shared/OpenApi/McpDocumentTransformer.cs b/shared/McpSamples.Shared/OpenApi/McpDocumentTransformer.cs
index 8beb9b07..9339f65a 100644
--- a/shared/McpSamples.Shared/OpenApi/McpDocumentTransformer.cs
+++ b/shared/McpSamples.Shared/OpenApi/McpDocumentTransformer.cs
@@ -1,3 +1,4 @@
+using System.Net.Mime;
 using System.Text.Json.Nodes;
 
 using McpSamples.Shared.Configurations;
@@ -6,6 +7,8 @@
 using Microsoft.AspNetCore.OpenApi;
 using Microsoft.OpenApi;
 
+using ModelContextProtocol.Protocol;
+
 namespace McpSamples.Shared.OpenApi;
 
 /// <summary>
@@ -16,7 +19,7 @@ namespace McpSamples.Shared.OpenApi;
 public class McpDocumentTransformer<T>(T appsettings, IHttpContextAccessor accessor) : IOpenApiDocumentTransformer where T : AppSettings, new()
 {
     /// <inheritdoc />
-    public Task TransformAsync(OpenApiDocument document, OpenApiDocumentTransformerContext context, CancellationToken cancellationToken)
+    public async Task TransformAsync(OpenApiDocument document, OpenApiDocumentTransformerContext context, CancellationToken cancellationToken)
     {
         document.Info = new OpenApiInfo
         {
@@ -33,6 +36,15 @@ public Task TransformAsync(OpenApiDocument document, OpenApiDocumentTransformerC
                     : "http://localhost:8080/"
             }
         ];
+
+        var jsonRpcResponse = await context.GetOrCreateSchemaAsync(typeof(JsonRpcResponse), cancellationToken: cancellationToken);
+
+        var jsonRpcRequest = await context.GetOrCreateSchemaAsync(typeof(JsonRpcRequest), cancellationToken: cancellationToken);
+
+        document.AddComponent(nameof(JsonRpcResponse), jsonRpcResponse);
+
+        document.AddComponent(nameof(JsonRpcRequest), jsonRpcRequest);
+
         var pathItem = new OpenApiPathItem();
         pathItem.AddOperation(HttpMethod.Post, new OpenApiOperation
         {
@@ -47,13 +59,29 @@ public Task TransformAsync(OpenApiDocument document, OpenApiDocumentTransformerC
                 ["200"] = new OpenApiResponse
                 {
                     Description = "Success",
+                    Content = new Dictionary<string, OpenApiMediaType>
+                    {
+                        [MediaTypeNames.Application.Json] = new()
+                        {
+                            Schema = new OpenApiSchemaReference(nameof(JsonRpcResponse), document),
+                        },
+                    },
                 }
-            }
+            },
+            RequestBody = new OpenApiRequestBody
+            {
+                Required = true,
+                Content = new Dictionary<string, OpenApiMediaType>
+                {
+                    [MediaTypeNames.Application.Json] = new()
+                    {
+                        Schema = new OpenApiSchemaReference(nameof(JsonRpcRequest), document),
+                    },
+                },
+            },
         });
 
         document.Paths ??= [];
         document.Paths.Add("/mcp", pathItem);
-
-        return Task.CompletedTask;
     }
 }
\ No newline at end of file