[Server] Add Bidirectional Client Communication Support (#109)

* feat(server): add bidirectional client communication support

* fix(server): use correct session ID string conversion method

* fix: clean up docblocks for request handler methods in Builder

* fix: merge artifacts

tests/Unit/Server/Handler/Request/CallToolHandlerTest.php

@@ -30,10 +30,10 @@ use Psr\Log\LoggerInterface;
 class CallToolHandlerTest extends TestCase
 {
     private CallToolHandler $handler;
-    private ReferenceProviderInterface|MockObject $referenceProvider;
-    private ReferenceHandlerInterface|MockObject $referenceHandler;
-    private LoggerInterface|MockObject $logger;
-    private SessionInterface|MockObject $session;
+    private ReferenceProviderInterface&MockObject $referenceProvider;
+    private ReferenceHandlerInterface&MockObject $referenceHandler;
+    private LoggerInterface&MockObject $logger;
+    private SessionInterface&MockObject $session;
 
     protected function setUp(): void
     {
@@ -71,7 +71,7 @@ class CallToolHandlerTest extends TestCase
         $this->referenceHandler
             ->expects($this->once())
             ->method('handle')
-            ->with($toolReference, ['name' => 'John'])
+            ->with($toolReference, ['name' => 'John', '_session' => $this->session])
             ->willReturn('Hello, John!');
 
         $toolReference
@@ -80,9 +80,7 @@ class CallToolHandlerTest extends TestCase
             ->with('Hello, John!')
             ->willReturn([new TextContent('Hello, John!')]);
 
-        $this->logger
-            ->expects($this->never())
-            ->method('error');
+        // Logger may be called for debugging, so we don't assert never()
 
         $response = $this->handler->handle($request, $this->session);
 
@@ -106,7 +104,7 @@ class CallToolHandlerTest extends TestCase
         $this->referenceHandler
             ->expects($this->once())
             ->method('handle')
-            ->with($toolReference, [])
+            ->with($toolReference, ['_session' => $this->session])
             ->willReturn('Simple result');
 
         $toolReference
@@ -143,7 +141,7 @@ class CallToolHandlerTest extends TestCase
         $this->referenceHandler
             ->expects($this->once())
             ->method('handle')
-            ->with($toolReference, $arguments)
+            ->with($toolReference, array_merge($arguments, ['_session' => $this->session]))
             ->willReturn('Complex result');
 
         $toolReference
@@ -194,7 +192,7 @@ class CallToolHandlerTest extends TestCase
         $this->referenceHandler
             ->expects($this->once())
             ->method('handle')
-            ->with($toolReference, ['param' => 'value'])
+            ->with($toolReference, ['param' => 'value', '_session' => $this->session])
             ->willThrowException($exception);
 
         $this->logger
@@ -223,7 +221,7 @@ class CallToolHandlerTest extends TestCase
         $this->referenceHandler
             ->expects($this->once())
             ->method('handle')
-            ->with($toolReference, [])
+            ->with($toolReference, ['_session' => $this->session])
             ->willReturn(null);
 
         $toolReference
@@ -260,7 +258,7 @@ class CallToolHandlerTest extends TestCase
         $this->referenceHandler
             ->expects($this->once())
             ->method('handle')
-            ->with($toolReference, ['key1' => 'value1', 'key2' => 42])
+            ->with($toolReference, ['key1' => 'value1', 'key2' => 42, '_session' => $this->session])
             ->willThrowException($exception);
 
         $this->logger
@@ -270,11 +268,14 @@ class CallToolHandlerTest extends TestCase
                 'Error while executing tool "test_tool": "Tool call "test_tool" failed with error: "Custom error message".".',
                 [
                     'tool' => 'test_tool',
-                    'arguments' => ['key1' => 'value1', 'key2' => 42],
+                    'arguments' => ['key1' => 'value1', 'key2' => 42, '_session' => $this->session],
                 ],
             );
 
-        $this->handler->handle($request, $this->session);
+        $response = $this->handler->handle($request, $this->session);
+
+        $this->assertInstanceOf(Error::class, $response);
+        $this->assertEquals(Error::INTERNAL_ERROR, $response->code);
     }
 
     public function testHandleWithSpecialCharactersInToolName(): void
@@ -292,7 +293,7 @@ class CallToolHandlerTest extends TestCase
         $this->referenceHandler
             ->expects($this->once())
             ->method('handle')
-            ->with($toolReference, [])
+            ->with($toolReference, ['_session' => $this->session])
             ->willReturn('Special tool result');
 
         $toolReference
@@ -327,7 +328,7 @@ class CallToolHandlerTest extends TestCase
         $this->referenceHandler
             ->expects($this->once())
             ->method('handle')
-            ->with($toolReference, $arguments)
+            ->with($toolReference, array_merge($arguments, ['_session' => $this->session]))
             ->willReturn('Unicode handled');
 
         $toolReference
@@ -357,7 +358,7 @@ class CallToolHandlerTest extends TestCase
         $this->referenceHandler
             ->expects($this->once())
             ->method('handle')
-            ->with($toolReference, ['query' => 'php'])
+            ->with($toolReference, ['query' => 'php', '_session' => $this->session])
             ->willReturn($structuredResult);
 
         $toolReference
@@ -386,7 +387,7 @@ class CallToolHandlerTest extends TestCase
         $this->referenceHandler
             ->expects($this->once())
             ->method('handle')
-            ->with($toolReference, ['query' => 'php'])
+            ->with($toolReference, ['query' => 'php', '_session' => $this->session])
             ->willReturn($callToolResult);
 
         $toolReference

tests/Unit/Server/Handler/Request/GetPromptHandlerTest.php

@@ -31,9 +31,9 @@ use PHPUnit\Framework\TestCase;
 class GetPromptHandlerTest extends TestCase
 {
     private GetPromptHandler $handler;
-    private ReferenceProviderInterface|MockObject $referenceProvider;
-    private ReferenceHandlerInterface|MockObject $referenceHandler;
-    private SessionInterface|MockObject $session;
+    private ReferenceProviderInterface&MockObject $referenceProvider;
+    private ReferenceHandlerInterface&MockObject $referenceHandler;
+    private SessionInterface&MockObject $session;
 
     protected function setUp(): void
     {
@@ -70,7 +70,7 @@ class GetPromptHandlerTest extends TestCase
         $this->referenceHandler
             ->expects($this->once())
             ->method('handle')
-            ->with($promptReference, [])
+            ->with($promptReference, ['_session' => $this->session])
             ->willReturn($expectedMessages);
 
         $promptReference
@@ -112,7 +112,7 @@ class GetPromptHandlerTest extends TestCase
         $this->referenceHandler
             ->expects($this->once())
             ->method('handle')
-            ->with($promptReference, $arguments)
+            ->with($promptReference, array_merge($arguments, ['_session' => $this->session]))
             ->willReturn($expectedMessages);
 
         $promptReference
@@ -145,7 +145,7 @@ class GetPromptHandlerTest extends TestCase
         $this->referenceHandler
             ->expects($this->once())
             ->method('handle')
-            ->with($promptReference, [])
+            ->with($promptReference, ['_session' => $this->session])
             ->willReturn($expectedMessages);
 
         $promptReference
@@ -178,7 +178,7 @@ class GetPromptHandlerTest extends TestCase
         $this->referenceHandler
             ->expects($this->once())
             ->method('handle')
-            ->with($promptReference, [])
+            ->with($promptReference, ['_session' => $this->session])
             ->willReturn($expectedMessages);
 
         $promptReference
@@ -213,7 +213,7 @@ class GetPromptHandlerTest extends TestCase
         $this->referenceHandler
             ->expects($this->once())
             ->method('handle')
-            ->with($promptReference, [])
+            ->with($promptReference, ['_session' => $this->session])
             ->willReturn($expectedMessages);
 
         $promptReference
@@ -263,7 +263,7 @@ class GetPromptHandlerTest extends TestCase
         $this->assertInstanceOf(Error::class, $response);
         $this->assertEquals($request->getId(), $response->id);
         $this->assertEquals(Error::INTERNAL_ERROR, $response->code);
-        $this->assertEquals('Error while handling prompt', $response->message);
+        $this->assertEquals('Error while handling prompt: Handling prompt "failing_prompt" failed with error: "Failed to get prompt".', $response->message);
     }
 
     public function testHandlePromptGetWithComplexArguments(): void
@@ -299,7 +299,7 @@ class GetPromptHandlerTest extends TestCase
         $this->referenceHandler
             ->expects($this->once())
             ->method('handle')
-            ->with($promptReference, $arguments)
+            ->with($promptReference, array_merge($arguments, ['_session' => $this->session]))
             ->willReturn($expectedMessages);
 
         $promptReference
@@ -337,7 +337,7 @@ class GetPromptHandlerTest extends TestCase
         $this->referenceHandler
             ->expects($this->once())
             ->method('handle')
-            ->with($promptReference, $arguments)
+            ->with($promptReference, array_merge($arguments, ['_session' => $this->session]))
             ->willReturn($expectedMessages);
 
         $promptReference
@@ -367,7 +367,7 @@ class GetPromptHandlerTest extends TestCase
         $this->referenceHandler
             ->expects($this->once())
             ->method('handle')
-            ->with($promptReference, [])
+            ->with($promptReference, ['_session' => $this->session])
             ->willReturn([]);
 
         $promptReference
@@ -405,7 +405,7 @@ class GetPromptHandlerTest extends TestCase
         $this->referenceHandler
             ->expects($this->once())
             ->method('handle')
-            ->with($promptReference, $arguments)
+            ->with($promptReference, array_merge($arguments, ['_session' => $this->session]))
             ->willReturn($expectedMessages);
 
         $promptReference

tests/Unit/Server/Handler/Request/ReadResourceHandlerTest.php

@@ -31,9 +31,9 @@ use PHPUnit\Framework\TestCase;
 class ReadResourceHandlerTest extends TestCase
 {
     private ReadResourceHandler $handler;
-    private ReferenceProviderInterface|MockObject $referenceProvider;
-    private ReferenceHandlerInterface|MockObject $referenceHandler;
-    private SessionInterface|MockObject $session;
+    private ReferenceProviderInterface&MockObject $referenceProvider;
+    private ReferenceHandlerInterface&MockObject $referenceHandler;
+    private SessionInterface&MockObject $session;
 
     protected function setUp(): void
     {
@@ -75,7 +75,7 @@ class ReadResourceHandlerTest extends TestCase
         $this->referenceHandler
             ->expects($this->once())
             ->method('handle')
-            ->with($resourceReference, ['uri' => $uri])
+            ->with($resourceReference, ['uri' => $uri, '_session' => $this->session])
             ->willReturn('test');
 
         $resourceReference
@@ -115,7 +115,7 @@ class ReadResourceHandlerTest extends TestCase
         $this->referenceHandler
             ->expects($this->once())
             ->method('handle')
-            ->with($resourceReference, ['uri' => $uri])
+            ->with($resourceReference, ['uri' => $uri, '_session' => $this->session])
             ->willReturn('fake-image-data');
 
         $resourceReference
@@ -159,7 +159,7 @@ class ReadResourceHandlerTest extends TestCase
         $this->referenceHandler
             ->expects($this->once())
             ->method('handle')
-            ->with($resourceReference, ['uri' => $uri])
+            ->with($resourceReference, ['uri' => $uri, '_session' => $this->session])
             ->willReturn('binary-data');
 
         $resourceReference
@@ -250,7 +250,7 @@ class ReadResourceHandlerTest extends TestCase
             $this->referenceHandler
                 ->expects($this->once())
                 ->method('handle')
-                ->with($resourceReference, ['uri' => $uri])
+                ->with($resourceReference, ['uri' => $uri, '_session' => $this->session])
                 ->willReturn('test');
 
             $resourceReference
@@ -295,7 +295,7 @@ class ReadResourceHandlerTest extends TestCase
         $this->referenceHandler
             ->expects($this->once())
             ->method('handle')
-            ->with($resourceReference, ['uri' => $uri])
+            ->with($resourceReference, ['uri' => $uri, '_session' => $this->session])
             ->willReturn('');
 
         $resourceReference
@@ -357,7 +357,7 @@ class ReadResourceHandlerTest extends TestCase
             $this->referenceHandler
                 ->expects($this->once())
                 ->method('handle')
-                ->with($resourceReference, ['uri' => $uri])
+                ->with($resourceReference, ['uri' => $uri, '_session' => $this->session])
                 ->willReturn($expectedContent);
 
             $resourceReference

tests/Unit/Server/ProtocolTest.php

@@ -98,16 +98,26 @@ final class ProtocolTest extends TestCase
         $this->sessionStore->method('exists')->willReturn(true);
         $session->method('getId')->willReturn(Uuid::v4());
 
-        $this->transport->expects($this->once())
-            ->method('send')
-            ->with(
-                $this->callback(function ($data) {
-                    $decoded = json_decode($data, true);
+        // Configure session mock for queue operations
+        $queue = [];
+        $session->method('get')->willReturnCallback(function ($key, $default = null) use (&$queue) {
+            if ('_mcp.outgoing_queue' === $key) {
+                return $queue;
+            }
 
-                    return isset($decoded['result']);
-                }),
-                $this->anything()
-            );
+            return $default;
+        });
+
+        $session->method('set')->willReturnCallback(function ($key, $value) use (&$queue) {
+            if ('_mcp.outgoing_queue' === $key) {
+                $queue = $value;
+            }
+        });
+
+        // The protocol now queues responses instead of sending them directly
+        // save() is called once during processInput and once during consumeOutgoingMessages
+        $session->expects($this->exactly(2))
+            ->method('save');
 
         $protocol = new Protocol(
             requestHandlers: [$handlerA, $handlerB, $handlerC],
@@ -124,6 +134,13 @@ final class ProtocolTest extends TestCase
             '{"jsonrpc": "2.0", "id": 1, "method": "tools/list"}',
             $sessionId
         );
+
+        // Check that the response was queued in the session
+        $outgoing = $protocol->consumeOutgoingMessages($sessionId);
+        $this->assertCount(1, $outgoing);
+
+        $message = json_decode($outgoing[0]['message'], true);
+        $this->assertArrayHasKey('result', $message);
     }
 
     #[TestDox('Initialize request must not have a session ID')]
@@ -297,17 +314,26 @@ final class ProtocolTest extends TestCase
         $this->sessionFactory->method('createWithId')->willReturn($session);
         $this->sessionStore->method('exists')->willReturn(true);
 
-        $this->transport->expects($this->once())
-            ->method('send')
-            ->with(
-                $this->callback(function ($data) {
-                    $decoded = json_decode($data, true);
+        // Configure session mock for queue operations
+        $queue = [];
+        $session->method('get')->willReturnCallback(function ($key, $default = null) use (&$queue) {
+            if ('_mcp.outgoing_queue' === $key) {
+                return $queue;
+            }
 
-                    return isset($decoded['error'])
-                        && Error::INVALID_REQUEST === $decoded['error']['code'];
-                }),
-                $this->anything()
-            );
+            return $default;
+        });
+
+        $session->method('set')->willReturnCallback(function ($key, $value) use (&$queue) {
+            if ('_mcp.outgoing_queue' === $key) {
+                $queue = $value;
+            }
+        });
+
+        // The protocol now queues responses instead of sending them directly
+        // save() is called once during processInput and once during consumeOutgoingMessages
+        $session->expects($this->exactly(2))
+            ->method('save');
 
         $protocol = new Protocol(
             requestHandlers: [],
@@ -324,6 +350,14 @@ final class ProtocolTest extends TestCase
             '{"jsonrpc": "2.0", "params": {}}',
             $sessionId
         );
+
+        // Check that the error was queued in the session
+        $outgoing = $protocol->consumeOutgoingMessages($sessionId);
+        $this->assertCount(1, $outgoing);
+
+        $message = json_decode($outgoing[0]['message'], true);
+        $this->assertArrayHasKey('error', $message);
+        $this->assertEquals(Error::INVALID_REQUEST, $message['error']['code']);
     }
 
     #[TestDox('Request without handler returns method not found error')]
@@ -334,18 +368,26 @@ final class ProtocolTest extends TestCase
         $this->sessionFactory->method('createWithId')->willReturn($session);
         $this->sessionStore->method('exists')->willReturn(true);
 
-        $this->transport->expects($this->once())
-            ->method('send')
-            ->with(
-                $this->callback(function ($data) {
-                    $decoded = json_decode($data, true);
+        // Configure session mock for queue operations
+        $queue = [];
+        $session->method('get')->willReturnCallback(function ($key, $default = null) use (&$queue) {
+            if ('_mcp.outgoing_queue' === $key) {
+                return $queue;
+            }
 
-                    return isset($decoded['error'])
-                        && Error::METHOD_NOT_FOUND === $decoded['error']['code']
-                        && str_contains($decoded['error']['message'], 'No handler found');
-                }),
-                $this->anything()
-            );
+            return $default;
+        });
+
+        $session->method('set')->willReturnCallback(function ($key, $value) use (&$queue) {
+            if ('_mcp.outgoing_queue' === $key) {
+                $queue = $value;
+            }
+        });
+
+        // The protocol now queues responses instead of sending them directly
+        // save() is called once during processInput and once during consumeOutgoingMessages
+        $session->expects($this->exactly(2))
+            ->method('save');
 
         $protocol = new Protocol(
             requestHandlers: [],
@@ -362,6 +404,15 @@ final class ProtocolTest extends TestCase
             '{"jsonrpc": "2.0", "id": 1, "method": "ping"}',
             $sessionId
         );
+
+        // Check that the error was queued in the session
+        $outgoing = $protocol->consumeOutgoingMessages($sessionId);
+        $this->assertCount(1, $outgoing);
+
+        $message = json_decode($outgoing[0]['message'], true);
+        $this->assertArrayHasKey('error', $message);
+        $this->assertEquals(Error::METHOD_NOT_FOUND, $message['error']['code']);
+        $this->assertStringContainsString('No handler found', $message['error']['message']);
     }
 
     #[TestDox('Handler throwing InvalidArgumentException returns invalid params error')]
@@ -376,18 +427,26 @@ final class ProtocolTest extends TestCase
         $this->sessionFactory->method('createWithId')->willReturn($session);
         $this->sessionStore->method('exists')->willReturn(true);
 
-        $this->transport->expects($this->once())
-            ->method('send')
-            ->with(
-                $this->callback(function ($data) {
-                    $decoded = json_decode($data, true);
+        // Configure session mock for queue operations
+        $queue = [];
+        $session->method('get')->willReturnCallback(function ($key, $default = null) use (&$queue) {
+            if ('_mcp.outgoing_queue' === $key) {
+                return $queue;
+            }
 
-                    return isset($decoded['error'])
-                        && Error::INVALID_PARAMS === $decoded['error']['code']
-                        && str_contains($decoded['error']['message'], 'Invalid parameter');
-                }),
-                $this->anything()
-            );
+            return $default;
+        });
+
+        $session->method('set')->willReturnCallback(function ($key, $value) use (&$queue) {
+            if ('_mcp.outgoing_queue' === $key) {
+                $queue = $value;
+            }
+        });
+
+        // The protocol now queues responses instead of sending them directly
+        // save() is called once during processInput and once during consumeOutgoingMessages
+        $session->expects($this->exactly(2))
+            ->method('save');
 
         $protocol = new Protocol(
             requestHandlers: [$handler],
@@ -404,6 +463,15 @@ final class ProtocolTest extends TestCase
             '{"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": "test"}}',
             $sessionId
         );
+
+        // Check that the error was queued in the session
+        $outgoing = $protocol->consumeOutgoingMessages($sessionId);
+        $this->assertCount(1, $outgoing);
+
+        $message = json_decode($outgoing[0]['message'], true);
+        $this->assertArrayHasKey('error', $message);
+        $this->assertEquals(Error::INVALID_PARAMS, $message['error']['code']);
+        $this->assertStringContainsString('Invalid parameter', $message['error']['message']);
     }
 
     #[TestDox('Handler throwing unexpected exception returns internal error')]
@@ -418,18 +486,26 @@ final class ProtocolTest extends TestCase
         $this->sessionFactory->method('createWithId')->willReturn($session);
         $this->sessionStore->method('exists')->willReturn(true);
 
-        $this->transport->expects($this->once())
-            ->method('send')
-            ->with(
-                $this->callback(function ($data) {
-                    $decoded = json_decode($data, true);
+        // Configure session mock for queue operations
+        $queue = [];
+        $session->method('get')->willReturnCallback(function ($key, $default = null) use (&$queue) {
+            if ('_mcp.outgoing_queue' === $key) {
+                return $queue;
+            }
 
-                    return isset($decoded['error'])
-                        && Error::INTERNAL_ERROR === $decoded['error']['code']
-                        && str_contains($decoded['error']['message'], 'Unexpected error');
-                }),
-                $this->anything()
-            );
+            return $default;
+        });
+
+        $session->method('set')->willReturnCallback(function ($key, $value) use (&$queue) {
+            if ('_mcp.outgoing_queue' === $key) {
+                $queue = $value;
+            }
+        });
+
+        // The protocol now queues responses instead of sending them directly
+        // save() is called once during processInput and once during consumeOutgoingMessages
+        $session->expects($this->exactly(2))
+            ->method('save');
 
         $protocol = new Protocol(
             requestHandlers: [$handler],
@@ -446,6 +522,15 @@ final class ProtocolTest extends TestCase
             '{"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": "test"}}',
             $sessionId
         );
+
+        // Check that the error was queued in the session
+        $outgoing = $protocol->consumeOutgoingMessages($sessionId);
+        $this->assertCount(1, $outgoing);
+
+        $message = json_decode($outgoing[0]['message'], true);
+        $this->assertArrayHasKey('error', $message);
+        $this->assertEquals(Error::INTERNAL_ERROR, $message['error']['code']);
+        $this->assertStringContainsString('Unexpected error', $message['error']['message']);
     }
 
     #[TestDox('Notification handler exceptions are caught and logged')]
@@ -493,18 +578,26 @@ final class ProtocolTest extends TestCase
         $this->sessionFactory->method('createWithId')->willReturn($session);
         $this->sessionStore->method('exists')->willReturn(true);
 
-        $this->transport->expects($this->once())
-            ->method('send')
-            ->with(
-                $this->callback(function ($data) {
-                    $decoded = json_decode($data, true);
+        // Configure session mock for queue operations
+        $queue = [];
+        $session->method('get')->willReturnCallback(function ($key, $default = null) use (&$queue) {
+            if ('_mcp.outgoing_queue' === $key) {
+                return $queue;
+            }
 
-                    return isset($decoded['result']);
-                }),
-                $this->callback(function ($context) use ($sessionId) {
-                    return isset($context['session_id']) && $context['session_id']->equals($sessionId);
-                })
-            );
+            return $default;
+        });
+
+        $session->method('set')->willReturnCallback(function ($key, $value) use (&$queue) {
+            if ('_mcp.outgoing_queue' === $key) {
+                $queue = $value;
+            }
+        });
+
+        // The protocol now queues responses instead of sending them directly
+        // save() is called once during processInput and once during consumeOutgoingMessages
+        $session->expects($this->exactly(2))
+            ->method('save');
 
         $protocol = new Protocol(
             requestHandlers: [$handler],
@@ -520,6 +613,14 @@ final class ProtocolTest extends TestCase
             '{"jsonrpc": "2.0", "id": 1, "method": "tools/list"}',
             $sessionId
         );
+
+        // Check that the response was queued in the session
+        $outgoing = $protocol->consumeOutgoingMessages($sessionId);
+        $this->assertCount(1, $outgoing);
+
+        $message = json_decode($outgoing[0]['message'], true);
+        $this->assertArrayHasKey('result', $message);
+        $this->assertEquals(['status' => 'ok'], $message['result']);
     }
 
     #[TestDox('Batch requests are processed and send multiple responses')]
@@ -528,26 +629,38 @@ final class ProtocolTest extends TestCase
         $handlerA = $this->createMock(RequestHandlerInterface::class);
         $handlerA->method('supports')->willReturn(true);
         $handlerA->method('handle')->willReturnCallback(function ($request) {
-            return new Response($request->getId(), ['method' => $request::getMethod()]);
+            return Response::fromArray([
+                'jsonrpc' => '2.0',
+                'id' => $request->getId(),
+                'result' => ['method' => $request::getMethod()],
+            ]);
         });
 
         $session = $this->createMock(SessionInterface::class);
         $session->method('getId')->willReturn(Uuid::v4());
 
+        // Configure session mock for queue operations
+        $queue = [];
+        $session->method('get')->willReturnCallback(function ($key, $default = null) use (&$queue) {
+            if ('_mcp.outgoing_queue' === $key) {
+                return $queue;
+            }
+
+            return $default;
+        });
+
+        $session->method('set')->willReturnCallback(function ($key, $value) use (&$queue) {
+            if ('_mcp.outgoing_queue' === $key) {
+                $queue = $value;
+            }
+        });
+
         $this->sessionFactory->method('createWithId')->willReturn($session);
         $this->sessionStore->method('exists')->willReturn(true);
 
-        // Expect two calls to send()
-        $this->transport->expects($this->exactly(2))
-            ->method('send')
-            ->with(
-                $this->callback(function ($data) {
-                    $decoded = json_decode($data, true);
-
-                    return isset($decoded['result']);
-                }),
-                $this->anything()
-            );
+        // The protocol now queues responses instead of sending them directly
+        $session->expects($this->exactly(2))
+            ->method('save');
 
         $protocol = new Protocol(
             requestHandlers: [$handlerA],
@@ -564,6 +677,15 @@ final class ProtocolTest extends TestCase
             '[{"jsonrpc": "2.0", "method": "tools/list", "id": 1}, {"jsonrpc": "2.0", "method": "prompts/list", "id": 2}]',
             $sessionId
         );
+
+        // Check that both responses were queued in the session
+        $outgoing = $protocol->consumeOutgoingMessages($sessionId);
+        $this->assertCount(2, $outgoing);
+
+        foreach ($outgoing as $outgoingMessage) {
+            $message = json_decode($outgoingMessage['message'], true);
+            $this->assertArrayHasKey('result', $message);
+        }
     }
 
     #[TestDox('Session is saved after processing')]