@@ -128,14 +128,18 @@ public async Task SampleAsync_Should_Throw_Exception_If_Client_Does_Not_Support_
128128 // Arrange
129129 await using var transport = new TestServerTransport ( ) ;
130130 await using var server = McpServer . Create ( transport , _options , LoggerFactory ) ;
131- SetClientCapabilities ( server , new ClientCapabilities ( ) ) ;
131+ var runTask = server . RunAsync ( TestContext . Current . CancellationToken ) ;
132+ await InitializeServerAsync ( transport , new ClientCapabilities ( ) , TestContext . Current . CancellationToken ) ;
132133
133134 var action = async ( ) => await server . SampleAsync (
134135 new CreateMessageRequestParams { Messages = [ ] , MaxTokens = 1000 } ,
135136 CancellationToken . None ) ;
136137
137138 // Act & Assert
138139 await Assert . ThrowsAsync < InvalidOperationException > ( action ) ;
140+
141+ await transport . DisposeAsync ( ) ;
142+ await runTask ;
139143 }
140144
141145 [ Fact ]
@@ -144,9 +148,8 @@ public async Task SampleAsync_Should_SendRequest()
144148 // Arrange
145149 await using var transport = new TestServerTransport ( ) ;
146150 await using var server = McpServer . Create ( transport , _options , LoggerFactory ) ;
147- SetClientCapabilities ( server , new ClientCapabilities { Sampling = new SamplingCapability ( ) } ) ;
148-
149151 var runTask = server . RunAsync ( TestContext . Current . CancellationToken ) ;
152+ await InitializeServerAsync ( transport , new ClientCapabilities { Sampling = new SamplingCapability ( ) } , TestContext . Current . CancellationToken ) ;
150153
151154 // Act
152155 var result = await server . SampleAsync (
@@ -155,8 +158,10 @@ public async Task SampleAsync_Should_SendRequest()
155158
156159 Assert . NotNull ( result ) ;
157160 Assert . NotEmpty ( transport . SentMessages ) ;
158- Assert . IsType < JsonRpcRequest > ( transport . SentMessages [ 0 ] ) ;
159- Assert . Equal ( RequestMethods . SamplingCreateMessage , ( ( JsonRpcRequest ) transport . SentMessages [ 0 ] ) . Method ) ;
161+ // First message is the initialize response, second is the sampling request
162+ Assert . True ( transport . SentMessages . Count >= 2 , "Expected at least 2 messages (initialize response and sampling request)" ) ;
163+ var samplingRequest = Assert . IsType < JsonRpcRequest > ( transport . SentMessages [ 1 ] ) ;
164+ Assert . Equal ( RequestMethods . SamplingCreateMessage , samplingRequest . Method ) ;
160165
161166 await transport . DisposeAsync ( ) ;
162167 await runTask ;
@@ -168,12 +173,16 @@ public async Task RequestRootsAsync_Should_Throw_Exception_If_Client_Does_Not_Su
168173 // Arrange
169174 await using var transport = new TestServerTransport ( ) ;
170175 await using var server = McpServer . Create ( transport , _options , LoggerFactory ) ;
171- SetClientCapabilities ( server , new ClientCapabilities ( ) ) ;
176+ var runTask = server . RunAsync ( TestContext . Current . CancellationToken ) ;
177+ await InitializeServerAsync ( transport , new ClientCapabilities ( ) , TestContext . Current . CancellationToken ) ;
172178
173179 // Act & Assert
174180 await Assert . ThrowsAsync < InvalidOperationException > ( async ( ) => await server . RequestRootsAsync (
175181 new ListRootsRequestParams ( ) ,
176182 CancellationToken . None ) ) ;
183+
184+ await transport . DisposeAsync ( ) ;
185+ await runTask ;
177186 }
178187
179188 [ Fact ]
@@ -182,17 +191,19 @@ public async Task RequestRootsAsync_Should_SendRequest()
182191 // Arrange
183192 await using var transport = new TestServerTransport ( ) ;
184193 await using var server = McpServer . Create ( transport , _options , LoggerFactory ) ;
185- SetClientCapabilities ( server , new ClientCapabilities { Roots = new RootsCapability ( ) } ) ;
186194 var runTask = server . RunAsync ( TestContext . Current . CancellationToken ) ;
195+ await InitializeServerAsync ( transport , new ClientCapabilities { Roots = new RootsCapability ( ) } , TestContext . Current . CancellationToken ) ;
187196
188197 // Act
189198 var result = await server . RequestRootsAsync ( new ListRootsRequestParams ( ) , CancellationToken . None ) ;
190199
191200 // Assert
192201 Assert . NotNull ( result ) ;
193202 Assert . NotEmpty ( transport . SentMessages ) ;
194- Assert . IsType < JsonRpcRequest > ( transport . SentMessages [ 0 ] ) ;
195- Assert . Equal ( RequestMethods . RootsList , ( ( JsonRpcRequest ) transport . SentMessages [ 0 ] ) . Method ) ;
203+ // First message is the initialize response, second is the roots request
204+ Assert . True ( transport . SentMessages . Count >= 2 , "Expected at least 2 messages (initialize response and roots request)" ) ;
205+ var rootsRequest = Assert . IsType < JsonRpcRequest > ( transport . SentMessages [ 1 ] ) ;
206+ Assert . Equal ( RequestMethods . RootsList , rootsRequest . Method ) ;
196207
197208 await transport . DisposeAsync ( ) ;
198209 await runTask ;
@@ -204,12 +215,16 @@ public async Task ElicitAsync_Should_Throw_Exception_If_Client_Does_Not_Support_
204215 // Arrange
205216 await using var transport = new TestServerTransport ( ) ;
206217 await using var server = McpServer . Create ( transport , _options , LoggerFactory ) ;
207- SetClientCapabilities ( server , new ClientCapabilities ( ) ) ;
218+ var runTask = server . RunAsync ( TestContext . Current . CancellationToken ) ;
219+ await InitializeServerAsync ( transport , new ClientCapabilities ( ) , TestContext . Current . CancellationToken ) ;
208220
209221 // Act & Assert
210222 await Assert . ThrowsAsync < InvalidOperationException > ( async ( ) => await server . ElicitAsync (
211223 new ElicitRequestParams { Message = "" } ,
212224 CancellationToken . None ) ) ;
225+
226+ await transport . DisposeAsync ( ) ;
227+ await runTask ;
213228 }
214229
215230 [ Fact ]
@@ -218,23 +233,25 @@ public async Task ElicitAsync_Should_SendRequest()
218233 // Arrange
219234 await using var transport = new TestServerTransport ( ) ;
220235 await using var server = McpServer . Create ( transport , _options , LoggerFactory ) ;
221- SetClientCapabilities ( server , new ClientCapabilities
236+ var runTask = server . RunAsync ( TestContext . Current . CancellationToken ) ;
237+ await InitializeServerAsync ( transport , new ClientCapabilities
222238 {
223239 Elicitation = new ( )
224240 {
225241 Form = new ( ) ,
226242 } ,
227- } ) ;
228- var runTask = server . RunAsync ( TestContext . Current . CancellationToken ) ;
243+ } , TestContext . Current . CancellationToken ) ;
229244
230245 // Act
231246 var result = await server . ElicitAsync ( new ElicitRequestParams { Message = "" , RequestedSchema = new ( ) } , CancellationToken . None ) ;
232247
233248 // Assert
234249 Assert . NotNull ( result ) ;
235250 Assert . NotEmpty ( transport . SentMessages ) ;
236- Assert . IsType < JsonRpcRequest > ( transport . SentMessages [ 0 ] ) ;
237- Assert . Equal ( RequestMethods . ElicitationCreate , ( ( JsonRpcRequest ) transport . SentMessages [ 0 ] ) . Method ) ;
251+ // First message is the initialize response, second is the elicit request
252+ Assert . True ( transport . SentMessages . Count >= 2 , "Expected at least 2 messages (initialize response and elicit request)" ) ;
253+ var elicitRequest = Assert . IsType < JsonRpcRequest > ( transport . SentMessages [ 1 ] ) ;
254+ Assert . Equal ( RequestMethods . ElicitationCreate , elicitRequest . Method ) ;
238255
239256 await transport . DisposeAsync ( ) ;
240257 await runTask ;
@@ -844,11 +861,33 @@ public async Task Can_SendMessage_Before_RunAsync()
844861 Assert . Same ( logNotification , transport . SentMessages [ 0 ] ) ;
845862 }
846863
847- private static void SetClientCapabilities ( McpServer server , ClientCapabilities capabilities )
864+ private static async Task InitializeServerAsync ( TestServerTransport transport , ClientCapabilities capabilities , CancellationToken cancellationToken = default )
848865 {
849- FieldInfo ? field = server . GetType ( ) . GetField ( "_clientCapabilities" , BindingFlags . NonPublic | BindingFlags . Instance ) ;
850- Assert . NotNull ( field ) ;
851- field . SetValue ( server , capabilities ) ;
866+ var initializeRequest = new JsonRpcRequest
867+ {
868+ Id = new RequestId ( "init-1" ) ,
869+ Method = RequestMethods . Initialize ,
870+ Params = JsonSerializer . SerializeToNode ( new InitializeRequestParams
871+ {
872+ ProtocolVersion = "2024-11-05" ,
873+ Capabilities = capabilities ,
874+ ClientInfo = new Implementation { Name = "test-client" , Version = "1.0.0" }
875+ } , McpJsonUtilities . DefaultOptions )
876+ } ;
877+
878+ var tcs = new TaskCompletionSource < bool > ( ) ;
879+ transport . OnMessageSent = ( message ) =>
880+ {
881+ if ( message is JsonRpcResponse response && response . Id == initializeRequest . Id )
882+ {
883+ tcs . TrySetResult ( true ) ;
884+ }
885+ } ;
886+
887+ await transport . SendClientMessageAsync ( initializeRequest , cancellationToken ) ;
888+
889+ // Wait for the initialize response to be sent
890+ await tcs . Task . WaitAsync ( TimeSpan . FromSeconds ( 5 ) , cancellationToken ) ;
852891 }
853892
854893 private sealed class TestServerForIChatClient ( bool supportsSampling ) : McpServer
0 commit comments