From c1bc464ebfcf1c4d67f3c22fe55ed658626336d0 Mon Sep 17 00:00:00 2001
From: Christopher Hertel <mail@christopher-hertel.de>
Date: Tue, 11 Nov 2025 22:35:26 +0100
Subject: [PATCH] Add support on Builder to inject Registry instance

---
 src/Capability/Registry.php                   | 14 +++++------
 .../Registry/ReferenceHandlerInterface.php    |  7 ++++--
 src/Capability/RegistryInterface.php          | 23 +++++++++++++++++++
 src/Server/Builder.php                        | 12 +++++++++-
 4 files changed, 45 insertions(+), 11 deletions(-)
 create mode 100644 src/Capability/RegistryInterface.php

diff --git a/src/Capability/Registry.php b/src/Capability/Registry.php
index 3a13c323..07380e4d 100644
--- a/src/Capability/Registry.php
+++ b/src/Capability/Registry.php
@@ -13,8 +13,6 @@
 
 use Mcp\Capability\Discovery\DiscoveryState;
 use Mcp\Capability\Registry\PromptReference;
-use Mcp\Capability\Registry\ReferenceProviderInterface;
-use Mcp\Capability\Registry\ReferenceRegistryInterface;
 use Mcp\Capability\Registry\ResourceReference;
 use Mcp\Capability\Registry\ResourceTemplateReference;
 use Mcp\Capability\Registry\ToolReference;
@@ -44,7 +42,7 @@
  * @author Kyrian Obikwelu <koshnawaza@gmail.com>
  * @author Pavel Buchnev   <butschster@gmail.com>
  */
-final class Registry implements ReferenceProviderInterface, ReferenceRegistryInterface
+final class Registry implements RegistryInterface
 {
     /**
      * @var array<string, ToolReference>
@@ -93,6 +91,11 @@ public function getCapabilities(): ServerCapabilities
         );
     }
 
+    public function setServerCapabilities(ServerCapabilities $serverCapabilities): void
+    {
+        $this->serverCapabilities = $serverCapabilities;
+    }
+
     public function registerTool(Tool $tool, callable|array|string $handler, bool $isManual = false): void
     {
         $toolName = $tool->name;
@@ -456,9 +459,4 @@ private function paginateResults(array $items, int $limit, ?string $cursor = nul
 
         return array_values(\array_slice($items, $offset, $limit));
     }
-
-    public function setServerCapabilities(ServerCapabilities $serverCapabilities): void
-    {
-        $this->serverCapabilities = $serverCapabilities;
-    }
 }
diff --git a/src/Capability/Registry/ReferenceHandlerInterface.php b/src/Capability/Registry/ReferenceHandlerInterface.php
index b8f565a4..c4e52f10 100644
--- a/src/Capability/Registry/ReferenceHandlerInterface.php
+++ b/src/Capability/Registry/ReferenceHandlerInterface.php
@@ -11,6 +11,9 @@
 
 namespace Mcp\Capability\Registry;
 
+use Mcp\Exception\InvalidArgumentException;
+use Mcp\Exception\RegistryException;
+
 /**
  * Interface for handling execution of MCP elements.
  * Allows custom implementations of element execution logic.
@@ -27,8 +30,8 @@ interface ReferenceHandlerInterface
      *
      * @return mixed the result of the element execution
      *
-     * @throws \Mcp\Exception\InvalidArgumentException if the handler is invalid
-     * @throws \Mcp\Exception\RegistryException        if execution fails
+     * @throws InvalidArgumentException if the handler is invalid
+     * @throws RegistryException        if execution fails
      */
     public function handle(ElementReference $reference, array $arguments): mixed;
 }
diff --git a/src/Capability/RegistryInterface.php b/src/Capability/RegistryInterface.php
new file mode 100644
index 00000000..9a5a6e0b
--- /dev/null
+++ b/src/Capability/RegistryInterface.php
@@ -0,0 +1,23 @@
+<?php
+
+/*
+ * This file is part of the official PHP MCP SDK.
+ *
+ * A collaboration between Symfony and the PHP Foundation.
+ *
+ * For the full copyright and license information, please view the LICENSE
+ * file that was distributed with this source code.
+ */
+
+namespace Mcp\Capability;
+
+use Mcp\Capability\Registry\ReferenceProviderInterface;
+use Mcp\Capability\Registry\ReferenceRegistryInterface;
+use Mcp\Schema\ServerCapabilities;
+
+interface RegistryInterface extends ReferenceRegistryInterface, ReferenceProviderInterface
+{
+    public function getCapabilities(): ServerCapabilities;
+
+    public function setServerCapabilities(ServerCapabilities $serverCapabilities): void;
+}
diff --git a/src/Server/Builder.php b/src/Server/Builder.php
index 9cc92f6a..67670fb0 100644
--- a/src/Server/Builder.php
+++ b/src/Server/Builder.php
@@ -18,6 +18,7 @@
 use Mcp\Capability\Registry\Loader\DiscoveryLoader;
 use Mcp\Capability\Registry\Loader\LoaderInterface;
 use Mcp\Capability\Registry\ReferenceHandler;
+use Mcp\Capability\RegistryInterface;
 use Mcp\JsonRpc\MessageFactory;
 use Mcp\Schema\Annotations;
 use Mcp\Schema\Enum\ProtocolVersion;
@@ -47,6 +48,8 @@ final class Builder
 {
     private ?Implementation $serverInfo = null;
 
+    private RegistryInterface $registry;
+
     private ?LoggerInterface $logger = null;
 
     private ?CacheInterface $discoveryCache = null;
@@ -245,6 +248,13 @@ public function addNotificationHandlers(iterable $handlers): self
         return $this;
     }
 
+    public function setRegistry(RegistryInterface $registry): self
+    {
+        $this->registry = $registry;
+
+        return $this;
+    }
+
     /**
      * Provides a PSR-3 logger instance. Defaults to NullLogger.
      */
@@ -401,7 +411,7 @@ public function build(): Server
     {
         $logger = $this->logger ?? new NullLogger();
         $container = $this->container ?? new Container();
-        $registry = new Registry($this->eventDispatcher, $logger);
+        $registry = $this->registry ?? new Registry($this->eventDispatcher, $logger);
 
         $loaders = [
             ...$this->loaders,