[Server] Fix: ensure JSON-RPC error responses include message IDs (#77)

* fix: ensure JSON-RPC error responses include message IDs

* fix: replace method_exists with instanceof for message ID extraction

Author: CodeWithKyrian

src/JsonRpc/Handler.php

@@ -23,6 +23,7 @@
 use Mcp\Schema\Implementation;
 use Mcp\Schema\JsonRpc\Error;
 use Mcp\Schema\JsonRpc\HasMethodInterface;
+use Mcp\Schema\JsonRpc\Request;
 use Mcp\Schema\JsonRpc\Response;
 use Mcp\Schema\Request\InitializeRequest;
 use Mcp\Server\MethodHandlerInterface;
@@ -94,9 +95,6 @@ public static function make(
 
     /**
      * @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
      */
     public function process(string $input, ?Uuid $sessionId): iterable
     {
@@ -163,7 +161,7 @@ public function process(string $input, ?Uuid $sessionId): iterable
         foreach ($messages as $message) {
             if ($message instanceof InvalidInputMessageException) {
                 $this->logger->warning('Failed to create message.', ['exception' => $message]);
-                $error = Error::forInvalidRequest($message->getMessage(), 0);
+                $error = Error::forInvalidRequest($message->getMessage());
                 yield [$this->encodeResponse($error), []];
                 continue;
             }
@@ -172,6 +170,8 @@ public function process(string $input, ?Uuid $sessionId): iterable
                 'method' => $message->getMethod(),
             ]);
 
+            $messageId = $message instanceof Request ? $message->getId() : 0;
+
             try {
                 $response = $this->handle($message, $session);
                 yield [$this->encodeResponse($response), ['session_id' => $session->getId()]];
@@ -183,17 +183,17 @@ public function process(string $input, ?Uuid $sessionId): iterable
                     ['exception' => $e],
                 );
 
-                $error = Error::forMethodNotFound($e->getMessage());
+                $error = Error::forMethodNotFound($e->getMessage(), $messageId);
                 yield [$this->encodeResponse($error), []];
             } catch (\InvalidArgumentException $e) {
                 $this->logger->warning(\sprintf('Invalid argument: %s', $e->getMessage()), ['exception' => $e]);
 
-                $error = Error::forInvalidParams($e->getMessage());
+                $error = Error::forInvalidParams($e->getMessage(), $messageId);
                 yield [$this->encodeResponse($error), []];
             } catch (\Throwable $e) {
                 $this->logger->critical(\sprintf('Uncaught exception: %s', $e->getMessage()), ['exception' => $e]);
 
-                $error = Error::forInternalError($e->getMessage());
+                $error = Error::forInternalError($e->getMessage(), $messageId);
                 yield [$this->encodeResponse($error), []];
             }
         }
@@ -202,7 +202,7 @@ public function process(string $input, ?Uuid $sessionId): iterable
     }
 
     /**
-     * @throws \JsonException When JSON encoding fails
+     * Encodes a response to JSON, handling encoding errors gracefully.
      */
     private function encodeResponse(Response|Error|null $response): ?string
     {
@@ -214,11 +214,26 @@ private function encodeResponse(Response|Error|null $response): ?string
 
         $this->logger->info('Encoding response.', ['response' => $response]);
 
-        if ($response instanceof Response && [] === $response->result) {
-            return json_encode($response, \JSON_THROW_ON_ERROR | \JSON_FORCE_OBJECT);
-        }
+        try {
+            if ($response instanceof Response && [] === $response->result) {
+                return json_encode($response, \JSON_THROW_ON_ERROR | \JSON_FORCE_OBJECT);
+            }
 
-        return json_encode($response, \JSON_THROW_ON_ERROR);
+            return json_encode($response, \JSON_THROW_ON_ERROR);
+        } catch (\JsonException $e) {
+            $this->logger->error('Failed to encode response to JSON.', [
+                'message_id' => $response->getId(),
+                'exception' => $e,
+            ]);
+
+            $fallbackError = new Error(
+                id: $response->getId(),
+                code: Error::INTERNAL_ERROR,
+                message: 'Response could not be encoded to JSON'
+            );
+
+            return json_encode($fallbackError, \JSON_THROW_ON_ERROR);
+        }
     }
 
     /**

src/Server.php

@@ -44,19 +44,12 @@ public function connect(TransportInterface $transport): void
         ]);
 
         $transport->onMessage(function (string $message, ?Uuid $sessionId) use ($transport) {
-            try {
-                foreach ($this->jsonRpcHandler->process($message, $sessionId) as [$response, $context]) {
-                    if (null === $response) {
-                        continue;
-                    }
-
-                    $transport->send($response, $context);
+            foreach ($this->jsonRpcHandler->process($message, $sessionId) as [$response, $context]) {
+                if (null === $response) {
+                    continue;
                 }
-            } catch (\JsonException $e) {
-                $this->logger->error('Failed to encode response to JSON.', [
-                    'message' => $message,
-                    'exception' => $e,
-                ]);
+
+                $transport->send($response, $context);
             }
         });
 

tests/ServerTest.php

@@ -14,37 +14,32 @@
 use Mcp\JsonRpc\Handler;
 use Mcp\Server;
 use Mcp\Server\Transport\InMemoryTransport;
-use PHPUnit\Framework\MockObject\Stub\Exception;
 use PHPUnit\Framework\TestCase;
-use Psr\Log\NullLogger;
 
 class ServerTest extends TestCase
 {
     public function testJsonExceptions()
     {
-        $logger = $this->getMockBuilder(NullLogger::class)
-            ->disableOriginalConstructor()
-            ->onlyMethods(['error'])
-            ->getMock();
-        $logger->expects($this->once())->method('error');
-
         $handler = $this->getMockBuilder(Handler::class)
             ->disableOriginalConstructor()
             ->onlyMethods(['process'])
             ->getMock();
 
         $handler->expects($this->exactly(2))->method('process')->willReturnOnConsecutiveCalls(
-            new Exception(new \JsonException('foobar')),
+            [['{"jsonrpc":"2.0","id":0,"error":{"code":-32700,"message":"Parse error"}}', []]],
             [['success', []]]
         );
 
         $transport = $this->getMockBuilder(InMemoryTransport::class)
             ->setConstructorArgs([['foo', 'bar']])
             ->onlyMethods(['send'])
             ->getMock();
-        $transport->expects($this->once())->method('send')->with('success', []);
+        $transport->expects($this->exactly(2))->method('send')->willReturnOnConsecutiveCalls(
+            null,
+            null
+        );
 
-        $server = new Server($handler, $logger);
+        $server = new Server($handler);
         $server->connect($transport);
 
         $transport->listen();