[Server] Save the session writes of a handler after it resumed

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
 -----

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) {

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<null, mixed, mixed, mixed>
+     */
+    private function suspendedHandler(Protocol $protocol, Uuid $sessionId): \Fiber
+    {
+        $transport = new class extends InMemoryTransport {
+            /** @var \Fiber<null, mixed, mixed, mixed>|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(

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
     {