From d393af634e4c90c5ac2c76760cf3f96f898a4932 Mon Sep 17 00:00:00 2001 From: soyuka Date: Wed, 1 Jul 2026 08:09:28 +0200 Subject: [PATCH] [Server] Defer element loading to first registry read MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Builder::build() runs all loaders eagerly and snapshots capabilities from the resulting registry. Both are computed once, when the server is built. Under a persistent runtime (e.g. FrankenPHP worker mode) the server is built a single time, so a loader whose data source is not yet ready at that moment (cold cache, un-warmed metadata) leaves the registry empty for the whole process — tools/list stays empty while tools/call still works. Wrap the registry in a LazyRegistry that runs the loaders on the first read, moving loading to request time when the application is initialized. The load is retried on the next read if it throws, so a transient failure at the first read does not freeze an empty registry for the whole process. Advertise capabilities from the configured element sources instead of the loaded registry, so the initialize handshake does not force an eager load. A registry supplied via setRegistry() counts as an opaque source too, so a pre-populated custom registry still advertises its capabilities. --- src/Capability/LazyRegistry.php | 206 ++++++++++++++++++ src/Server/Builder.php | 20 +- tests/Unit/Capability/LazyRegistryTest.php | 127 +++++++++++ .../Loader/ExplicitElementLoaderTest.php | 19 +- tests/Unit/Server/BuilderTest.php | 21 ++ 5 files changed, 383 insertions(+), 10 deletions(-) create mode 100644 src/Capability/LazyRegistry.php create mode 100644 tests/Unit/Capability/LazyRegistryTest.php diff --git a/src/Capability/LazyRegistry.php b/src/Capability/LazyRegistry.php new file mode 100644 index 00000000..a371f7a1 --- /dev/null +++ b/src/Capability/LazyRegistry.php @@ -0,0 +1,206 @@ + + */ +final class LazyRegistry implements RegistryInterface +{ + private bool $loaded = false; + + public function __construct( + private readonly RegistryInterface $registry, + private readonly LoaderInterface $loader, + ) { + } + + public function registerTool(Tool $tool, callable|array|string $handler): ToolReference + { + return $this->registry->registerTool($tool, $handler); + } + + public function registerResource(ResourceDefinition $resource, callable|array|string $handler): ResourceReference + { + return $this->registry->registerResource($resource, $handler); + } + + public function registerResourceTemplate(ResourceTemplate $template, callable|array|string $handler, array $completionProviders = []): ResourceTemplateReference + { + return $this->registry->registerResourceTemplate($template, $handler, $completionProviders); + } + + public function registerPrompt(Prompt $prompt, callable|array|string $handler, array $completionProviders = []): PromptReference + { + return $this->registry->registerPrompt($prompt, $handler, $completionProviders); + } + + public function unregisterTool(string $name): void + { + $this->registry->unregisterTool($name); + } + + public function unregisterResource(string $uri): void + { + $this->registry->unregisterResource($uri); + } + + public function unregisterResourceTemplate(string $uriTemplate): void + { + $this->registry->unregisterResourceTemplate($uriTemplate); + } + + public function unregisterPrompt(string $name): void + { + $this->registry->unregisterPrompt($name); + } + + public function hasTool(string $name): bool + { + $this->load(); + + return $this->registry->hasTool($name); + } + + public function hasResource(string $uri): bool + { + $this->load(); + + return $this->registry->hasResource($uri); + } + + public function hasResourceTemplate(string $uriTemplate): bool + { + $this->load(); + + return $this->registry->hasResourceTemplate($uriTemplate); + } + + public function hasPrompt(string $name): bool + { + $this->load(); + + return $this->registry->hasPrompt($name); + } + + public function hasTools(): bool + { + $this->load(); + + return $this->registry->hasTools(); + } + + public function getTools(?int $limit = null, ?string $cursor = null): Page + { + $this->load(); + + return $this->registry->getTools($limit, $cursor); + } + + public function getTool(string $name): ToolReference + { + $this->load(); + + return $this->registry->getTool($name); + } + + public function hasResources(): bool + { + $this->load(); + + return $this->registry->hasResources(); + } + + public function getResources(?int $limit = null, ?string $cursor = null): Page + { + $this->load(); + + return $this->registry->getResources($limit, $cursor); + } + + public function getResource(string $uri, bool $includeTemplates = true): ResourceReference|ResourceTemplateReference + { + $this->load(); + + return $this->registry->getResource($uri, $includeTemplates); + } + + public function hasResourceTemplates(): bool + { + $this->load(); + + return $this->registry->hasResourceTemplates(); + } + + public function getResourceTemplates(?int $limit = null, ?string $cursor = null): Page + { + $this->load(); + + return $this->registry->getResourceTemplates($limit, $cursor); + } + + public function getResourceTemplate(string $uriTemplate): ResourceTemplateReference + { + $this->load(); + + return $this->registry->getResourceTemplate($uriTemplate); + } + + public function hasPrompts(): bool + { + $this->load(); + + return $this->registry->hasPrompts(); + } + + public function getPrompts(?int $limit = null, ?string $cursor = null): Page + { + $this->load(); + + return $this->registry->getPrompts($limit, $cursor); + } + + public function getPrompt(string $name): PromptReference + { + $this->load(); + + return $this->registry->getPrompt($name); + } + + private function load(): void + { + if ($this->loaded) { + return; + } + + $this->loader->load($this->registry); + $this->loaded = true; + } +} diff --git a/src/Server/Builder.php b/src/Server/Builder.php index 198a3497..4351c219 100644 --- a/src/Server/Builder.php +++ b/src/Server/Builder.php @@ -16,6 +16,7 @@ use Mcp\Capability\Discovery\Discoverer; use Mcp\Capability\Discovery\DiscovererInterface; use Mcp\Capability\Discovery\SchemaGeneratorInterface; +use Mcp\Capability\LazyRegistry; use Mcp\Capability\Registry; use Mcp\Capability\Registry\Container; use Mcp\Capability\Registry\ElementReference; @@ -219,6 +220,8 @@ final class Builder */ private array $loaders = []; + private bool $hasCustomRegistry = false; + /** * Sets the server's identity. Required. * @@ -344,6 +347,7 @@ public function addNotificationHandlers(iterable $handlers): self public function setRegistry(RegistryInterface $registry): self { $this->registry = $registry; + $this->hasCustomRegistry = true; return $this; } @@ -674,18 +678,22 @@ public function build(): Server } } - $loader = new ChainLoader($loaders); - $loader->load($registry); + // Defer loading to the first registry read (request time) instead of eagerly here. @see LazyRegistry + $registry = new LazyRegistry($registry, new ChainLoader($loaders)); $messageFactory = MessageFactory::make(); + // Advertise from configured sources, not the now-lazy registry, so initialize does not force a load. + // Opaque sources (custom loaders, discovery, a caller registry) advertise all kinds; over-advertising is harmless. + $hasOpaqueSources = [] !== $this->loaders || null !== $this->discoveryBasePath || $this->hasCustomRegistry; + $capabilities = $this->serverCapabilities ?? new ServerCapabilities( - tools: $registry->hasTools(), + tools: [] !== $this->tools || [] !== $this->explicitTools || $hasOpaqueSources, toolsListChanged: $this->eventDispatcher instanceof EventDispatcherInterface, - resources: $registry->hasResources() || $registry->hasResourceTemplates(), - resourcesSubscribe: $registry->hasResources() || $registry->hasResourceTemplates(), + resources: [] !== $this->resources || [] !== $this->explicitResources || [] !== $this->resourceTemplates || [] !== $this->explicitResourceTemplates || $hasOpaqueSources, + resourcesSubscribe: [] !== $this->resources || [] !== $this->explicitResources || [] !== $this->resourceTemplates || [] !== $this->explicitResourceTemplates || $hasOpaqueSources, resourcesListChanged: $this->eventDispatcher instanceof EventDispatcherInterface, - prompts: $registry->hasPrompts(), + prompts: [] !== $this->prompts || [] !== $this->explicitPrompts || $hasOpaqueSources, promptsListChanged: $this->eventDispatcher instanceof EventDispatcherInterface, logging: true, completions: true, diff --git a/tests/Unit/Capability/LazyRegistryTest.php b/tests/Unit/Capability/LazyRegistryTest.php new file mode 100644 index 00000000..c98cf9ba --- /dev/null +++ b/tests/Unit/Capability/LazyRegistryTest.php @@ -0,0 +1,127 @@ +createMock(LoaderInterface::class); + $loader->expects($this->never())->method('load'); + + // Constructing (and registering) must not trigger the loader. + $registry = new LazyRegistry(new Registry(), $loader); + $registry->registerTool($this->tool('manual'), 'handler'); + } + + public function testLoaderRunsOnFirstReadAndPopulatesTheRegistry(): void + { + $inner = new Registry(); + $loader = new class($this->tool('loaded')) implements LoaderInterface { + public function __construct(private readonly Tool $tool) + { + } + + public function load(RegistryInterface $registry): void + { + $registry->registerTool($this->tool, 'handler'); + } + }; + + $registry = new LazyRegistry($inner, $loader); + + $this->assertTrue($registry->hasTools()); + $tools = $registry->getTools()->references; + $this->assertArrayHasKey('loaded', $tools); + } + + public function testLoaderRunsExactlyOnceAcrossManyReads(): void + { + $loader = $this->createMock(LoaderInterface::class); + $loader->expects($this->once())->method('load'); + + $registry = new LazyRegistry(new Registry(), $loader); + $registry->hasTools(); + $registry->getTools(); + $registry->hasResources(); + $registry->getPrompts(); + } + + public function testRuntimeRegistrationsSurviveTheDeferredLoad(): void + { + $inner = new Registry(); + $loader = new class($this->tool('loaded')) implements LoaderInterface { + public function __construct(private readonly Tool $tool) + { + } + + public function load(RegistryInterface $registry): void + { + $registry->registerTool($this->tool, 'handler'); + } + }; + + $registry = new LazyRegistry($inner, $loader); + // Registered before the first read; the deferred load must be additive, not replacing. + $registry->registerTool($this->tool('runtime'), 'handler'); + + $tools = $registry->getTools()->references; + $this->assertArrayHasKey('runtime', $tools); + $this->assertArrayHasKey('loaded', $tools); + } + + public function testLoaderRetriesAfterAFailedLoad(): void + { + $inner = new Registry(); + $loader = new class($this->tool('loaded')) implements LoaderInterface { + private int $calls = 0; + + public function __construct(private readonly Tool $tool) + { + } + + public function load(RegistryInterface $registry): void + { + ++$this->calls; + if (1 === $this->calls) { + throw new \RuntimeException('data source not ready'); + } + + $registry->registerTool($this->tool, 'handler'); + } + }; + + $registry = new LazyRegistry($inner, $loader); + + try { + $registry->hasTools(); + $this->fail('Expected the first load to throw.'); + } catch (\RuntimeException $e) { + $this->assertSame('data source not ready', $e->getMessage()); + } + + $tools = $registry->getTools()->references; + $this->assertArrayHasKey('loaded', $tools); + } + + private function tool(string $name): Tool + { + return new Tool($name, null, ['type' => 'object', 'properties' => [], 'required' => null], null, null); + } +} diff --git a/tests/Unit/Capability/Registry/Loader/ExplicitElementLoaderTest.php b/tests/Unit/Capability/Registry/Loader/ExplicitElementLoaderTest.php index 0066a318..29a895e8 100644 --- a/tests/Unit/Capability/Registry/Loader/ExplicitElementLoaderTest.php +++ b/tests/Unit/Capability/Registry/Loader/ExplicitElementLoaderTest.php @@ -23,6 +23,7 @@ use Mcp\Server; use Mcp\Server\ClientGateway; use Mcp\Server\Handler\PromptHandlerInterface; +use Mcp\Server\Handler\Request\CallToolHandler; use Mcp\Server\Handler\ResourceHandlerInterface; use Mcp\Server\Handler\ResourceTemplateHandlerInterface; use Mcp\Server\Handler\ToolHandlerInterface; @@ -373,12 +374,22 @@ public function execute(array $arguments, ClientGateway $gateway): mixed */ private function buildAndGetRegistry(callable $configure): RegistryInterface { - $registry = new Registry(); $builder = Server::builder() ->setServerInfo('test', '1.0.0') - ->setRegistry($registry); - $configure($builder)->build(); + ->setRegistry(new Registry()); + $server = $configure($builder)->build(); - return $registry; + // build() wraps the registry in a LazyRegistry that loads on first read; read through the + // server's registry rather than the injected instance so the deferred load runs. + $protocol = (new \ReflectionClass($server))->getProperty('protocol')->getValue($server); + $requestHandlers = (new \ReflectionClass($protocol))->getProperty('requestHandlers')->getValue($protocol); + + foreach ($requestHandlers as $handler) { + if ($handler instanceof CallToolHandler) { + return (new \ReflectionClass($handler))->getProperty('registry')->getValue($handler); + } + } + + $this->fail('CallToolHandler not found in request handlers'); } } diff --git a/tests/Unit/Server/BuilderTest.php b/tests/Unit/Server/BuilderTest.php index 4435639c..d917383a 100644 --- a/tests/Unit/Server/BuilderTest.php +++ b/tests/Unit/Server/BuilderTest.php @@ -11,6 +11,7 @@ namespace Mcp\Tests\Unit\Server; +use Mcp\Capability\Registry; use Mcp\Capability\Registry\ElementReference; use Mcp\Capability\Registry\ReferenceHandlerInterface; use Mcp\Exception\LogicException; @@ -19,6 +20,7 @@ use Mcp\Schema\JsonRpc\Response; use Mcp\Schema\Request\CallToolRequest; use Mcp\Schema\ServerCapabilities; +use Mcp\Schema\Tool; use Mcp\Server; use Mcp\Server\Handler\Request\CallToolHandler; use Mcp\Server\Handler\Request\InitializeHandler; @@ -122,6 +124,25 @@ public function testEnableExtensionMergesIntoCustomCapabilities(): void $this->assertArrayHasKey(McpApps::EXTENSION_ID, $capabilities->extensions); } + #[TestDox('build() advertises tools capability for a pre-populated registry set via setRegistry()')] + public function testBuildAdvertisesToolsForPreloadedCustomRegistry(): void + { + $registry = new Registry(); + $registry->registerTool( + new Tool(name: 'test_tool', title: null, inputSchema: ['type' => 'object', 'properties' => [], 'required' => null], description: 'A test tool', annotations: null), + static fn (): string => 'result', + ); + + $server = Server::builder() + ->setServerInfo('test', '1.0.0') + ->setRegistry($registry) + ->build(); + + $capabilities = $this->extractServerCapabilities($server); + + $this->assertTrue($capabilities->tools); + } + private function extractServerCapabilities(Server $server): ServerCapabilities { $protocol = (new \ReflectionClass($server))->getProperty('protocol')->getValue($server);