Skip to content

Commit 2bc2bdb

Browse files
soyukasoyuka
committed
[Server][Capability] Defer loading into a custom registry
`setLazyLoading()` defaults to true but never applied to a registry supplied through `setRegistry()`: the builder ran the chain loader itself and set `$eagerlyLoaded = true`, because it could not hand a loader to an instance it did not construct. Two consequences under a persistent runtime. The loaders run once at build, so a source that is not ready yet (a cold metadata cache) freezes the registry empty for the whole process and `tools/list` keeps returning `[]`. And `detectCapabilities()` reads `hasTools()` off that cold registry, so the server can advertise `tools: false` and a client that respects capabilities never calls `tools/list` at all. `Registry::deferLoadingFrom()` adopts a loader after construction, so the builder can defer instead of loading eagerly. It chains behind a loader the constructor took while that one is still owed its run, and replaces it once it has already run, since chaining would run it twice and discovery would rescan. Resetting `loaded` is what lets the adopted loader run on the next read: without it, a registry read before `build()` would keep the deferred loader forever unrun. Deferral is not conditional on the registry being empty. Gating it that way would mean a caller who hand-registers a single element before `setRegistry()` silently falls back to eager loading and gets the cold-source bug back. Instead `Registry::isEmpty()` — which reads the element arrays directly and so does not trigger the loader — tells `detectCapabilities()` that a deferred registry already holds elements, and it advertises them as one more opaque source. That over-advertises a registry holding only one kind, which is harmless per MCP semantics and already how custom loaders and discovery are treated. A foreign `RegistryInterface` cannot be deferred into and still loads eagerly. `RegistryInterface` is unchanged: both methods live on `Registry`, which is what the builder already type-checks for `loadFrom()`.
1 parent c5dbfb6 commit 2bc2bdb

4 files changed

Lines changed: 233 additions & 15 deletions

File tree

‎src/Capability/Registry.php‎

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111

1212
namespace Mcp\Capability;
1313

14+
use Mcp\Capability\Registry\Loader\ChainLoader;
1415
use Mcp\Capability\Registry\Loader\LoaderInterface;
1516
use Mcp\Capability\Registry\PromptReference;
1617
use Mcp\Capability\Registry\ResourceReference;
@@ -69,7 +70,7 @@ public function __construct(
6970
private readonly ?EventDispatcherInterface $eventDispatcher = null,
7071
private readonly LoggerInterface $logger = new NullLogger(),
7172
private readonly NameValidator $nameValidator = new NameValidator(),
72-
private readonly ?LoaderInterface $loader = null,
73+
private ?LoaderInterface $loader = null,
7374
) {
7475
}
7576

@@ -114,6 +115,28 @@ public function loadFrom(LoaderInterface $loader): void
114115
}
115116
}
116117

118+
/**
119+
* Adopts $loader for the deferred load, so a registry the caller constructed can still load at
120+
* first read instead of at build time. Chains behind a loader the constructor already took,
121+
* but only when that loader is still owed its run: once it has already run, chaining would run
122+
* it a second time — discovery would rescan — so $loader replaces it instead. Resetting $loaded
123+
* is what makes the adopted loader actually run on the next read.
124+
*/
125+
public function deferLoadingFrom(LoaderInterface $loader): void
126+
{
127+
$this->loader = null === $this->loader || $this->loaded ? $loader : new ChainLoader([$this->loader, $loader]);
128+
$this->loaded = false;
129+
}
130+
131+
/**
132+
* True when nothing is registered yet. Reads the backing arrays directly, so unlike has*() it
133+
* never triggers the loader.
134+
*/
135+
public function isEmpty(): bool
136+
{
137+
return [] === $this->tools && [] === $this->resources && [] === $this->resourceTemplates && [] === $this->prompts;
138+
}
139+
117140
public function registerTool(Tool $tool, callable|array|string $handler): ToolReference
118141
{
119142
if (!$this->nameValidator->isValid($tool->name)) {

‎src/Server/Builder.php‎

Lines changed: 23 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -522,8 +522,9 @@ public function setRegistry(RegistryInterface $registry): self
522522
*
523523
* Lazy (the default) defers loading to the first registry read so a persistent runtime does not
524524
* freeze the registry to a source not yet ready at build time. Disable to load eagerly at build.
525-
* A registry supplied via setRegistry() is always loaded eagerly; its own constructor loader,
526-
* if it has one, still runs on the first read.
525+
* A registry supplied via setRegistry() is deferred the same way, whatever it already holds is
526+
* still advertised by capability detection. Either way its own constructor loader, if it has
527+
* one, still runs on the first read.
527528
*/
528529
public function setLazyLoading(bool $lazyLoading = true): self
529530
{
@@ -1043,17 +1044,25 @@ private function resolve(): array
10431044
}
10441045

10451046
$chainLoader = new ChainLoader($loaders);
1047+
$hasPreloadedElements = false;
10461048

10471049
if ($this->hasCustomRegistry) {
1048-
// Builder can't inject the loader into an already-constructed instance, so load it eagerly.
1049-
// Via loadFrom(), which suppresses the change events the load would otherwise dispatch.
10501050
$registry = $this->registry;
1051-
if ($registry instanceof Registry) {
1052-
$registry->loadFrom($chainLoader);
1051+
1052+
if ($this->lazyLoading && $registry instanceof Registry) {
1053+
$hasPreloadedElements = !$registry->isEmpty();
1054+
$registry->deferLoadingFrom($chainLoader);
1055+
$eagerlyLoaded = false;
10531056
} else {
1054-
$chainLoader->load($registry);
1057+
// A foreign RegistryInterface cannot be deferred into, so it is loaded eagerly here.
1058+
// loadFrom() suppresses the change events the load would otherwise dispatch.
1059+
if ($registry instanceof Registry) {
1060+
$registry->loadFrom($chainLoader);
1061+
} else {
1062+
$chainLoader->load($registry);
1063+
}
1064+
$eagerlyLoaded = true;
10551065
}
1056-
$eagerlyLoaded = true;
10571066
} else {
10581067
$registry = new Registry($eventDispatcher, $logger, loader: $chainLoader);
10591068
if (!$this->lazyLoading) {
@@ -1064,7 +1073,7 @@ private function resolve(): array
10641073

10651074
$messageFactory = MessageFactory::make(additional: $this->extensionMessages);
10661075

1067-
$capabilities = $this->serverCapabilities ?? $this->detectCapabilities($registry, $eagerlyLoaded, $eventDispatcher);
1076+
$capabilities = $this->serverCapabilities ?? $this->detectCapabilities($registry, $eagerlyLoaded, $eventDispatcher, $hasPreloadedElements);
10681077

10691078
// Extensions enabled via enableExtension() are folded into caller-supplied
10701079
// capabilities too, so setCapabilities() does not silently drop them.
@@ -1118,9 +1127,11 @@ private function resolve(): array
11181127
/**
11191128
* When loaded, capabilities are read from the registry. When deferred, reading it would force
11201129
* the load, so they are advertised from the configured sources instead — opaque sources (custom
1121-
* loaders, discovery) advertise all kinds, and over-advertising is harmless per MCP semantics.
1130+
* loaders, discovery) advertise all kinds, and over-advertising is harmless per MCP semantics. A
1131+
* custom registry deferred while already holding elements ($hasPreloadedElements) counts as an
1132+
* opaque source too, for the same reason: reading it would force the load it is deferred to avoid.
11221133
*/
1123-
private function detectCapabilities(RegistryInterface $registry, bool $eagerlyLoaded, ?EventDispatcherInterface $eventDispatcher): ServerCapabilities
1134+
private function detectCapabilities(RegistryInterface $registry, bool $eagerlyLoaded, ?EventDispatcherInterface $eventDispatcher, bool $hasPreloadedElements): ServerCapabilities
11241135
{
11251136
// Without a dispatcher the registry announces nothing, so there is no
11261137
// list-changed notification to advertise.
@@ -1143,7 +1154,7 @@ private function detectCapabilities(RegistryInterface $registry, bool $eagerlyLo
11431154
);
11441155
}
11451156

1146-
$hasOpaqueSources = [] !== $this->loaders || null !== $this->discoveryBasePath;
1157+
$hasOpaqueSources = [] !== $this->loaders || null !== $this->discoveryBasePath || $hasPreloadedElements;
11471158
$hasResources = [] !== $this->resources || [] !== $this->explicitResources || [] !== $this->resourceTemplates || [] !== $this->explicitResourceTemplates || $hasOpaqueSources;
11481159

11491160
return new ServerCapabilities(

‎tests/Unit/Capability/RegistryTest.php‎

Lines changed: 100 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -853,6 +853,104 @@ public function testLoadIsANoopWithoutAConfiguredLoader(): void
853853
$this->assertFalse($registry->hasTools());
854854
}
855855

856+
public function testDeferLoadingFromDoesNotRunUntilFirstRead(): void
857+
{
858+
$loader = $this->createMock(LoaderInterface::class);
859+
$loader->expects($this->never())->method('load');
860+
861+
$registry = new Registry(null, $this->logger);
862+
$registry->deferLoadingFrom($loader);
863+
}
864+
865+
public function testDeferLoadingFromRunsOnFirstReadAndPopulatesTheRegistry(): void
866+
{
867+
$registry = new Registry(null, $this->logger);
868+
$registry->deferLoadingFrom($this->toolLoader($this->createValidTool('deferred')));
869+
870+
$this->assertTrue($registry->hasTools());
871+
$this->assertArrayHasKey('deferred', $registry->getTools()->references);
872+
}
873+
874+
public function testDeferLoadingFromRunsTheLoaderExactlyOnceAcrossManyReads(): void
875+
{
876+
$loader = $this->createMock(LoaderInterface::class);
877+
$loader->expects($this->once())->method('load');
878+
879+
$registry = new Registry(null, $this->logger);
880+
$registry->deferLoadingFrom($loader);
881+
882+
$registry->hasTools();
883+
$registry->getTools();
884+
$registry->hasResources();
885+
}
886+
887+
public function testDeferLoadingFromChainsBehindTheConstructorLoader(): void
888+
{
889+
// Both register 'shared'; last-write-wins proves the run order, since
890+
// ChainLoader lets the later loader overwrite the earlier one's registration.
891+
$constructorLoader = $this->toolLoader($this->createValidTool('shared', null, 'from constructor'));
892+
$deferredLoader = $this->toolLoader($this->createValidTool('shared', null, 'from deferred'));
893+
894+
$registry = new Registry(null, $this->logger, loader: $constructorLoader);
895+
$registry->deferLoadingFrom($deferredLoader);
896+
897+
$tools = $registry->getTools()->references;
898+
899+
$this->assertArrayHasKey('shared', $tools);
900+
$this->assertSame('from deferred', $tools['shared']->description);
901+
}
902+
903+
public function testDeferLoadingFromRunsTheAdoptedLoaderAfterTheConstructorLoaderAlreadyRan(): void
904+
{
905+
// A read before deferLoadingFrom() runs the constructor loader and sets $loaded, the bug
906+
// this covers: the adopted loader was then stored but never run because load() returned on
907+
// $loaded before consulting it.
908+
$constructorLoader = new class implements LoaderInterface {
909+
public int $calls = 0;
910+
911+
public function load(RegistryInterface $registry): void
912+
{
913+
++$this->calls;
914+
}
915+
};
916+
$adoptedLoader = new class implements LoaderInterface {
917+
public int $calls = 0;
918+
919+
public function load(RegistryInterface $registry): void
920+
{
921+
++$this->calls;
922+
}
923+
};
924+
925+
$registry = new Registry(null, $this->logger, loader: $constructorLoader);
926+
$registry->hasTools();
927+
928+
$registry->deferLoadingFrom($adoptedLoader);
929+
$registry->hasTools();
930+
$registry->hasResources();
931+
932+
$this->assertSame(1, $constructorLoader->calls);
933+
$this->assertSame(1, $adoptedLoader->calls);
934+
}
935+
936+
public function testIsEmptyIsTrueForAFreshRegistryAndDoesNotTriggerTheLoader(): void
937+
{
938+
$loader = $this->createMock(LoaderInterface::class);
939+
$loader->expects($this->never())->method('load');
940+
941+
$registry = new Registry(null, $this->logger, loader: $loader);
942+
943+
$this->assertTrue($registry->isEmpty());
944+
}
945+
946+
public function testIsEmptyIsFalseAfterRegisterTool(): void
947+
{
948+
$registry = new Registry(null, $this->logger);
949+
$registry->registerTool($this->createValidTool('registered'), 'handler');
950+
951+
$this->assertFalse($registry->isEmpty());
952+
}
953+
856954
private function toolLoader(Tool $tool): LoaderInterface
857955
{
858956
return new class($tool) implements LoaderInterface {
@@ -907,7 +1005,7 @@ public function jsonSerialize(): float
9071005
$this->assertNull($toolRef->extractStructuredContent($result, ProtocolVersion::V2025_11_25));
9081006
}
9091007

910-
private function createValidTool(string $name, ?array $outputSchema = null): Tool
1008+
private function createValidTool(string $name, ?array $outputSchema = null, ?string $description = null): Tool
9111009
{
9121010
return new Tool(
9131011
name: $name,
@@ -919,7 +1017,7 @@ private function createValidTool(string $name, ?array $outputSchema = null): Too
9191017
],
9201018
'required' => null,
9211019
],
922-
description: "Test tool: {$name}",
1020+
description: $description ?? "Test tool: {$name}",
9231021
annotations: null,
9241022
icons: null,
9251023
meta: null,

‎tests/Unit/Server/BuilderTest.php‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -443,6 +443,92 @@ public function testAThirdPartyRegistryIsStillLoadedThroughThePlainLoader(): voi
443443
->addTool(static fn (): string => 'ok', 'alpha')
444444
->build();
445445
}
446+
447+
#[TestDox('An empty custom registry with lazy loading defers its configured loader past build(), running it on the first registry read')]
448+
public function testEmptyCustomRegistryDefersLoaderPastBuild(): void
449+
{
450+
$loader = new class implements LoaderInterface {
451+
public int $calls = 0;
452+
453+
public function load(RegistryInterface $registry): void
454+
{
455+
++$this->calls;
456+
}
457+
};
458+
459+
$registry = new Registry();
460+
461+
Server::builder()
462+
->setRegistry($registry)
463+
->addLoader($loader)
464+
->build();
465+
466+
$this->assertSame(0, $loader->calls);
467+
468+
$registry->hasTools();
469+
470+
$this->assertSame(1, $loader->calls);
471+
}
472+
473+
#[TestDox('An empty custom registry with lazy loading advertises tools from the configured loader without forcing a load')]
474+
public function testEmptyCustomRegistryAdvertisesToolsFromConfiguredLoaderWithoutLoading(): void
475+
{
476+
$loader = $this->createMock(LoaderInterface::class);
477+
$loader->expects($this->never())->method('load');
478+
479+
$registry = new Registry();
480+
481+
$server = Server::builder()
482+
->setServerInfo('test', '1.0.0')
483+
->setRegistry($registry)
484+
->addLoader($loader)
485+
->build();
486+
487+
$capabilities = $this->extractServerCapabilities($server);
488+
489+
$this->assertTrue($capabilities->tools);
490+
}
491+
492+
#[TestDox('setLazyLoading(false) with an empty custom registry loads it eagerly during build()')]
493+
public function testSetLazyLoadingFalseWithEmptyCustomRegistryLoadsEagerly(): void
494+
{
495+
$loader = $this->createMock(LoaderInterface::class);
496+
$loader->expects($this->once())->method('load');
497+
498+
$registry = new Registry();
499+
500+
Server::builder()
501+
->setRegistry($registry)
502+
->setLazyLoading(false)
503+
->addLoader($loader)
504+
->build();
505+
}
506+
507+
#[TestDox('A pre-populated custom registry with lazy loading also defers its configured loader past build(), instead of the old isEmpty() gate loading it eagerly')]
508+
public function testPreloadedCustomRegistryDefersLoaderPastBuild(): void
509+
{
510+
$loader = new class implements LoaderInterface {
511+
public int $calls = 0;
512+
513+
public function load(RegistryInterface $registry): void
514+
{
515+
++$this->calls;
516+
}
517+
};
518+
519+
$registry = new Registry();
520+
$registry->registerTool(
521+
new Tool(name: 'preloaded_tool', title: null, inputSchema: ['type' => 'object', 'properties' => [], 'required' => null], description: 'A preloaded tool', annotations: null),
522+
static fn (): string => 'result',
523+
);
524+
525+
Server::builder()
526+
->setRegistry($registry)
527+
->addLoader($loader)
528+
->build();
529+
530+
$this->assertSame(0, $loader->calls);
531+
}
446532
}
447533

448534
/**

0 commit comments

Comments
 (0)