[Server] add a reference handler setter to server builder

CHANGELOG.md

@@ -7,6 +7,7 @@ All notable changes to `mcp/sdk` will be documented in this file.
 
 * Add built-in authentication middleware for HTTP transport using OAuth
 * Add client component for building MCP clients
+* Add `Builder::setReferenceHandler()` to allow custom `ReferenceHandlerInterface` implementations (e.g. authorization decorators)
 
 0.4.0
 -----

src/Server/Builder.php

@@ -22,6 +22,7 @@
 use Mcp\Capability\Registry\Loader\DiscoveryLoader;
 use Mcp\Capability\Registry\Loader\LoaderInterface;
 use Mcp\Capability\Registry\ReferenceHandler;
+use Mcp\Capability\Registry\ReferenceHandlerInterface;
 use Mcp\Capability\RegistryInterface;
 use Mcp\Exception\InvalidArgumentException;
 use Mcp\JsonRpc\MessageFactory;
@@ -69,6 +70,8 @@ final class Builder
 
     private ?SchemaGeneratorInterface $schemaGenerator = null;
 
+    private ?ReferenceHandlerInterface $referenceHandler = null;
+
     private ?DiscovererInterface $discoverer = null;
 
     private ?SessionManagerInterface $sessionManager = null;
@@ -305,6 +308,13 @@ public function setSchemaGenerator(SchemaGeneratorInterface $schemaGenerator): s
         return $this;
     }
 
+    public function setReferenceHandler(ReferenceHandlerInterface $referenceHandler): self
+    {
+        $this->referenceHandler = $referenceHandler;
+
+        return $this;
+    }
+
     public function setDiscoverer(DiscovererInterface $discoverer): self
     {
         $this->discoverer = $discoverer;
@@ -537,7 +547,7 @@ public function build(): Server
 
         $serverInfo = $this->serverInfo ?? new Implementation();
         $configuration = new Configuration($serverInfo, $capabilities, $this->paginationLimit, $this->instructions, $this->protocolVersion);
-        $referenceHandler = new ReferenceHandler($container);
+        $referenceHandler = $this->referenceHandler ?? new ReferenceHandler($container);
 
         $requestHandlers = array_merge($this->requestHandlers, [
             new Handler\Request\CallToolHandler($registry, $referenceHandler, $logger),

tests/Unit/Server/BuilderTest.php

@@ -0,0 +1,111 @@
+<?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\Tests\Unit\Server;
+
+use Mcp\Capability\Registry\ElementReference;
+use Mcp\Capability\Registry\ReferenceHandlerInterface;
+use Mcp\Schema\Content\TextContent;
+use Mcp\Schema\JsonRpc\Response;
+use Mcp\Schema\Request\CallToolRequest;
+use Mcp\Server;
+use Mcp\Server\Handler\Request\CallToolHandler;
+use Mcp\Server\Session\SessionInterface;
+use PHPUnit\Framework\Attributes\TestDox;
+use PHPUnit\Framework\TestCase;
+
+final class BuilderTest extends TestCase
+{
+    #[TestDox('setReferenceHandler() returns the builder for fluent chaining')]
+    public function testSetReferenceHandlerReturnsSelf(): void
+    {
+        $referenceHandler = $this->createStub(ReferenceHandlerInterface::class);
+
+        $builder = Server::builder();
+        $result = $builder->setReferenceHandler($referenceHandler);
+
+        $this->assertSame($builder, $result);
+    }
+
+    #[TestDox('build() succeeds with a custom ReferenceHandler')]
+    public function testBuildWithCustomReferenceHandler(): void
+    {
+        $referenceHandler = $this->createStub(ReferenceHandlerInterface::class);
+
+        $server = Server::builder()
+            ->setServerInfo('test', '1.0.0')
+            ->setReferenceHandler($referenceHandler)
+            ->build();
+
+        $this->assertInstanceOf(Server::class, $server);
+    }
+
+    #[TestDox('build() succeeds without a custom ReferenceHandler (uses default)')]
+    public function testBuildWithoutCustomReferenceHandler(): void
+    {
+        $server = Server::builder()
+            ->setServerInfo('test', '1.0.0')
+            ->build();
+
+        $this->assertInstanceOf(Server::class, $server);
+    }
+
+    #[TestDox('Custom ReferenceHandler is used when calling a tool')]
+    public function testCustomReferenceHandlerIsUsedForToolCalls(): void
+    {
+        $referenceHandler = $this->createMock(ReferenceHandlerInterface::class);
+        $referenceHandler->expects($this->once())
+            ->method('handle')
+            ->willReturnCallback(static function (ElementReference $reference, array $arguments): string {
+                return 'intercepted';
+            });
+
+        $server = Server::builder()
+            ->setServerInfo('test', '1.0.0')
+            ->setReferenceHandler($referenceHandler)
+            ->addTool(static fn (): string => 'original', 'test_tool', 'A test tool')
+            ->build();
+
+        $result = $this->callTool($server, 'test_tool');
+
+        $this->assertSame('intercepted', $result);
+    }
+
+    private function callTool(Server $server, string $toolName): mixed
+    {
+        $protocol = (new \ReflectionClass($server))->getProperty('protocol')->getValue($server);
+        $requestHandlers = (new \ReflectionClass($protocol))->getProperty('requestHandlers')->getValue($protocol);
+
+        foreach ($requestHandlers as $handler) {
+            if ($handler instanceof CallToolHandler) {
+                $request = CallToolRequest::fromArray([
+                    'jsonrpc' => '2.0',
+                    'method' => 'tools/call',
+                    'id' => 'test-1',
+                    'params' => ['name' => $toolName, 'arguments' => []],
+                ]);
+                $session = $this->createStub(SessionInterface::class);
+
+                $response = $handler->handle($request, $session);
+
+                if ($response instanceof Response) {
+                    $content = $response->result->content[0] ?? null;
+
+                    return $content instanceof TextContent ? $content->text : null;
+                }
+
+                $this->fail('Expected Response, got '.$response::class);
+            }
+        }
+
+        $this->fail('CallToolHandler not found in request handlers');
+    }
+}