feat(server): Integrate session management in message handling

- Added session support to the `Server` and `Handler` classes, allowing for session data to be managed during message processing.
- Updated `TransportInterface` to include session context in the `send` method.
- Refactored various request handlers to utilize session information, ensuring proper session handling for incoming requests.
- Introduced a file-based session store for persistent session data management

Author: CodeWithKyrian

examples/10-simple-http-transport/.gitignore

@@ -0,0 +1 @@
+sessions/

examples/10-simple-http-transport/server.php

@@ -8,6 +8,7 @@
 use Nyholm\Psr7\Factory\Psr17Factory;
 use Nyholm\Psr7Server\ServerRequestCreator;
 use Laminas\HttpHandlerRunner\Emitter\SapiEmitter;
+use Mcp\Server\Session\FileSessionStore;
 
 $psr17Factory = new Psr17Factory();
 $creator = new ServerRequestCreator($psr17Factory, $psr17Factory, $psr17Factory, $psr17Factory);
@@ -17,6 +18,7 @@
 $server = Server::make()
     ->withServerInfo('HTTP MCP Server', '1.0.0', 'MCP Server over HTTP transport')
     ->withContainer(container())
+    ->withSession(new FileSessionStore(__DIR__ . '/sessions'))
     ->withDiscovery(__DIR__, ['.'])
     ->build();
 

src/JsonRpc/Handler.php

@@ -25,8 +25,12 @@
 use Mcp\Schema\JsonRpc\HasMethodInterface;
 use Mcp\Schema\JsonRpc\Response;
 use Mcp\Server\MethodHandlerInterface;
+use Mcp\Server\Session\SessionInterface;
 use Mcp\Server\NotificationHandler;
 use Mcp\Server\RequestHandler;
+use Mcp\Server\Session\SessionFactoryInterface;
+use Mcp\Server\Session\SessionStoreInterface;
+use Mcp\Schema\Request\InitializeRequest;
 use Symfony\Component\Uid\Uuid;
 use Psr\Log\LoggerInterface;
 use Psr\Log\NullLogger;
@@ -48,6 +52,8 @@ class Handler
      */
     public function __construct(
         private readonly MessageFactory $messageFactory,
+        private readonly SessionFactoryInterface $sessionFactory,
+        private readonly SessionStoreInterface $sessionStore,
         iterable $methodHandlers,
         private readonly LoggerInterface $logger = new NullLogger(),
     ) {
@@ -63,10 +69,14 @@ public static function make(
         ToolCallerInterface $toolCaller,
         ResourceReaderInterface $resourceReader,
         PromptGetterInterface $promptGetter,
+        SessionStoreInterface $sessionStore,
+        SessionFactoryInterface $sessionFactory,
         LoggerInterface $logger = new NullLogger(),
     ): self {
         return new self(
             messageFactory: MessageFactory::make(),
+            sessionFactory: $sessionFactory,
+            sessionStore: $sessionStore,
             methodHandlers: [
                 new NotificationHandler\InitializedHandler(),
                 new RequestHandler\InitializeHandler($registry->getCapabilities(), $implementation),
@@ -83,7 +93,7 @@ public static function make(
     }
 
     /**
-     * @return iterable<string|null>
+     * @return iterable<array{string|null, array<string, mixed>}>
      *
      * @throws ExceptionInterface When a handler throws an exception during message processing
      * @throws \JsonException     When JSON encoding of the response fails
@@ -92,20 +102,65 @@ public function process(string $input, ?Uuid $sessionId): iterable
     {
         $this->logger->info('Received message to process.', ['message' => $input]);
 
+        $this->runGarbageCollection();
+
         try {
-            $messages = $this->messageFactory->create($input);
+            $messages = iterator_to_array($this->messageFactory->create($input));
         } catch (\JsonException $e) {
             $this->logger->warning('Failed to decode json message.', ['exception' => $e]);
-
-            yield $this->encodeResponse(Error::forParseError($e->getMessage()));
+            $error = Error::forParseError($e->getMessage());
+            yield [$this->encodeResponse($error), []];
 
             return;
         }
 
+        $hasInitializeRequest = false;
+        foreach ($messages as $message) {
+            if ($message instanceof InitializeRequest) {
+                $hasInitializeRequest = true;
+                break;
+            }
+        }
+
+        $session = null;
+
+        if ($hasInitializeRequest) {
+            // Spec: An initialize request must not be part of a batch.
+            if (count($messages) > 1) {
+                $error = Error::forInvalidRequest('The "initialize" request MUST NOT be part of a batch.');
+                yield [$this->encodeResponse($error), []];
+                return;
+            }
+
+            // Spec: An initialize request must not have a session ID.
+            if ($sessionId) {
+                $error = Error::forInvalidRequest('A session ID MUST NOT be sent with an "initialize" request.');
+                yield [$this->encodeResponse($error), []];
+                return;
+            }
+
+            $session = $this->sessionFactory->create($this->sessionStore);
+        } else {
+            if (!$sessionId) {
+                $error = Error::forInvalidRequest('A valid session id is REQUIRED for non-initialize requests.');
+                yield [$this->encodeResponse($error), ['status_code' => 400]];
+                return;
+            }
+
+            if (!$this->sessionStore->exists($sessionId)) {
+                $error = Error::forInvalidRequest('Session not found or has expired.');
+                yield [$this->encodeResponse($error), ['status_code' => 404]];
+                return;
+            }
+
+            $session = $this->sessionFactory->createWithId($sessionId, $this->sessionStore);
+        }
+
         foreach ($messages as $message) {
             if ($message instanceof InvalidInputMessageException) {
                 $this->logger->warning('Failed to create message.', ['exception' => $message]);
-                yield $this->encodeResponse(Error::forInvalidRequest($message->getMessage(), 0));
+                $error = Error::forInvalidRequest($message->getMessage(), 0);
+                yield [$this->encodeResponse($error), []];
                 continue;
             }
 
@@ -114,26 +169,32 @@ public function process(string $input, ?Uuid $sessionId): iterable
             ]);
 
             try {
-                yield $this->encodeResponse($this->handle($message));
+                $response = $this->encodeResponse($this->handle($message, $session));
+                yield [$response, ['session_id' => $session->getId()]];
             } catch (\DomainException) {
-                yield null;
+                yield [null, []];
             } catch (NotFoundExceptionInterface $e) {
                 $this->logger->warning(
                     \sprintf('Failed to create response: %s', $e->getMessage()),
                     ['exception' => $e],
                 );
 
-                yield $this->encodeResponse(Error::forMethodNotFound($e->getMessage()));
+                $error = Error::forMethodNotFound($e->getMessage());
+                yield [$this->encodeResponse($error), []];
             } catch (\InvalidArgumentException $e) {
                 $this->logger->warning(\sprintf('Invalid argument: %s', $e->getMessage()), ['exception' => $e]);
 
-                yield $this->encodeResponse(Error::forInvalidParams($e->getMessage()));
+                $error = Error::forInvalidParams($e->getMessage());
+                yield [$this->encodeResponse($error), []];
             } catch (\Throwable $e) {
                 $this->logger->critical(\sprintf('Uncaught exception: %s', $e->getMessage()), ['exception' => $e]);
 
-                yield $this->encodeResponse(Error::forInternalError($e->getMessage()));
+                $error = Error::forInternalError($e->getMessage());
+                yield [$this->encodeResponse($error), []];
             }
         }
+
+        $session->save();
     }
 
     /**
@@ -162,26 +223,28 @@ private function encodeResponse(Response|Error|null $response): ?string
      * @throws NotFoundExceptionInterface When no handler is found for the request method
      * @throws ExceptionInterface         When a request handler throws an exception
      */
-    private function handle(HasMethodInterface $message): Response|Error|null
+    private function handle(HasMethodInterface $message, SessionInterface $session): Response|Error|null
     {
         $this->logger->info(\sprintf('Handling message for method "%s".', $message::getMethod()), [
             'message' => $message,
         ]);
 
         $handled = false;
         foreach ($this->methodHandlers as $handler) {
-            if ($handler->supports($message)) {
-                $return = $handler->handle($message);
-                $handled = true;
-
-                $this->logger->debug(\sprintf('Message handled by "%s".', $handler::class), [
-                    'method' => $message::getMethod(),
-                    'response' => $return,
-                ]);
-
-                if (null !== $return) {
-                    return $return;
-                }
+            if (!$handler->supports($message)) {
+                continue;
+            }
+
+            $return = $handler->handle($message, $session);
+            $handled = true;
+
+            $this->logger->debug(\sprintf('Message handled by "%s".', $handler::class), [
+                'method' => $message::getMethod(),
+                'response' => $return,
+            ]);
+
+            if (null !== $return) {
+                return $return;
             }
         }
 
@@ -191,4 +254,32 @@ private function handle(HasMethodInterface $message): Response|Error|null
 
         throw new HandlerNotFoundException(\sprintf('No handler found for method "%s".', $message::getMethod()));
     }
+
+    /**
+     * Run garbage collection on expired sessions.
+     * Uses the session store's internal TTL configuration.
+     */
+    private function runGarbageCollection(): void
+    {
+        if (random_int(0, 100) > 1) {
+            return;
+        }
+
+        $deletedSessions = $this->sessionStore->gc();
+        if (!empty($deletedSessions)) {
+            $this->logger->debug('Garbage collected expired sessions.', [
+                'count' => count($deletedSessions),
+                'session_ids' => array_map(fn(Uuid $id) => $id->toRfc4122(), $deletedSessions),
+            ]);
+        }
+    }
+
+    /**
+     * Destroy a specific session.
+     */
+    public function destroySession(Uuid $sessionId): void
+    {
+        $this->sessionStore->destroy($sessionId);
+        $this->logger->info('Session destroyed.', ['session_id' => $sessionId->toRfc4122()]);
+    }
 }

src/Server.php

@@ -13,23 +13,19 @@
 
 use Mcp\JsonRpc\Handler;
 use Mcp\Server\ServerBuilder;
-use Mcp\Server\Session\SessionFactoryInterface;
-use Mcp\Server\Session\SessionStoreInterface;
 use Mcp\Server\TransportInterface;
 use Psr\Log\LoggerInterface;
 use Psr\Log\NullLogger;
 use Symfony\Component\Uid\Uuid;
 
 /**
  * @author Christopher Hertel <mail@christopher-hertel.de>
+ * @author Kyrian Obikwelu <koshnawaza@gmail.com>
  */
 final class Server
 {
     public function __construct(
         private readonly Handler $jsonRpcHandler,
-        private readonly SessionFactoryInterface $sessionFactory,
-        private readonly SessionStoreInterface $sessionStore,
-        private readonly int $sessionTtl,
         private readonly LoggerInterface $logger = new NullLogger(),
     ) {}
 
@@ -47,25 +43,24 @@ public function connect(TransportInterface $transport): void
         ]);
 
         $transport->onMessage(function (string $message, ?Uuid $sessionId) use ($transport) {
-            $this->handleMessage($message, $sessionId, $transport);
-        });
-    }
+            try {
+                foreach ($this->jsonRpcHandler->process($message, $sessionId) as [$response, $context]) {
+                    if (null === $response) {
+                        continue;
+                    }
 
-    private function handleMessage(string $message, ?Uuid $sessionId, TransportInterface $transport): void
-    {
-        try {
-            foreach ($this->jsonRpcHandler->process($message, $sessionId) as $response) {
-                if (null === $response) {
-                    continue;
+                    $transport->send($response, $context);
                 }
-
-                $transport->send($response);
+            } catch (\JsonException $e) {
+                $this->logger->error('Failed to encode response to JSON.', [
+                    'message' => $message,
+                    'exception' => $e,
+                ]);
             }
-        } catch (\JsonException $e) {
-            $this->logger->error('Failed to encode response to JSON.', [
-                'message' => $message,
-                'exception' => $e,
-            ]);
-        }
+        });
+
+        $transport->onSessionEnd(function (Uuid $sessionId) {
+            $this->jsonRpcHandler->destroySession($sessionId);
+        });
     }
 }

src/Server/MethodHandlerInterface.php

@@ -16,6 +16,7 @@
 use Mcp\Schema\JsonRpc\HasMethodInterface;
 use Mcp\Schema\JsonRpc\Request;
 use Mcp\Schema\JsonRpc\Response;
+use Mcp\Server\Session\SessionInterface;
 
 /**
  * @author Christopher Hertel <mail@christopher-hertel.de>
@@ -27,5 +28,5 @@ public function supports(HasMethodInterface $message): bool;
     /**
      * @throws ExceptionInterface When the handler encounters an error processing the request
      */
-    public function handle(HasMethodInterface $message): Response|Error|null;
+    public function handle(HasMethodInterface $message, SessionInterface $session): Response|Error|null;
 }

src/Server/NotificationHandler/InitializedHandler.php

@@ -16,6 +16,7 @@
 use Mcp\Schema\JsonRpc\Response;
 use Mcp\Schema\Notification\InitializedNotification;
 use Mcp\Server\MethodHandlerInterface;
+use Mcp\Server\Session\SessionInterface;
 
 /**
  * @author Christopher Hertel <mail@christopher-hertel.de>
@@ -27,8 +28,10 @@ public function supports(HasMethodInterface $message): bool
         return $message instanceof InitializedNotification;
     }
 
-    public function handle(InitializedNotification|HasMethodInterface $message): Response|Error|null
+    public function handle(InitializedNotification|HasMethodInterface $message, SessionInterface $session): Response|Error|null
     {
+        $session->set('initialized', true);
+
         return null;
     }
 }

src/Server/RequestHandler/CallToolHandler.php

@@ -18,6 +18,7 @@
 use Mcp\Schema\JsonRpc\Response;
 use Mcp\Schema\Request\CallToolRequest;
 use Mcp\Server\MethodHandlerInterface;
+use Mcp\Server\Session\SessionInterface;
 use Psr\Log\LoggerInterface;
 use Psr\Log\NullLogger;
 
@@ -30,15 +31,14 @@ final class CallToolHandler implements MethodHandlerInterface
     public function __construct(
         private readonly ToolCallerInterface $toolCaller,
         private readonly LoggerInterface $logger = new NullLogger(),
-    ) {
-    }
+    ) {}
 
     public function supports(HasMethodInterface $message): bool
     {
         return $message instanceof CallToolRequest;
     }
 
-    public function handle(CallToolRequest|HasMethodInterface $message): Response|Error
+    public function handle(CallToolRequest|HasMethodInterface $message, SessionInterface $session): Response|Error
     {
         \assert($message instanceof CallToolRequest);