[Server] Fix StreamableHttpTransport mutable state and clock (#487)

src/Server/Transport/StreamableHttpTransport.php

@@ -15,6 +15,7 @@ use Http\Discovery\Psr17FactoryDiscovery;
 use Mcp\Exception\InvalidArgumentException;
 use Mcp\Schema\Enum\ProtocolVersion;
 use Mcp\Schema\JsonRpc\Error;
+use Mcp\Server\NativeClock;
 use Mcp\Server\Stateless\StatelessProtocol;
 use Mcp\Server\Transport\Http\Middleware\CorsMiddleware;
 use Mcp\Server\Transport\Http\Middleware\DnsRebindingProtectionMiddleware;
@@ -22,6 +23,7 @@ use Mcp\Server\Transport\Http\Middleware\ProtocolVersionMiddleware;
 use Mcp\Server\Transport\Http\MiddlewareRequestHandler;
 use Mcp\Server\Transport\Http\StatelessResponder;
 use Mcp\Server\Wire\InboundClassifier;
+use Psr\Clock\ClockInterface;
 use Psr\Http\Message\ResponseFactoryInterface;
 use Psr\Http\Message\ResponseInterface;
 use Psr\Http\Message\ServerRequestInterface;
@@ -79,12 +81,13 @@ class StreamableHttpTransport extends BaseTransport implements StatelessAwareTra
      * @param iterable<MiddlewareInterface>|null $middleware `null` installs {@see self::defaultMiddleware()}; `[]` disables all middleware
      */
     public function __construct(
-        private ServerRequestInterface $request,
+        private readonly ServerRequestInterface $request,
         ?ResponseFactoryInterface $responseFactory = null,
         ?StreamFactoryInterface $streamFactory = null,
         ?LoggerInterface $logger = null,
         ?iterable $middleware = null,
         private readonly int $maxBodyBytes = self::DEFAULT_MAX_BODY_BYTES,
+        private readonly ClockInterface $clock = new NativeClock(),
     ) {
         parent::__construct($logger);
 
@@ -185,12 +188,17 @@ class StreamableHttpTransport extends BaseTransport implements StatelessAwareTra
     {
         $this->handleMessage($body, $this->sessionId);
 
-        if (null !== $this->immediateResponse) {
+        // Consume the immediate response exactly once, so a transport instance
-            $response = $this->responseFactory->createResponse($this->immediateStatusCode ?? 200)
+        // reused for a later POST does not replay it.
-                ->withHeader('Content-Type', 'application/json')
+        $immediateResponse = $this->immediateResponse;
-                ->withBody($this->streamFactory->createStream($this->immediateResponse));
+        $immediateStatusCode = $this->immediateStatusCode;
+        $this->immediateResponse = null;
+        $this->immediateStatusCode = null;
 
-            return $response;
+        if (null !== $immediateResponse) {
+            return $this->responseFactory->createResponse($immediateStatusCode ?? 200)
+                ->withHeader('Content-Type', 'application/json')
+                ->withBody($this->streamFactory->createStream($immediateResponse));
         }
 
         if (null !== $this->sessionFiber) {
@@ -268,7 +276,7 @@ class StreamableHttpTransport extends BaseTransport implements StatelessAwareTra
                             break;
                         }
 
-                        if (time() - $timestamp >= $timeout) {
+                        if ($this->clock->now()->getTimestamp() - $timestamp >= $timeout) {
                             $error = Error::forInternalError('Request timed out', $requestId);
                             $yielded = $this->sessionFiber->resume($error);
                             $this->handleFiberYield($yielded, $this->sessionId);
@@ -376,8 +384,6 @@ class StreamableHttpTransport extends BaseTransport implements StatelessAwareTra
 
     private function handleRequest(ServerRequestInterface $request): ResponseInterface
     {
-        $this->request = $request;
-
         if ('OPTIONS' === $request->getMethod()) {
             return $this->handleOptionsRequest();
         }
@@ -405,7 +411,7 @@ class StreamableHttpTransport extends BaseTransport implements StatelessAwareTra
         }
 
         if ($classification->modern) {
-            return $this->handleModernRequest($body ?? '', $classification->claimedVersion ?? '');
+            return $this->handleModernRequest($request, $body ?? '', $classification->claimedVersion ?? '');
         }
 
         // The version-header rule only reaches the traffic it is about. Running
@@ -443,7 +449,7 @@ class StreamableHttpTransport extends BaseTransport implements StatelessAwareTra
     /**
      * Answers a request that claimed the modern era's per-request envelope.
      */
-    private function handleModernRequest(string $body, string $claimedVersion): ResponseInterface
+    private function handleModernRequest(ServerRequestInterface $request, string $body, string $claimedVersion): ResponseInterface
     {
         if (null === $this->stateless) {
             return $this->responder->error(
@@ -452,7 +458,7 @@ class StreamableHttpTransport extends BaseTransport implements StatelessAwareTra
             );
         }
 
-        return $this->responder->respond($this->stateless->handle($body, self::headers($this->request)));
+        return $this->responder->respond($this->stateless->handle($body, self::headers($request)));
     }
 
     /**

tests/Unit/Server/Transport/StreamableHttpTransportTest.php

@@ -12,19 +12,23 @@
 namespace Mcp\Tests\Unit\Server\Transport;
 
 use Mcp\Exception\InvalidArgumentException;
+use Mcp\Schema\JsonRpc\Error;
 use Mcp\Server\Transport\Http\Middleware\CorsMiddleware;
 use Mcp\Server\Transport\Http\Middleware\DnsRebindingProtectionMiddleware;
 use Mcp\Server\Transport\Http\Middleware\ProtocolVersionMiddleware;
 use Mcp\Server\Transport\StreamableHttpTransport;
+use Mcp\Server\Transport\TransportInterface;
 use Nyholm\Psr7\Factory\Psr17Factory;
 use PHPUnit\Framework\Attributes\TestDox;
 use PHPUnit\Framework\TestCase;
+use Psr\Clock\ClockInterface;
 use Psr\Http\Message\ResponseFactoryInterface;
 use Psr\Http\Message\ResponseInterface;
 use Psr\Http\Message\ServerRequestInterface;
 use Psr\Http\Server\MiddlewareInterface;
 use Psr\Http\Server\RequestHandlerInterface;
 use Psr\Log\LoggerInterface;
+use Symfony\Component\Uid\Uuid;
 
 final class StreamableHttpTransportTest extends TestCase
 {
@@ -347,6 +351,92 @@ final class StreamableHttpTransportTest extends TestCase
         new StreamableHttpTransport($request, $this->factory, $this->factory, null, [], maxBodyBytes: 0);
     }
 
+    #[TestDox('an immediate response is consumed once and not replayed on a later POST')]
+    public function testImmediateResponseIsNotReplayedOnSecondPost(): void
+    {
+        $request = $this->factory
+            ->createServerRequest('POST', 'http://localhost/')
+            ->withHeader('Host', 'localhost')
+            ->withBody($this->factory->createStream('{"jsonrpc":"2.0","id":1,"method":"ping"}'));
+
+        $transport = new StreamableHttpTransport($request, $this->factory, $this->factory);
+
+        $calls = 0;
+        $transport->onMessage(static function (TransportInterface $transport, string $payload) use (&$calls): void {
+            if (1 === ++$calls) {
+                $transport->send('{"jsonrpc":"2.0","id":1,"result":{}}', ['status_code' => 200]);
+            }
+        });
+
+        $first = $transport->listen();
+
+        $this->assertSame(200, $first->getStatusCode());
+        $this->assertSame('{"jsonrpc":"2.0","id":1,"result":{}}', (string) $first->getBody());
+
+        $second = $transport->listen();
+
+        $this->assertSame(2, $calls);
+        $this->assertSame(202, $second->getStatusCode());
+        $this->assertSame('', (string) $second->getBody());
+    }
+
+    #[TestDox('the polling loop times out a pending request via the injected clock')]
+    public function testPollingLoopTimesOutPendingRequestViaInjectedClock(): void
+    {
+        $request = $this->factory
+            ->createServerRequest('POST', 'http://localhost/')
+            ->withHeader('Host', 'localhost')
+            ->withBody($this->factory->createStream('{"jsonrpc":"2.0","id":1,"method":"ping"}'));
+
+        $requestedAt = 1_000_000;
+
+        // Frozen 121s after the pending request was issued — past its 120s timeout.
+        $clock = new class($requestedAt + 121) implements ClockInterface {
+            public function __construct(private readonly int $timestamp)
+            {
+            }
+
+            public function now(): \DateTimeImmutable
+            {
+                return (new \DateTimeImmutable())->setTimestamp($this->timestamp);
+            }
+        };
+
+        $transport = new StreamableHttpTransport(
+            $request,
+            $this->factory,
+            $this->factory,
+            clock: $clock,
+        );
+
+        $received = null;
+        $fiber = new \Fiber(static function () use (&$received) {
+            $received = \Fiber::suspend();
+
+            return null;
+        });
+        $fiber->start();
+
+        $transport->onMessage(static function (TransportInterface $transport) use ($fiber): void {
+            $transport->attachFiberToSession($fiber, Uuid::v4());
+        });
+        $transport->setOutgoingMessagesProvider(static fn (): array => []);
+        $transport->setResponseFinder(static fn () => null);
+        $transport->setPendingRequestsProvider(static fn (): array => [
+            ['request_id' => 1, 'timestamp' => $requestedAt, 'timeout' => 120],
+        ]);
+
+        $response = $transport->listen();
+
+        $this->assertSame('text/event-stream', $response->getHeaderLine('Content-Type'));
+
+        $this->expectOutputString('');
+        $response->getBody()->getContents();
+
+        $this->assertTrue($fiber->isTerminated());
+        $this->assertInstanceOf(Error::class, $received);
+    }
+
     private function stubAuth401(): MiddlewareInterface
     {
         return new class($this->factory) implements MiddlewareInterface {