Refactor MCP transport with ping/pong handling and client tracking

Added support for server-client connection health monitoring through ping/pong mechanism integrated with `SseTransportInterface`. Updated handlers to manage message IDs, ensure client tracking, and maintain active connections. Introduced tests to validate ping intervals, connection states, and message composition.

@see https://modelcontextprotocol.io/specification/2024-11-05/basic/utilities/ping

Author: klapaudius

src/Data/Requests/NotificationData.php

@@ -10,10 +10,17 @@
  */
 class NotificationData
 {
+
+    /**
+     * The client ID of the client that sent the notification.
+     *
+     */
+    public string $clientId;
+
     /**
      * The method to be invoked.
      */
-    public string $method;
+    public ?string $method = null;
 
     /**
      * The JSON-RPC version string. MUST be "2.0".
@@ -30,11 +37,11 @@ class NotificationData
     /**
      * Constructor for NotificationData.
      *
-     * @param  string  $method  The notification method name.
+     * @param  string|null  $method  The notification method name.
      * @param  string  $jsonRpc  The JSON-RPC version (should be "2.0").
      * @param  array<mixed>|null  $params  The notification parameters.
      */
-    public function __construct(string $method, string $jsonRpc, ?array $params)
+    public function __construct(?string $method, string $jsonRpc, ?array $params)
     {
         $this->method = $method;
         $this->jsonRpc = $jsonRpc;
@@ -49,8 +56,11 @@ public function __construct(string $method, string $jsonRpc, ?array $params)
      */
     public static function fromArray(array $data): self
     {
+        if (isset($data['clientId']) && empty($data['params'])) {
+            $data['params'] = ['clientId' => $data['clientId']];
+        }
         return new self(
-            method: $data['method'],
+            method: $data['method'] ?? null,
             jsonRpc: $data['jsonrpc'],
             params: $data['params'] ?? null
         );

src/Protocol/Handlers/NotificationHandler.php

@@ -6,5 +6,5 @@ interface NotificationHandler
 {
     public function execute(?array $params = null): array;
 
-    public function isHandle(string $method): bool;
+    public function isHandle(?string $method): bool;
 }

src/Protocol/Handlers/RequestHandler.php

@@ -4,7 +4,7 @@
 
 interface RequestHandler
 {
-    public function execute(string $method, ?array $params = null): array;
+    public function execute(string $method, string|int $messageId, ?array $params = null): array;
 
     public function isHandle(string $method): bool;
 }

src/Protocol/MCPProtocol.php

@@ -12,6 +12,9 @@
 use KLP\KlpMcpServer\Exceptions\ToolParamsValidatorException;
 use KLP\KlpMcpServer\Protocol\Handlers\NotificationHandler;
 use KLP\KlpMcpServer\Protocol\Handlers\RequestHandler;
+use KLP\KlpMcpServer\Server\Notification\PongHandler;
+use KLP\KlpMcpServer\Server\Request\PingHandler;
+use KLP\KlpMcpServer\Transports\SseTransportInterface;
 use KLP\KlpMcpServer\Transports\TransportInterface;
 use KLP\KlpMcpServer\Utils\DataUtil;
 
@@ -41,6 +44,10 @@ final class MCPProtocol implements MCPProtocolInterface
     public function __construct(private readonly TransportInterface $transport)
     {
         $this->transport->onMessage([$this, 'handleMessage']);
+        if ($this->transport instanceof SseTransportInterface) {
+            $this->registerNotificationHandler(new PongHandler($this->transport->getAdapter()));
+            $this->registerRequestHandler(new PingHandler( $this->transport ));
+        }
     }
 
     /**
@@ -125,7 +132,7 @@ public function handleMessage(string $clientId, array $message): void
                 throw new JsonRpcErrorException(message: 'Invalid Request: Not a valid JSON-RPC 2.0 message', code: JsonRpcErrorCode::INVALID_REQUEST, data: $message);
             }
 
-            $requestData = DataUtil::makeRequestData(message: $message);
+            $requestData = DataUtil::makeRequestData(clientId: $clientId, message: $message);
             if ($requestData instanceof RequestData) {
                 $this->handleRequestProcess(clientId: $clientId, requestData: $requestData);
 
@@ -159,7 +166,7 @@ private function handleRequestProcess(string $clientId, RequestData $requestData
         try {
             foreach ($this->requestHandlers as $handler) {
                 if ($handler->isHandle(method: $requestData->method)) {
-                    $result = $handler->execute(method: $requestData->method, params: $requestData->params);
+                    $result = $handler->execute(method: $requestData->method, messageId: $requestData->id, params: $requestData->params);
 
                     $resultResource = new JsonRpcResultResource(id: $requestData->id, result: $result);
                     $this->pushMessage(clientId: $clientId, message: $resultResource);

src/Server/Notification/PongHandler.php

@@ -0,0 +1,43 @@
+<?php
+
+namespace KLP\KlpMcpServer\Server\Notification;
+
+use Exception;
+use KLP\KlpMcpServer\Protocol\Handlers\NotificationHandler;
+use KLP\KlpMcpServer\Transports\SseAdapters\SseAdapterInterface;
+
+
+/**
+ * Handles the processing of "Pong" notifications for managing client-server interactions.
+ * This class specifically implements the NotificationHandler interface.
+ *
+ * @see https://modelcontextprotocol.io/specification/2024-11-05/basic/utilities/ping
+ */
+readonly class PongHandler implements NotificationHandler
+{
+    public function __construct(private ?SseAdapterInterface $adapter = null)
+    {
+    }
+
+    public function isHandle(?string $method): bool
+    {
+        return is_null($method);
+    }
+
+    /**
+     * Executes a specified method with optional parameters.
+     *
+     * @param array|null $params Optional parameters to pass to the method. Null if not provided.
+     * @return array The result of the execution, returned as an array.
+     */
+    public function execute(?array $params = null): array
+    {
+        try {
+            $this->adapter?->storeLastPongResponseTimestamp($params['clientId'], time());
+        } catch (Exception) {
+            // Nothing to do here the client will be disconnected anyway
+        }
+
+        return [];
+    }
+}

src/Server/Request/InitializeHandler.php

@@ -24,7 +24,7 @@ public function isHandle(string $method): bool
     /**
      * @throws JsonRpcErrorException
      */
-    public function execute(string $method, ?array $params = null): array
+    public function execute(string $method, string|int $messageId, ?array $params = null): array
     {
         $data = InitializeData::fromArray(data: $params);
         $result = $this->server->initialize(data: $data);

src/Server/Request/PingHandler.php

@@ -0,0 +1,29 @@
+<?php
+
+namespace KLP\KlpMcpServer\Server\Request;
+
+use KLP\KlpMcpServer\Data\Requests\InitializeData;
+use KLP\KlpMcpServer\Exceptions\JsonRpcErrorException;
+use KLP\KlpMcpServer\Protocol\Handlers\RequestHandler;
+use KLP\KlpMcpServer\Server\MCPServer;
+use KLP\KlpMcpServer\Transports\SseTransportInterface;
+
+class PingHandler implements RequestHandler
+{
+
+    public function __construct(private readonly SseTransportInterface $transport)
+    {
+    }
+
+    public function isHandle(string $method): bool
+    {
+        return $method === 'ping';
+    }
+
+
+    public function execute(string $method, string|int $messageId, ?array $params = null): array
+    {
+        $this->transport->send(["id" => $messageId, "jsonrpc" => "2.0", "result" => []]);
+        return [];
+    }
+}

src/Server/Request/ToolsCallHandler.php

@@ -27,13 +27,14 @@ public function isHandle(string $method): bool
      * Executes a specified method with provided parameters and returns the result.
      *
      * @param  string  $method  The method to be executed.
+     * @param  string|int $messageId  The ID of the request message. Used for response identification.
      * @param  array|null  $params  An associative array of parameters required for execution. Must include 'name' as the tool identifier and optionally 'arguments'.
      * @return array The response array containing the execution result, which may vary based on the method.
      *
      * @throws JsonRpcErrorException If the tool name is missing or the tool is not found
      * @throws ToolParamsValidatorException If the provided arguments are invalid.
      */
-    public function execute(string $method, ?array $params = null): array
+    public function execute(string $method, string|int $messageId, ?array $params = null): array
     {
         $name = $params['name'] ?? null;
         if ($name === null) {

src/Server/Request/ToolsListHandler.php

@@ -19,7 +19,7 @@ public function isHandle(string $method): bool
         return $method === 'tools/list';
     }
 
-    public function execute(string $method, ?array $params = null): array
+    public function execute(string $method, string|int $messageId, ?array $params = null): array
     {
         return [
             'tools' => $this->toolRepository->getToolSchemas(),

src/Transports/SseAdapters/RedisAdapter.php

@@ -194,4 +194,54 @@ public function getMessageCount(string $clientId): int
             return 0;
         }
     }
+
+    /**
+     * Store the last pong response timestamp for a specific client
+     *
+     * @param  string  $clientId  The unique identifier for the client
+     * @param  int|null  $timestamp  The timestamp to store (defaults to current time if null)
+     *
+     * @throws Exception If the timestamp cannot be stored
+     */
+    public function storeLastPongResponseTimestamp(string $clientId, ?int $timestamp = null): void
+    {
+        try {
+            $key = $this->getQueueKey($clientId).":last_pong";
+            $timestamp = $timestamp ?? time();
+
+            $this->redis->set($key, $timestamp);
+            $this->redis->expire($key, $this->messageTtl);
+
+        } catch (Exception $e) {
+            $this->logger?->error('Failed to store last pong timestamp: '.$e->getMessage());
+            throw new Exception('Failed to store last pong timestamp: '.$e->getMessage());
+        }
+    }
+
+    /**
+     * Get the last pong response timestamp for a specific client
+     *
+     * @param  string  $clientId  The unique identifier for the client
+     * @return int|null The timestamp or null if no timestamp is stored
+     *
+     * @throws Exception If the timestamp cannot be retrieved
+     */
+    public function getLastPongResponseTimestamp(string $clientId): ?int
+    {
+        try {
+            $key = $this->getQueueKey($clientId).":last_pong";
+
+            $timestamp = $this->redis->get($key);
+
+            if ($timestamp === false) {
+                return null;
+            }
+
+            return (int) $timestamp;
+
+        } catch (Exception $e) {
+            $this->logger?->error('Failed to get last pong timestamp: '.$e->getMessage());
+            throw new Exception('Failed to get last pong timestamp: '.$e->getMessage());
+        }
+    }
 }