From 45dcfe07ca51487030c87d29c1539e294aeb8992 Mon Sep 17 00:00:00 2001
From: Stephen Toub <stoub@microsoft.com>
Date: Thu, 10 Apr 2025 10:51:41 -0400
Subject: [PATCH] Fix CreateMessageRule.Role to be of type Role (#264)

---
 src/ModelContextProtocol/Client/McpClientExtensions.cs      | 2 +-
 .../Protocol/Types/CreateMessageResult.cs                   | 2 +-
 src/ModelContextProtocol/Server/McpServerExtensions.cs      | 2 +-
 .../SseServerIntegrationTests.cs                            | 2 +-
 .../Utils/TestServerTransport.cs                            | 2 +-
 .../Client/McpClientExtensionsTests.cs                      | 6 +++---
 .../Client/McpClientFactoryTests.cs                         | 2 +-
 tests/ModelContextProtocol.Tests/ClientIntegrationTests.cs  | 2 +-
 .../DockerEverythingServerTests.cs                          | 2 +-
 tests/ModelContextProtocol.Tests/Server/McpServerTests.cs   | 2 +-
 .../ModelContextProtocol.Tests/Utils/TestServerTransport.cs | 2 +-
 11 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/src/ModelContextProtocol/Client/McpClientExtensions.cs b/src/ModelContextProtocol/Client/McpClientExtensions.cs
index 9cc8ab2..aaaa426 100644
--- a/src/ModelContextProtocol/Client/McpClientExtensions.cs
+++ b/src/ModelContextProtocol/Client/McpClientExtensions.cs
@@ -926,7 +926,7 @@ public static class McpClientExtensions
         {
             Content = content,
             Model = chatResponse.ModelId ?? "unknown",
-            Role = lastMessage?.Role == ChatRole.User ? "user" : "assistant",
+            Role = lastMessage?.Role == ChatRole.User ? Role.User : Role.Assistant,
             StopReason = chatResponse.FinishReason == ChatFinishReason.Length ? "maxTokens" : "endTurn",
         };
     }
diff --git a/src/ModelContextProtocol/Protocol/Types/CreateMessageResult.cs b/src/ModelContextProtocol/Protocol/Types/CreateMessageResult.cs
index 04397be..4ae6890 100644
--- a/src/ModelContextProtocol/Protocol/Types/CreateMessageResult.cs
+++ b/src/ModelContextProtocol/Protocol/Types/CreateMessageResult.cs
@@ -50,5 +50,5 @@ public class CreateMessageResult
     /// Gets or sets the role of the user who generated the message.
     /// </summary>
     [JsonPropertyName("role")]
-    public required string Role { get; init; }
+    public required Role Role { get; init; }
 }
diff --git a/src/ModelContextProtocol/Server/McpServerExtensions.cs b/src/ModelContextProtocol/Server/McpServerExtensions.cs
index 17b4b35..aa103c3 100644
--- a/src/ModelContextProtocol/Server/McpServerExtensions.cs
+++ b/src/ModelContextProtocol/Server/McpServerExtensions.cs
@@ -142,7 +142,7 @@ public static class McpServerExtensions
                 ModelPreferences = modelPreferences,
             }, cancellationToken).ConfigureAwait(false);
 
-        return new(new ChatMessage(new(result.Role), [result.Content.ToAIContent()]))
+        return new(new ChatMessage(result.Role is Role.User ? ChatRole.User : ChatRole.Assistant, [result.Content.ToAIContent()]))
         {
             ModelId = result.Model,
             FinishReason = result.StopReason switch
diff --git a/tests/ModelContextProtocol.AspNetCore.Tests/SseServerIntegrationTests.cs b/tests/ModelContextProtocol.AspNetCore.Tests/SseServerIntegrationTests.cs
index 7913bbb..280a757 100644
--- a/tests/ModelContextProtocol.AspNetCore.Tests/SseServerIntegrationTests.cs
+++ b/tests/ModelContextProtocol.AspNetCore.Tests/SseServerIntegrationTests.cs
@@ -219,7 +219,7 @@ public class SseServerIntegrationTests : LoggedTest, IClassFixture<SseServerInte
             return new CreateMessageResult
             {
                 Model = "test-model",
-                Role = "assistant",
+                Role = Role.Assistant,
                 Content = new Content
                 {
                     Type = "text",
diff --git a/tests/ModelContextProtocol.AspNetCore.Tests/Utils/TestServerTransport.cs b/tests/ModelContextProtocol.AspNetCore.Tests/Utils/TestServerTransport.cs
index f216601..0c9cb43 100644
--- a/tests/ModelContextProtocol.AspNetCore.Tests/Utils/TestServerTransport.cs
+++ b/tests/ModelContextProtocol.AspNetCore.Tests/Utils/TestServerTransport.cs
@@ -72,7 +72,7 @@ public class TestServerTransport : ITransport
         await WriteMessageAsync(new JsonRpcResponse
         {
             Id = request.Id,
-            Result = JsonSerializer.SerializeToNode(new CreateMessageResult { Content = new(), Model = "model", Role = "role" }),
+            Result = JsonSerializer.SerializeToNode(new CreateMessageResult { Content = new(), Model = "model", Role = Role.Assistant }),
         }, cancellationToken);
     }
 
diff --git a/tests/ModelContextProtocol.Tests/Client/McpClientExtensionsTests.cs b/tests/ModelContextProtocol.Tests/Client/McpClientExtensionsTests.cs
index 58864a1..c9cb374 100644
--- a/tests/ModelContextProtocol.Tests/Client/McpClientExtensionsTests.cs
+++ b/tests/ModelContextProtocol.Tests/Client/McpClientExtensionsTests.cs
@@ -84,7 +84,7 @@ public class McpClientExtensionsTests : ClientServerTestBase
         Assert.NotNull(result);
         Assert.Equal("Hello, World!", result.Content.Text);
         Assert.Equal("test-model", result.Model);
-        Assert.Equal("assistant", result.Role);
+        Assert.Equal(Role.Assistant, result.Role);
         Assert.Equal("endTurn", result.StopReason);
     }
 
@@ -139,7 +139,7 @@ public class McpClientExtensionsTests : ClientServerTestBase
         Assert.NotNull(result);
         Assert.Equal(expectedData, result.Content.Data);
         Assert.Equal("test-model", result.Model);
-        Assert.Equal("assistant", result.Role);
+        Assert.Equal(Role.Assistant, result.Role);
         Assert.Equal("endTurn", result.StopReason);
     }
 
@@ -201,7 +201,7 @@ public class McpClientExtensionsTests : ClientServerTestBase
         // Assert
         Assert.NotNull(result);
         Assert.Equal("test-model", result.Model);
-        Assert.Equal(ChatRole.Assistant.ToString(), result.Role);
+        Assert.Equal(Role.Assistant, result.Role);
         Assert.Equal("endTurn", result.StopReason);
     }
 
diff --git a/tests/ModelContextProtocol.Tests/Client/McpClientFactoryTests.cs b/tests/ModelContextProtocol.Tests/Client/McpClientFactoryTests.cs
index 1414f65..38e01eb 100644
--- a/tests/ModelContextProtocol.Tests/Client/McpClientFactoryTests.cs
+++ b/tests/ModelContextProtocol.Tests/Client/McpClientFactoryTests.cs
@@ -44,7 +44,7 @@ public class McpClientFactoryTests
                         new CreateMessageResult { 
                             Content = new Content { Text = "result" }, 
                             Model = "test-model", 
-                            Role = "test-role", 
+                            Role = Role.User, 
                             StopReason = "endTurn" 
                     }),
                 },
diff --git a/tests/ModelContextProtocol.Tests/ClientIntegrationTests.cs b/tests/ModelContextProtocol.Tests/ClientIntegrationTests.cs
index baca8f1..46e928e 100644
--- a/tests/ModelContextProtocol.Tests/ClientIntegrationTests.cs
+++ b/tests/ModelContextProtocol.Tests/ClientIntegrationTests.cs
@@ -374,7 +374,7 @@ public class ClientIntegrationTests : LoggedTest, IClassFixture<ClientIntegratio
                         return Task.FromResult(new CreateMessageResult
                         {
                             Model = "test-model",
-                            Role = "assistant",
+                            Role = Role.Assistant,
                             Content = new Content
                             {
                                 Type = "text",
diff --git a/tests/ModelContextProtocol.Tests/DockerEverythingServerTests.cs b/tests/ModelContextProtocol.Tests/DockerEverythingServerTests.cs
index b5947ea..149c7d6 100644
--- a/tests/ModelContextProtocol.Tests/DockerEverythingServerTests.cs
+++ b/tests/ModelContextProtocol.Tests/DockerEverythingServerTests.cs
@@ -83,7 +83,7 @@ public class DockerEverythingServerTests(ITestOutputHelper testOutputHelper) : L
                         return Task.FromResult(new CreateMessageResult
                         {
                             Model = "test-model",
-                            Role = "assistant",
+                            Role = Role.Assistant,
                             Content = new Content
                             {
                                 Type = "text",
diff --git a/tests/ModelContextProtocol.Tests/Server/McpServerTests.cs b/tests/ModelContextProtocol.Tests/Server/McpServerTests.cs
index 054a479..dcb688c 100644
--- a/tests/ModelContextProtocol.Tests/Server/McpServerTests.cs
+++ b/tests/ModelContextProtocol.Tests/Server/McpServerTests.cs
@@ -601,7 +601,7 @@ public class McpServerTests : LoggedTest
             {
                 Content = new() { Text = "The Eiffel Tower.", Type = "text" },
                 Model = "amazingmodel",
-                Role = "assistant",
+                Role = Role.Assistant,
                 StopReason = "endTurn",
             };
 
diff --git a/tests/ModelContextProtocol.Tests/Utils/TestServerTransport.cs b/tests/ModelContextProtocol.Tests/Utils/TestServerTransport.cs
index ed9b2e0..aae1438 100644
--- a/tests/ModelContextProtocol.Tests/Utils/TestServerTransport.cs
+++ b/tests/ModelContextProtocol.Tests/Utils/TestServerTransport.cs
@@ -72,7 +72,7 @@ public class TestServerTransport : ITransport
         await WriteMessageAsync(new JsonRpcResponse
         {
             Id = request.Id,
-            Result = JsonSerializer.SerializeToNode(new CreateMessageResult { Content = new(), Model = "model", Role = "role" }),
+            Result = JsonSerializer.SerializeToNode(new CreateMessageResult { Content = new(), Model = "model", Role = Role.User }),
         }, cancellationToken);
     }