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..f17a4e5c 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 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): 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; + $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 ($suspended) { + $saveResumedSession($session); + } + + 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.'); } + $suspended = true; + $this->trackAwaitedRequest($fiber, $awaitedRequestId); $transport->attachFiberToSession($fiber, $session->getId()); @@ -951,6 +962,35 @@ private function dispatchResult(Response|Error $result, Request $request, Sessio } } + /** + * 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): void + { + $sessionId = $session->getId(); + $locked = false; + + try { + 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; + } + + // 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()]); + } + } 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 { + if ($locked) { + $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..ca62a625 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,111 @@ 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", + "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..2185c724 100644 --- a/tests/Unit/Server/ProtocolTest.php +++ b/tests/Unit/Server/ProtocolTest.php @@ -699,6 +699,69 @@ 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->set('shared.handler', 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); + $other = $sessionManager->createWithId($sessionId); + $other->set('shared.other', true); + $other->save(); + + $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')); + $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)); + $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 {