From 80dc6c29f59964c2cb31e878caef66db12889475 Mon Sep 17 00:00:00 2001 From: Christopher Hertel Date: Wed, 14 Oct 2026 01:39:07 +0200 Subject: [PATCH 1/2] [Server] Save the session writes of a handler after it resumed --- CHANGELOG.md | 1 + src/Server/Protocol.php | 67 ++++++++++- tests/Unit/Server/ProtocolSessionLockTest.php | 110 ++++++++++++++++++ tests/Unit/Server/ProtocolTest.php | 58 +++++++++ 4 files changed, 233 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8bd4226c..24099b61 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -81,6 +81,7 @@ All notable changes to `mcp/sdk` will be documented in this file. * `Protocol::__construct()` takes an optional `SessionLockInterface $sessionLock` as its last argument. * `FileSessionStore::gc()` also removes the lock files of expired sessions a `FileSessionLock` in the same directory left behind, each only while no request holds it. * Resume a suspended handler with an error when the transport fails to store what it sends to the client next, instead of resuming it with `null`; `ClientGateway::notify()` then throws a `ClientException`. +* Fix session writes of a handler after it resumed from an elicitation, sampling request or notification being lost. 0.8.1 ----- diff --git a/src/Server/Protocol.php b/src/Server/Protocol.php index 834981e9..97cf592a 100644 --- a/src/Server/Protocol.php +++ b/src/Server/Protocol.php @@ -391,20 +391,29 @@ private function handleRequest(TransportInterface $transport, Request $request, $codec = $this->requestStateCodec; $dispatchResult = $this->dispatchResult(...); $errorForThrowable = $this->errorForThrowable(...); + $saveResumedSession = $this->saveResumedSession(...); + /** @var array|null $suspendedState what the session held once the fiber first suspended */ + $suspendedState = null; // One fiber for the whole exchange: with the shim, the handler // re-enters inside it each round rather than needing a new one. // Handled in the fiber, as the transport resumes it outside the catch below. /** @var McpFiber $fiber */ - $fiber = new \Fiber(static function () use ($handler, $request, $session, $shim, $codec, $dispatchResult, $errorForThrowable): Response|Error { + $fiber = new \Fiber(static function () use ($handler, $request, $session, $shim, $codec, $dispatchResult, $errorForThrowable, $saveResumedSession, &$suspendedState): Response|Error { try { $result = $handler->handle($request, $session); $result = $shim?->fulfill($result, $handler, $request, $session, $codec) ?? $result; + $result = $dispatchResult($result, $request, $session); } catch (\Throwable $e) { - return $errorForThrowable($e, $request, $session); + $result = $errorForThrowable($e, $request, $session); } - return $dispatchResult($result, $request, $session); + // A fiber that never suspended is saved by doProcessInput(); after a resume nothing else saves it. + if (null !== $suspendedState) { + $saveResumedSession($session, $suspendedState); + } + + return $result; }); $result = $fiber->start(); @@ -438,6 +447,8 @@ private function handleRequest(TransportInterface $transport, Request $request, throw new RuntimeException('Failed to save the session of a suspended request.'); } + $suspendedState = $session->all(); + $this->trackAwaitedRequest($fiber, $awaitedRequestId); $transport->attachFiberToSession($fiber, $session->getId()); @@ -951,6 +962,56 @@ private function dispatchResult(Response|Error $result, Request $request, Sessio } } + /** + * Stores what a handler wrote to its session after it resumed. + * + * The fiber's session holds the state it suspended with, while the stored one has moved + * on: the client's answer was consumed, the outgoing queue drained, and other requests of + * the session saved theirs. So only the keys the handler changed since it suspended are + * written onto the stored session, under the lock, as the code after a resume runs without it. + * + * @param array $suspendedState + */ + private function saveResumedSession(SessionInterface $session, array $suspendedState): void + { + $state = $session->all(); + $changed = array_filter($state, static fn (mixed $value, int|string $key): bool => !\array_key_exists($key, $suspendedState) || $suspendedState[$key] !== $value, \ARRAY_FILTER_USE_BOTH); + $removed = array_diff_key($suspendedState, $state); + + if ([] === $changed && [] === $removed) { + return; + } + + $sessionId = $session->getId(); + + if (!$this->acquireSessionLock($sessionId)) { + $this->logger->error('Lost the session writes of a resumed request, its session stayed locked.', ['session_id' => $sessionId->toRfc4122()]); + + return; + } + + try { + // A session the client ended while the handler waited must not be brought back. + if (!$this->sessionManager->exists($sessionId)) { + return; + } + + $stored = $this->sessionManager->createWithId($sessionId); + foreach ($changed as $key => $value) { + $stored->set((string) $key, $value); + } + foreach (array_keys($removed) as $key) { + $stored->forget((string) $key); + } + $stored->save(); + } catch (\Throwable $e) { + // The response is already decided: an error for the same request would contradict it. + $this->logger->error(\sprintf('Failed to save the session of a resumed request: %s', $e->getMessage()), ['exception' => $e]); + } finally { + $this->releaseSessionLock($sessionId); + } + } + private function errorForThrowable(\Throwable $e, Request $request, SessionInterface $session): Error { if ($e instanceof \InvalidArgumentException) { diff --git a/tests/Unit/Server/ProtocolSessionLockTest.php b/tests/Unit/Server/ProtocolSessionLockTest.php index d4dee178..dd057b61 100644 --- a/tests/Unit/Server/ProtocolSessionLockTest.php +++ b/tests/Unit/Server/ProtocolSessionLockTest.php @@ -14,14 +14,18 @@ use Mcp\Exception\SessionLockException; use Mcp\JsonRpc\MessageFactory; use Mcp\Schema\JsonRpc\Error; +use Mcp\Schema\JsonRpc\Request; use Mcp\Schema\JsonRpc\Response; use Mcp\Schema\Request\PingRequest; use Mcp\Server\Handler\Request\PingHandler; +use Mcp\Server\Handler\Request\RequestHandlerInterface; use Mcp\Server\Protocol; use Mcp\Server\Session\Session; +use Mcp\Server\Session\SessionInterface; use Mcp\Server\Session\SessionManager; use Mcp\Server\Session\SessionManagerInterface; use Mcp\Server\Suspension\RequestSuspension; +use Mcp\Server\Transport\InMemoryTransport; use Mcp\Tests\Unit\Fixtures\RecordingTransport; use Mcp\Tests\Unit\Server\Session\Fixture\RecordingLogger; use Mcp\Tests\Unit\Server\Session\Fixture\RecordingSessionLock; @@ -453,6 +457,112 @@ public function testWithoutALockOnlyTheStoreIsTouched(): void $this->assertSame(["read $sessionId", "read $sessionId", "write $sessionId"], $this->log); } + #[TestDox('what a handler writes after it resumed is saved under the lock, as the code after a resume runs without it')] + public function testSessionWritesAfterResumeAreSavedUnderTheLock(): void + { + $lock = new RecordingSessionLock($this->log); + $sessionId = $this->existingSession(); + $fiber = $this->suspendedHandler($this->resumingProtocol($lock), $sessionId); + + $this->log = []; + $result = $fiber->resume(new Response(1000, [])); + + $this->assertTrue($fiber->isTerminated()); + $this->assertNull($result); + $this->assertInstanceOf(Response::class, $fiber->getReturn()); + $this->assertSame([ + "acquire $sessionId", + "read $sessionId", + "read $sessionId", + "write $sessionId", + "release $sessionId", + ], $this->log); + $this->assertTrue((new Session($this->store, $sessionId))->get('after')); + } + + #[TestDox('a resumed handler whose session stays locked still gets its response out, without its later writes')] + public function testResumedHandlerOnABusySessionIsStillAnswered(): void + { + $lock = new RecordingSessionLock($this->log); + $sessionId = $this->existingSession(); + $fiber = $this->suspendedHandler($this->resumingProtocol($lock), $sessionId); + + $lock->acquirable = false; + $this->log = []; + $fiber->resume(new Response(1000, [])); + + $this->assertInstanceOf(Response::class, $fiber->getReturn()); + $this->assertSame([], $this->log); + $this->assertSame([], $lock->heldSessions()); + $this->assertNull((new Session($this->store, $sessionId))->get('after')); + } + + #[TestDox('a session the client ended while the handler waited is not brought back by its later writes')] + public function testSessionDestroyedWhileTheHandlerWaitedStaysDestroyed(): void + { + $lock = new RecordingSessionLock($this->log); + $sessionId = $this->existingSession(); + $protocol = $this->resumingProtocol($lock); + $fiber = $this->suspendedHandler($protocol, $sessionId); + + $protocol->destroySession($sessionId); + $fiber->resume(new Response(1000, [])); + + $this->assertInstanceOf(Response::class, $fiber->getReturn()); + $this->assertFalse($this->store->exists($sessionId)); + $this->assertSame([], $lock->heldSessions()); + } + + private function resumingProtocol(RecordingSessionLock $lock): Protocol + { + $handler = new class implements RequestHandlerInterface { + public function supports(Request $request): bool + { + return true; + } + + public function handle(Request $request, SessionInterface $session): Response + { + $session->set('before', true); + \Fiber::suspend(new RequestSuspension(new PingRequest(), $session->getId()->toRfc4122(), 5)); + $session->set('after', true); + + return new Response($request->getId(), []); + } + }; + + return new Protocol( + requestHandlers: [$handler], + notificationHandlers: [], + messageFactory: MessageFactory::make(), + sessionManager: new SessionManager($this->store, gcProbability: 0), + sessionLock: $lock, + ); + } + + /** + * @return \Fiber + */ + private function suspendedHandler(Protocol $protocol, Uuid $sessionId): \Fiber + { + $transport = new class extends InMemoryTransport { + /** @var \Fiber|null */ + public ?\Fiber $fiber = null; + + public function attachFiberToSession(\Fiber $fiber, Uuid $sessionId): void + { + $this->fiber = $fiber; + } + }; + + $protocol->processInput($transport, '{"jsonrpc":"2.0","id":1,"method":"ping"}', $sessionId); + + $this->assertNotNull($transport->fiber); + $this->assertTrue($transport->fiber->isSuspended()); + + return $transport->fiber; + } + private function protocol(RecordingSessionLock $lock, LoggerInterface $logger = new NullLogger()): Protocol { return new Protocol( diff --git a/tests/Unit/Server/ProtocolTest.php b/tests/Unit/Server/ProtocolTest.php index 38c15ee3..cb5c204c 100644 --- a/tests/Unit/Server/ProtocolTest.php +++ b/tests/Unit/Server/ProtocolTest.php @@ -699,6 +699,64 @@ public function write(Uuid $id, string $data): bool $this->assertSame([], $protocol->getPendingRequests($sessionId)); } + #[TestDox('What a handler writes to its session after it resumed is stored, onto what the session holds by then')] + public function testSessionWritesAfterResumeAreSaved(): void + { + $handler = $this->createMock(RequestHandlerInterface::class); + $handler->method('supports')->willReturn(true); + $handler->method('handle')->willReturnCallback(static function (Request $request, SessionInterface $session): Response { + $session->set('before', true); + \Fiber::suspend(new RequestSuspension(new PingRequest(), $session->getId()->toRfc4122(), 5)); + $session->set('after', true); + $session->forget('dropped'); + + return new Response($request->getId(), []); + }); + + $sessionManager = new SessionManager(new InMemorySessionStore(), gcProbability: 0); + $sessionId = Uuid::v4(); + $session = $sessionManager->createWithId($sessionId); + $session->set('dropped', true); + $session->save(); + + $fiber = null; + $this->transport->method('attachFiberToSession')->willReturnCallback(static function (\Fiber $attached) use (&$fiber): void { + $fiber = $attached; + }); + + $protocol = new Protocol( + requestHandlers: [$handler], + notificationHandlers: [], + messageFactory: MessageFactory::make(), + sessionManager: $sessionManager, + ); + + $protocol->processInput($this->transport, '{"jsonrpc": "2.0", "id": 1, "method": "ping"}', $sessionId); + $this->assertInstanceOf(\Fiber::class, $fiber); + + // The request reaches the client and its answer comes back while the handler waits. + $pendingIds = array_keys($protocol->getPendingRequests($sessionId)); + $this->assertCount(1, $pendingIds); + $protocol->processInput($this->transport, \sprintf('{"jsonrpc": "2.0", "id": %s, "result": {}}', json_encode($pendingIds[0])), $sessionId); + $answer = $protocol->checkResponse(Uuid::fromString($pendingIds[0]), $sessionId); + $this->assertInstanceOf(Response::class, $answer); + $protocol->handleFiberYield(new NotificationSuspension(new LoggingMessageNotification(LoggingLevel::Info, 'ping'), $sessionId->toRfc4122()), $sessionId); + + $fiber->resume($answer); + $this->assertTrue($fiber->isTerminated()); + + $stored = $sessionManager->createWithId($sessionId); + $this->assertTrue($stored->get('before')); + $this->assertTrue($stored->get('after')); + $this->assertFalse($stored->has('dropped')); + // What happened to the session while the handler waited stays: the answered request is + // not pending again, the sent request is not queued again, the queued notification stays. + $this->assertSame([], $protocol->getPendingRequests($sessionId)); + $queued = $protocol->consumeOutgoingMessages($sessionId); + $this->assertCount(1, $queued); + $this->assertSame('notification', $queued[0]['context']['type']); + } + #[TestDox('A failing notification event listener does not produce a response')] public function testFailingNotificationListenerDoesNotProduceResponse(): void { From 2c87717880aacb991d7cd6cfc2be0d716e9d58cc Mon Sep 17 00:00:00 2001 From: Christopher Hertel Date: Thu, 15 Oct 2026 00:48:08 +0200 Subject: [PATCH 2/2] [Server] Address review: save the resumed session with Session::save() --- src/Server/Protocol.php | 57 ++++++------------- tests/Unit/Server/ProtocolSessionLockTest.php | 1 - tests/Unit/Server/ProtocolTest.php | 5 ++ 3 files changed, 23 insertions(+), 40 deletions(-) diff --git a/src/Server/Protocol.php b/src/Server/Protocol.php index 97cf592a..f17a4e5c 100644 --- a/src/Server/Protocol.php +++ b/src/Server/Protocol.php @@ -392,14 +392,14 @@ private function handleRequest(TransportInterface $transport, Request $request, $dispatchResult = $this->dispatchResult(...); $errorForThrowable = $this->errorForThrowable(...); $saveResumedSession = $this->saveResumedSession(...); - /** @var array|null $suspendedState what the session held once the fiber first suspended */ - $suspendedState = null; + /** @var bool $suspended set by reference once the fiber first suspended */ + $suspended = false; // One fiber for the whole exchange: with the shim, the handler // re-enters inside it each round rather than needing a new one. // Handled in the fiber, as the transport resumes it outside the catch below. /** @var McpFiber $fiber */ - $fiber = new \Fiber(static function () use ($handler, $request, $session, $shim, $codec, $dispatchResult, $errorForThrowable, $saveResumedSession, &$suspendedState): Response|Error { + $fiber = new \Fiber(static function () use ($handler, $request, $session, $shim, $codec, $dispatchResult, $errorForThrowable, $saveResumedSession, &$suspended): Response|Error { try { $result = $handler->handle($request, $session); $result = $shim?->fulfill($result, $handler, $request, $session, $codec) ?? $result; @@ -409,8 +409,8 @@ private function handleRequest(TransportInterface $transport, Request $request, } // A fiber that never suspended is saved by doProcessInput(); after a resume nothing else saves it. - if (null !== $suspendedState) { - $saveResumedSession($session, $suspendedState); + if ($suspended) { + $saveResumedSession($session); } return $result; @@ -447,7 +447,7 @@ private function handleRequest(TransportInterface $transport, Request $request, throw new RuntimeException('Failed to save the session of a suspended request.'); } - $suspendedState = $session->all(); + $suspended = true; $this->trackAwaitedRequest($fiber, $awaitedRequestId); $transport->attachFiberToSession($fiber, $session->getId()); @@ -963,52 +963,31 @@ private function dispatchResult(Response|Error $result, Request $request, Sessio } /** - * Stores what a handler wrote to its session after it resumed. - * - * The fiber's session holds the state it suspended with, while the stored one has moved - * on: the client's answer was consumed, the outgoing queue drained, and other requests of - * the session saved theirs. So only the keys the handler changed since it suspended are - * written onto the stored session, under the lock, as the code after a resume runs without it. - * - * @param array $suspendedState + * The code after a resume runs without the lock, so the session's changes since it suspended are saved under it. */ - private function saveResumedSession(SessionInterface $session, array $suspendedState): void + private function saveResumedSession(SessionInterface $session): void { - $state = $session->all(); - $changed = array_filter($state, static fn (mixed $value, int|string $key): bool => !\array_key_exists($key, $suspendedState) || $suspendedState[$key] !== $value, \ARRAY_FILTER_USE_BOTH); - $removed = array_diff_key($suspendedState, $state); - - if ([] === $changed && [] === $removed) { - return; - } - $sessionId = $session->getId(); - - if (!$this->acquireSessionLock($sessionId)) { - $this->logger->error('Lost the session writes of a resumed request, its session stayed locked.', ['session_id' => $sessionId->toRfc4122()]); - - return; - } + $locked = false; try { - // A session the client ended while the handler waited must not be brought back. - if (!$this->sessionManager->exists($sessionId)) { + if (!$locked = $this->acquireSessionLock($sessionId)) { + $this->logger->error('Lost the session writes of a resumed request, its session stayed locked.', ['session_id' => $sessionId->toRfc4122()]); + return; } - $stored = $this->sessionManager->createWithId($sessionId); - foreach ($changed as $key => $value) { - $stored->set((string) $key, $value); - } - foreach (array_keys($removed) as $key) { - $stored->forget((string) $key); + // A session the client ended while the handler waited is not saved back, and that is no failure. + if ($this->sessionManager->exists($sessionId) && !$session->save()) { + $this->logger->error('Failed to save the session of a resumed request.', ['session_id' => $sessionId->toRfc4122()]); } - $stored->save(); } catch (\Throwable $e) { // The response is already decided: an error for the same request would contradict it. $this->logger->error(\sprintf('Failed to save the session of a resumed request: %s', $e->getMessage()), ['exception' => $e]); } finally { - $this->releaseSessionLock($sessionId); + if ($locked) { + $this->releaseSessionLock($sessionId); + } } } diff --git a/tests/Unit/Server/ProtocolSessionLockTest.php b/tests/Unit/Server/ProtocolSessionLockTest.php index dd057b61..ca62a625 100644 --- a/tests/Unit/Server/ProtocolSessionLockTest.php +++ b/tests/Unit/Server/ProtocolSessionLockTest.php @@ -473,7 +473,6 @@ public function testSessionWritesAfterResumeAreSavedUnderTheLock(): void $this->assertSame([ "acquire $sessionId", "read $sessionId", - "read $sessionId", "write $sessionId", "release $sessionId", ], $this->log); diff --git a/tests/Unit/Server/ProtocolTest.php b/tests/Unit/Server/ProtocolTest.php index cb5c204c..2185c724 100644 --- a/tests/Unit/Server/ProtocolTest.php +++ b/tests/Unit/Server/ProtocolTest.php @@ -708,6 +708,7 @@ public function testSessionWritesAfterResumeAreSaved(): void $session->set('before', true); \Fiber::suspend(new RequestSuspension(new PingRequest(), $session->getId()->toRfc4122(), 5)); $session->set('after', true); + $session->set('shared.handler', true); $session->forget('dropped'); return new Response($request->getId(), []); @@ -741,6 +742,9 @@ public function testSessionWritesAfterResumeAreSaved(): void $answer = $protocol->checkResponse(Uuid::fromString($pendingIds[0]), $sessionId); $this->assertInstanceOf(Response::class, $answer); $protocol->handleFiberYield(new NotificationSuspension(new LoggingMessageNotification(LoggingLevel::Info, 'ping'), $sessionId->toRfc4122()), $sessionId); + $other = $sessionManager->createWithId($sessionId); + $other->set('shared.other', true); + $other->save(); $fiber->resume($answer); $this->assertTrue($fiber->isTerminated()); @@ -749,6 +753,7 @@ public function testSessionWritesAfterResumeAreSaved(): void $this->assertTrue($stored->get('before')); $this->assertTrue($stored->get('after')); $this->assertFalse($stored->has('dropped')); + $this->assertSame(['other' => true, 'handler' => true], $stored->get('shared')); // What happened to the session while the handler waited stays: the answered request is // not pending again, the sent request is not queued again, the queued notification stays. $this->assertSame([], $protocol->getPendingRequests($sessionId));