Add Middleware handlers to StreamableHttpTransport (#218)

src/Server/Transport/StreamableHttpTransport.php

@@ -12,11 +12,14 @@
 namespace Mcp\Server\Transport;
 
 use Http\Discovery\Psr17FactoryDiscovery;
+use Mcp\Exception\InvalidArgumentException;
 use Mcp\Schema\JsonRpc\Error;
+use Mcp\Server\Transport\Http\MiddlewareRequestHandler;
 use Psr\Http\Message\ResponseFactoryInterface;
 use Psr\Http\Message\ResponseInterface;
 use Psr\Http\Message\ServerRequestInterface;
 use Psr\Http\Message\StreamFactoryInterface;
+use Psr\Http\Server\MiddlewareInterface;
 use Psr\Log\LoggerInterface;
 use Symfony\Component\Uid\Uuid;
 
@@ -36,19 +39,22 @@ class StreamableHttpTransport extends BaseTransport
     /** @var array<string, string> */
     private array $corsHeaders;
 
+    /** @var list<MiddlewareInterface> */
+    private array $middleware = [];
+
     /**
-     * @param array<string, string> $corsHeaders
+     * @param array<string, string>         $corsHeaders
+     * @param iterable<MiddlewareInterface> $middleware
      */
     public function __construct(
-        private readonly ServerRequestInterface $request,
+        private ServerRequestInterface $request,
         ?ResponseFactoryInterface $responseFactory = null,
         ?StreamFactoryInterface $streamFactory = null,
         array $corsHeaders = [],
         ?LoggerInterface $logger = null,
+        iterable $middleware = [],
     ) {
         parent::__construct($logger);
-        $sessionIdString = $this->request->getHeaderLine('Mcp-Session-Id');
-        $this->sessionId = $sessionIdString ? Uuid::fromString($sessionIdString) : null;
 
         $this->responseFactory = $responseFactory ?? Psr17FactoryDiscovery::findResponseFactory();
         $this->streamFactory = $streamFactory ?? Psr17FactoryDiscovery::findStreamFactory();
@@ -59,6 +65,13 @@ class StreamableHttpTransport extends BaseTransport
             'Access-Control-Allow-Headers' => 'Content-Type, Mcp-Session-Id, Mcp-Protocol-Version, Last-Event-ID, Authorization, Accept',
             'Access-Control-Expose-Headers' => 'Mcp-Session-Id',
         ], $corsHeaders);
+
+        foreach ($middleware as $m) {
+            if (!$m instanceof MiddlewareInterface) {
+                throw new InvalidArgumentException('Streamable HTTP middleware must implement Psr\\Http\\Server\\MiddlewareInterface.');
+            }
+            $this->middleware[] = $m;
+        }
     }
 
     public function send(string $data, array $context): void
@@ -69,17 +82,17 @@ class StreamableHttpTransport extends BaseTransport
 
     public function listen(): ResponseInterface
     {
-        return match ($this->request->getMethod()) {
-            'OPTIONS' => $this->handleOptionsRequest(),
-            'POST' => $this->handlePostRequest(),
-            'DELETE' => $this->handleDeleteRequest(),
-            default => $this->createErrorResponse(Error::forInvalidRequest('Method Not Allowed'), 405),
-        };
+        $handler = new MiddlewareRequestHandler(
+            $this->middleware,
+            \Closure::fromCallable([$this, 'handleRequest']),
+        );
+
+        return $this->withCorsHeaders($handler->handle($this->request));
     }
 
     protected function handleOptionsRequest(): ResponseInterface
     {
-        return $this->withCorsHeaders($this->responseFactory->createResponse(204));
+        return $this->responseFactory->createResponse(204);
     }
 
     protected function handlePostRequest(): ResponseInterface
@@ -92,7 +105,7 @@ class StreamableHttpTransport extends BaseTransport
                 ->withHeader('Content-Type', 'application/json')
                 ->withBody($this->streamFactory->createStream($this->immediateResponse));
 
-            return $this->withCorsHeaders($response);
+            return $response;
         }
 
         if (null !== $this->sessionFiber) {
@@ -112,7 +125,7 @@ class StreamableHttpTransport extends BaseTransport
 
         $this->handleSessionEnd($this->sessionId);
 
-        return $this->withCorsHeaders($this->responseFactory->createResponse(200));
+        return $this->responseFactory->createResponse(200);
     }
 
     protected function createJsonResponse(): ResponseInterface
@@ -120,7 +133,7 @@ class StreamableHttpTransport extends BaseTransport
         $outgoingMessages = $this->getOutgoingMessages($this->sessionId);
 
         if (empty($outgoingMessages)) {
-            return $this->withCorsHeaders($this->responseFactory->createResponse(202));
+            return $this->responseFactory->createResponse(202);
         }
 
         $messages = array_column($outgoingMessages, 'message');
@@ -134,7 +147,7 @@ class StreamableHttpTransport extends BaseTransport
             $response = $response->withHeader('Mcp-Session-Id', $this->sessionId->toRfc4122());
         }
 
-        return $this->withCorsHeaders($response);
+        return $response;
     }
 
     protected function createStreamedResponse(): ResponseInterface
@@ -201,7 +214,7 @@ class StreamableHttpTransport extends BaseTransport
             $response = $response->withHeader('Mcp-Session-Id', $this->sessionId->toRfc4122());
         }
 
-        return $this->withCorsHeaders($response);
+        return $response;
     }
 
     protected function handleFiberTermination(): void
@@ -246,15 +259,31 @@ class StreamableHttpTransport extends BaseTransport
             $response = $response->withHeader('Allow', 'POST, DELETE, OPTIONS');
         }
 
-        return $this->withCorsHeaders($response);
+        return $response;
     }
 
     protected function withCorsHeaders(ResponseInterface $response): ResponseInterface
     {
         foreach ($this->corsHeaders as $name => $value) {
-            $response = $response->withHeader($name, $value);
+            if (!$response->hasHeader($name)) {
+                $response = $response->withHeader($name, $value);
+            }
         }
 
         return $response;
     }
+
+    private function handleRequest(ServerRequestInterface $request): ResponseInterface
+    {
+        $this->request = $request;
+        $sessionIdString = $request->getHeaderLine('Mcp-Session-Id');
+        $this->sessionId = $sessionIdString ? Uuid::fromString($sessionIdString) : null;
+
+        return match ($request->getMethod()) {
+            'OPTIONS' => $this->handleOptionsRequest(),
+            'POST' => $this->handlePostRequest(),
+            'DELETE' => $this->handleDeleteRequest(),
+            default => $this->createErrorResponse(Error::forInvalidRequest('Method Not Allowed'), 405),
+        };
+    }
 }