[Server] Drop pending client requests of a vanished stream

CHANGELOG.md

@@ -85,6 +85,7 @@ All notable changes to `mcp/sdk` will be documented in this file.
 * `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.
+* Fix pending client requests of a vanished stream, and their late answers, staying in the session until it expires: they are dropped once a minute past their timeout.
 
 0.8.1
 -----

src/Server/Protocol.php

@@ -80,6 +80,12 @@ class Protocol
     private const SESSION_BUSY_MESSAGE = 'Another request of this session is still being processed.';
     private const SESSION_LOCK_UNAVAILABLE_MESSAGE = 'The session lock is unavailable.';
 
+    /**
+     * Seconds past its timeout after which a pending client request counts as abandoned by its stream.
+     * Until then, the stream that sent it expires it itself and reports the timeout.
+     */
+    private const ABANDONED_REQUEST_GRACE = 60;
+
     /**
      * The client request each suspended fiber waits on. Pending requests are stored in the session,
      * which concurrent calls share, so a fiber must only be resumed with the answer to the one it sent.
@@ -96,6 +102,14 @@ class Protocol
      */
     private \WeakMap $fiberOutgoingMessages;
 
+    /**
+     * Awaited requests another worker dropped as abandoned, e.g. while this one could not get the session
+     * lock to expire them: their fiber is still owed the timeout.
+     *
+     * @var array<string, true>
+     */
+    private array $droppedAwaitedRequestIds = [];
+
     /**
      * @param array<int, RequestHandlerInterface<ResultInterface|array<string, mixed>>> $requestHandlers
      * @param array<int, NotificationHandlerInterface>                                  $notificationHandlers
@@ -559,6 +573,8 @@ private function registerRequest(Request $request, int $timeout, SessionInterfac
             'method' => $request::getMethod(),
         ]);
 
+        $this->dropAbandonedRequests($session);
+
         // Per entry, so saving does not write back what other streams changed in the others.
         $session->set(self::SESSION_PENDING_REQUESTS.".{$id}", [
             'request_id' => $id,
@@ -736,6 +752,17 @@ public function checkResponse(Uuid $requestId, Uuid $sessionId): Response|Error|
     {
         $id = $requestId->toRfc4122();
 
+        if (isset($this->droppedAwaitedRequestIds[$id])) {
+            unset($this->droppedAwaitedRequestIds[$id]);
+            $this->settleAwaitedRequest($id);
+            $this->logger->warning('Client request timed out.', [
+                'request_id' => $id,
+                'session_id' => $sessionId->toRfc4122(),
+            ]);
+
+            return Error::forInternalError('Request timed out', $id);
+        }
+
         if (!$this->tryAcquireSessionLockForPoll($sessionId)) {
             return null;
         }
@@ -758,6 +785,7 @@ public function checkResponse(Uuid $requestId, Uuid $sessionId): Response|Error|
             $session->set(self::SESSION_RESPONSES.".{$id}", null);
             $session->forget(self::SESSION_PENDING_REQUESTS.".{$id}");
             $session->save();
+            $this->settleAwaitedRequest($id);
         } finally {
             $this->releaseSessionLock($sessionId);
         }
@@ -803,10 +831,43 @@ private function expireRequest(Uuid $requestId, Uuid $sessionId): ?Error
 
         $session->forget(self::SESSION_PENDING_REQUESTS.".{$id}");
         $session->save();
+        $this->settleAwaitedRequest($id);
 
         return Error::forInternalError('Request timed out', $id);
     }
 
+    /**
+     * Drops the pending client requests, and their unconsumed answers, that no stream waits for anymore.
+     *
+     * Only the stream that sent a request clears it, so one whose stream is gone, through a client
+     * disconnect or a dead worker, would stay until the session expires.
+     */
+    private function dropAbandonedRequests(SessionInterface $session): void
+    {
+        // A fiber of this process still waits on these, e.g. on stdio while another handler blocked the loop.
+        $awaited = [];
+        foreach ($this->awaitedRequestIds as $requestId) {
+            $awaited[$requestId->toRfc4122()] = true;
+        }
+
+        foreach ($session->get(self::SESSION_PENDING_REQUESTS, []) as $id => $pending) {
+            if (isset($awaited[$id]) || !\is_array($pending) || time() < (int) ($pending['timestamp'] ?? 0) + (int) ($pending['timeout'] ?? 120) + self::ABANDONED_REQUEST_GRACE) {
+                continue;
+            }
+
+            $this->logger->debug('Dropping abandoned client request.', [
+                'request_id' => $id,
+                'session_id' => $session->getId()->toRfc4122(),
+            ]);
+
+            // Per entry, so saving does not write back what other streams changed in the others. The answer
+            // may have been stored after this session was loaded: set() makes forget() record its removal.
+            $session->forget(self::SESSION_PENDING_REQUESTS.".{$id}");
+            $session->set(self::SESSION_RESPONSES.".{$id}", null);
+            $session->forget(self::SESSION_RESPONSES.".{$id}");
+        }
+    }
+
     /**
      * @param array<mixed> $pending
      */
@@ -928,6 +989,18 @@ private function trackAwaitedRequest(\Fiber $fiber, ?Uuid $requestId): void
         $this->awaitedRequestIds[$fiber] = $requestId;
     }
 
+    /**
+     * The fiber that awaited this request gets its answer now, so a later poll must not take it for one dropped elsewhere.
+     */
+    private function settleAwaitedRequest(string $id): void
+    {
+        foreach ($this->awaitedRequestIds as $fiber => $requestId) {
+            if ($requestId->toRfc4122() === $id) {
+                unset($this->awaitedRequestIds[$fiber]);
+            }
+        }
+    }
+
     /**
      * @param McpFiber $fiber
      *
@@ -940,7 +1013,16 @@ private function getAwaitedPendingRequests(\Fiber $fiber, Uuid $sessionId): arra
             return [];
         }
 
-        return array_intersect_key($this->getPendingRequests($sessionId), [$requestId->toRfc4122() => true]);
+        $id = $requestId->toRfc4122();
+        $pending = $this->getPendingRequests($sessionId)[$id] ?? null;
+
+        if (null === $pending) {
+            $this->droppedAwaitedRequestIds[$id] = true;
+
+            return [$id => ['request_id' => $id]];
+        }
+
+        return [$id => $pending];
     }
 
     /**

tests/Unit/Server/ProtocolSessionLockTest.php

@@ -26,6 +26,7 @@
 use Mcp\Server\Session\SessionManagerInterface;
 use Mcp\Server\Suspension\RequestSuspension;
 use Mcp\Server\Transport\InMemoryTransport;
+use Mcp\Tests\Unit\Fixtures\PollingLoopTransport;
 use Mcp\Tests\Unit\Fixtures\RecordingTransport;
 use Mcp\Tests\Unit\Server\Session\Fixture\RecordingLogger;
 use Mcp\Tests\Unit\Server\Session\Fixture\RecordingSessionLock;
@@ -512,6 +513,41 @@ public function testSessionDestroyedWhileTheHandlerWaitedStaysDestroyed(): void
         $this->assertSame([], $lock->heldSessions());
     }
 
+    #[TestDox('a request another worker dropped as abandoned while the lock kept its stream from expiring it still times out')]
+    public function testAwaitedRequestDroppedByAnotherWorkerTimesOut(): void
+    {
+        $lock = new RecordingSessionLock($this->log);
+        $waiting = $this->resumingProtocol($lock);
+        $sessionId = $this->existingSession();
+        $stream = new PollingLoopTransport();
+        $waiting->connect($stream);
+        $waiting->processInput($stream, '{"jsonrpc":"2.0","id":1,"method":"ping"}', $sessionId);
+        $this->assertCount(1, $pendingRequestIds = $stream->getPendingRequestIds());
+        $requestId = Uuid::fromString($pendingRequestIds[0]);
+
+        $session = new Session($this->store, $sessionId);
+        $session->set("_mcp.pending_requests.{$pendingRequestIds[0]}.timestamp", time() - 100);
+        $session->save();
+
+        // Another worker holds the session past the request's timeout, then sends a request of its own.
+        $lock->acquirable = false;
+        $this->assertNull($waiting->checkResponse($requestId, $sessionId));
+        $lock->acquirable = true;
+
+        $sweeping = $this->protocol($lock);
+        $session = new Session($this->store, $sessionId);
+        $sweeping->sendRequest(new PingRequest(), 5, $session);
+        $session->save();
+        $this->assertArrayNotHasKey($pendingRequestIds[0], (new Session($this->store, $sessionId))->get('_mcp.pending_requests'));
+
+        $this->assertSame($pendingRequestIds, $stream->getPendingRequestIds());
+        $result = $waiting->checkResponse($requestId, $sessionId);
+        $this->assertInstanceOf(Error::class, $result);
+        $this->assertSame('Request timed out', $result->message);
+        $this->assertSame([], $stream->getPendingRequestIds());
+        $this->assertNull($waiting->checkResponse($requestId, $sessionId), 'The timeout is reported only once.');
+    }
+
     private function resumingProtocol(RecordingSessionLock $lock): Protocol
     {
         $handler = new class implements RequestHandlerInterface {

tests/Unit/Server/ProtocolTest.php

@@ -2291,6 +2291,104 @@ public function testAnswerAfterReportedTimeoutIsDropped(): void
         $this->assertNull($protocol->checkResponse(Uuid::fromString('00000000-0000-4000-8000-000000001000'), $session->getId()));
     }
 
+    #[TestDox('A client request its stream abandoned is dropped with its answer once another one is sent')]
+    public function testAbandonedRequestIsDroppedWhenAnotherRequestIsSent(): void
+    {
+        [$protocol, $sessionManager, $session] = $this->createProtocolWithPendingRequests([
+            '00000000-0000-4000-8000-000000001000' => ['request_id' => '00000000-0000-4000-8000-000000001000', 'timeout' => 5, 'timestamp' => time() - 100],
+        ]);
+        $session->set('_mcp.responses.00000000-0000-4000-8000-000000001000', ['jsonrpc' => '2.0', 'id' => '00000000-0000-4000-8000-000000001000', 'result' => []]);
+        $session->save();
+
+        $session = $sessionManager->createWithId($session->getId());
+        $requestId = $protocol->sendRequest(new PingRequest(), 5, $session);
+        $session->save();
+
+        $session = $sessionManager->createWithId($session->getId());
+        $this->assertSame([$requestId->toRfc4122()], array_keys($session->get('_mcp.pending_requests')));
+        $this->assertArrayNotHasKey('00000000-0000-4000-8000-000000001000', $session->get('_mcp.responses', []));
+    }
+
+    #[TestDox('An answer to an abandoned client request is dropped even if it arrived after the session was loaded')]
+    public function testAnswerStoredAfterLoadingIsDroppedWithAbandonedRequest(): void
+    {
+        [$protocol, $sessionManager, $session] = $this->createProtocolWithPendingRequests([
+            '00000000-0000-4000-8000-000000001000' => ['request_id' => '00000000-0000-4000-8000-000000001000', 'timeout' => 5, 'timestamp' => time() - 100],
+        ]);
+
+        $snapshot = $sessionManager->createWithId($session->getId());
+        $snapshot->get('_mcp.pending_requests');
+
+        $other = $sessionManager->createWithId($session->getId());
+        $other->set('_mcp.responses.00000000-0000-4000-8000-000000001000', ['jsonrpc' => '2.0', 'id' => '00000000-0000-4000-8000-000000001000', 'result' => []]);
+        $other->save();
+
+        $protocol->sendRequest(new PingRequest(), 5, $snapshot);
+        $snapshot->save();
+
+        $this->assertArrayNotHasKey('00000000-0000-4000-8000-000000001000', $sessionManager->createWithId($session->getId())->get('_mcp.responses', []));
+    }
+
+    #[TestDox('A client request a suspended fiber of the same process still waits on is not dropped as abandoned')]
+    public function testRequestAwaitedByLocalFiberIsKeptPastGracePeriod(): void
+    {
+        $handler = $this->createMock(RequestHandlerInterface::class);
+        $handler->method('supports')->willReturn(true);
+        $handler->method('handle')->willReturnCallback(static function (Request $request, SessionInterface $session): Response {
+            \Fiber::suspend(new RequestSuspension(new PingRequest(), $session->getId()->toRfc4122(), 5));
+
+            return new Response($request->getId(), []);
+        });
+
+        $sessionManager = new SessionManager(new InMemorySessionStore());
+        $protocol = new Protocol(
+            requestHandlers: [$handler],
+            notificationHandlers: [],
+            messageFactory: MessageFactory::make(),
+            sessionManager: $sessionManager,
+        );
+
+        $session = $sessionManager->create();
+        $session->save();
+        $sessionId = $session->getId();
+
+        $stream = new PollingLoopTransport();
+        $protocol->connect($stream);
+        $protocol->processInput($stream, '{"jsonrpc": "2.0", "id": 1, "method": "ping"}', $sessionId);
+        $this->assertCount(1, $pendingRequestIds = $stream->getPendingRequestIds());
+
+        // The loop that would poll the fiber was blocked well past the request's timeout.
+        $session = $sessionManager->createWithId($sessionId);
+        $session->set("_mcp.pending_requests.{$pendingRequestIds[0]}.timestamp", time() - 100);
+        $session->save();
+
+        $session = $sessionManager->createWithId($sessionId);
+        $protocol->sendRequest(new PingRequest(), 5, $session);
+        $session->save();
+
+        $this->assertArrayHasKey($pendingRequestIds[0], $protocol->getPendingRequests($sessionId));
+        $this->assertSame($pendingRequestIds, $stream->getPendingRequestIds());
+    }
+
+    #[TestDox('A timed out client request within the grace period is left to its own stream')]
+    public function testTimedOutRequestWithinGracePeriodIsKeptWhenAnotherRequestIsSent(): void
+    {
+        [$protocol, $sessionManager, $session] = $this->createProtocolWithPendingRequests([
+            '00000000-0000-4000-8000-000000001000' => ['request_id' => '00000000-0000-4000-8000-000000001000', 'timeout' => 5, 'timestamp' => time() - 10],
+        ]);
+
+        $session = $sessionManager->createWithId($session->getId());
+        $protocol->sendRequest(new PingRequest(), 5, $session);
+        $session->save();
+
+        $this->assertArrayHasKey('00000000-0000-4000-8000-000000001000', $protocol->getPendingRequests($session->getId()));
+
+        $result = $protocol->checkResponse(Uuid::fromString('00000000-0000-4000-8000-000000001000'), $session->getId());
+
+        $this->assertInstanceOf(Error::class, $result);
+        $this->assertSame('Request timed out', $result->message);
+    }
+
     #[TestDox('ClientResponseEvent::getId() is null when the error has no id')]
     public function testClientResponseEventGetIdCanBeNull(): void
     {