Skip to content

Commit c8c59a4

Browse files
soyukasoyuka
committed
fix: clear() wiping dynamically registered tools by tracking discovered state on references
clear() previously removed all non-manual elements, which destroyed dynamically registered tools when setDiscoveryState() was called on the next request cycle. Add isDiscovered flag to ElementReference so clear() only removes elements that were imported via setDiscoveryState(), preserving both manual and dynamic registrations.
1 parent 6a99196 commit c8c59a4

9 files changed

Lines changed: 202 additions & 65 deletions

‎src/Capability/Registry.php‎

Lines changed: 21 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -161,36 +161,10 @@ public function registerPrompt(
161161

162162
public function clear(): void
163163
{
164-
$clearCount = 0;
165-
166-
foreach ($this->tools as $name => $tool) {
167-
if (!$tool->isManual) {
168-
unset($this->tools[$name]);
169-
++$clearCount;
170-
}
171-
}
172-
foreach ($this->resources as $uri => $resource) {
173-
if (!$resource->isManual) {
174-
unset($this->resources[$uri]);
175-
++$clearCount;
176-
}
177-
}
178-
foreach ($this->prompts as $name => $prompt) {
179-
if (!$prompt->isManual) {
180-
unset($this->prompts[$name]);
181-
++$clearCount;
182-
}
183-
}
184-
foreach ($this->resourceTemplates as $uriTemplate => $template) {
185-
if (!$template->isManual) {
186-
unset($this->resourceTemplates[$uriTemplate]);
187-
++$clearCount;
188-
}
189-
}
190-
191-
if ($clearCount > 0) {
192-
$this->logger->debug(\sprintf('Removed %d discovered elements from internal registry.', $clearCount));
193-
}
164+
$this->tools = array_filter($this->tools, static fn ($t) => !$t->isDiscovered);
165+
$this->resources = array_filter($this->resources, static fn ($r) => !$r->isDiscovered);
166+
$this->prompts = array_filter($this->prompts, static fn ($p) => !$p->isDiscovered);
167+
$this->resourceTemplates = array_filter($this->resourceTemplates, static fn ($t) => !$t->isDiscovered);
194168
}
195169

196170
public function hasTools(): bool
@@ -344,10 +318,10 @@ public function getPrompt(string $name): PromptReference
344318
public function getDiscoveryState(): DiscoveryState
345319
{
346320
return new DiscoveryState(
347-
tools: array_filter($this->tools, static fn ($tool) => !$tool->isManual),
348-
resources: array_filter($this->resources, static fn ($resource) => !$resource->isManual),
349-
prompts: array_filter($this->prompts, static fn ($prompt) => !$prompt->isManual),
350-
resourceTemplates: array_filter($this->resourceTemplates, static fn ($template) => !$template->isManual),
321+
tools: array_filter($this->tools, static fn ($t) => $t->isDiscovered),
322+
resources: array_filter($this->resources, static fn ($r) => $r->isDiscovered),
323+
prompts: array_filter($this->prompts, static fn ($p) => $p->isDiscovered),
324+
resourceTemplates: array_filter($this->resourceTemplates, static fn ($t) => $t->isDiscovered),
351325
);
352326
}
353327

@@ -360,21 +334,29 @@ public function setDiscoveryState(DiscoveryState $state): void
360334
// Clear existing discovered elements
361335
$this->clear();
362336

363-
// Import new discovered elements
337+
// Import new discovered elements — skip any that conflict with manual or dynamic registrations
364338
foreach ($state->getTools() as $name => $tool) {
365-
$this->tools[$name] = $tool;
339+
if (!isset($this->tools[$name])) {
340+
$this->tools[$name] = new ToolReference($tool->tool, $tool->handler, isDiscovered: true);
341+
}
366342
}
367343

368344
foreach ($state->getResources() as $uri => $resource) {
369-
$this->resources[$uri] = $resource;
345+
if (!isset($this->resources[$uri])) {
346+
$this->resources[$uri] = new ResourceReference($resource->resource, $resource->handler, isDiscovered: true);
347+
}
370348
}
371349

372350
foreach ($state->getPrompts() as $name => $prompt) {
373-
$this->prompts[$name] = $prompt;
351+
if (!isset($this->prompts[$name])) {
352+
$this->prompts[$name] = new PromptReference($prompt->prompt, $prompt->handler, completionProviders: $prompt->completionProviders, isDiscovered: true);
353+
}
374354
}
375355

376356
foreach ($state->getResourceTemplates() as $uriTemplate => $template) {
377-
$this->resourceTemplates[$uriTemplate] = $template;
357+
if (!isset($this->resourceTemplates[$uriTemplate])) {
358+
$this->resourceTemplates[$uriTemplate] = new ResourceTemplateReference($template->resourceTemplate, $template->handler, completionProviders: $template->completionProviders, isDiscovered: true);
359+
}
378360
}
379361

380362
// Dispatch events for the imported elements

‎src/Capability/Registry/ElementReference.php‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ class ElementReference
2424
public function __construct(
2525
public readonly \Closure|array|string $handler,
2626
public readonly bool $isManual = false,
27+
public readonly bool $isDiscovered = false,
2728
) {
2829
}
2930
}

‎src/Capability/Registry/PromptReference.php‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,8 +31,9 @@ public function __construct(
3131
\Closure|array|string $handler,
3232
bool $isManual = false,
3333
public readonly array $completionProviders = [],
34+
bool $isDiscovered = false,
3435
) {
35-
parent::__construct($handler, $isManual);
36+
parent::__construct($handler, $isManual, $isDiscovered);
3637
}
3738

3839
/**

‎src/Capability/Registry/ResourceReference.php‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,9 @@ public function __construct(
2929
public readonly Resource $resource,
3030
callable|array|string $handler,
3131
bool $isManual = false,
32+
bool $isDiscovered = false,
3233
) {
33-
parent::__construct($handler, $isManual);
34+
parent::__construct($handler, $isManual, $isDiscovered);
3435
}
3536

3637
/**

‎src/Capability/Registry/ResourceTemplateReference.php‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,8 +38,9 @@ public function __construct(
3838
callable|array|string $handler,
3939
bool $isManual = false,
4040
public readonly array $completionProviders = [],
41+
bool $isDiscovered = false,
4142
) {
42-
parent::__construct($handler, $isManual);
43+
parent::__construct($handler, $isManual, $isDiscovered);
4344

4445
$this->compileTemplate();
4546
}

‎src/Capability/Registry/ToolReference.php‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,9 @@ public function __construct(
2929
public readonly Tool $tool,
3030
callable|array|string $handler,
3131
bool $isManual = false,
32+
bool $isDiscovered = false,
3233
) {
33-
parent::__construct($handler, $isManual);
34+
parent::__construct($handler, $isManual, $isDiscovered);
3435
}
3536

3637
/**

‎tests/Inspector/Http/snapshots/HttpCombinedRegistrationTest-resources_list.json‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
{
22
"resources": [
33
{
4-
"name": "priority_config_discovered",
4+
"name": "priority_config_manual",
55
"uri": "config://priority",
6-
"description": "A resource discovered via attributes.\n\nThis will be overridden by a manual registration with the same URI."
6+
"description": "Manually registered resource that overrides a discovered one."
77
}
88
]
99
}

‎tests/Inspector/Http/snapshots/HttpCombinedRegistrationTest-resources_read-config_priority.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
{
44
"uri": "config://priority",
55
"mimeType": "text/plain",
6-
"text": "Discovered Priority Config: Low"
6+
"text": "Manual Priority Config: HIGH (overrides discovered)"
77
}
88
]
99
}

‎tests/Unit/Capability/RegistryTest.php‎

Lines changed: 169 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
namespace Mcp\Tests\Unit\Capability;
1313

1414
use Mcp\Capability\Completion\EnumCompletionProvider;
15+
use Mcp\Capability\Discovery\DiscoveryState;
1516
use Mcp\Capability\Registry;
1617
use Mcp\Capability\Registry\PromptReference;
1718
use Mcp\Capability\Registry\ResourceReference;
@@ -426,42 +427,70 @@ public function testClearRemovesOnlyDiscoveredElements(): void
426427
$manualTemplate = $this->createValidResourceTemplate('manual://{id}');
427428
$discoveredTemplate = $this->createValidResourceTemplate('discovered://{id}');
428429

430+
// Register manual elements directly
429431
$this->registry->registerTool($manualTool, static fn () => 'manual', true);
430-
$this->registry->registerTool($discoveredTool, static fn () => 'discovered');
431432
$this->registry->registerResource($manualResource, static fn () => 'manual', true);
432-
$this->registry->registerResource($discoveredResource, static fn () => 'discovered');
433433
$this->registry->registerPrompt($manualPrompt, static fn () => [], [], true);
434-
$this->registry->registerPrompt($discoveredPrompt, static fn () => []);
435434
$this->registry->registerResourceTemplate($manualTemplate, static fn () => 'manual', [], true);
436-
$this->registry->registerResourceTemplate($discoveredTemplate, static fn () => 'discovered');
437435

438-
// Test that all elements exist
439-
$this->registry->getTool('manual_tool');
440-
$this->registry->getResource('test://manual');
441-
$this->registry->getPrompt('manual_prompt');
442-
$this->registry->getResourceTemplate('manual://{id}');
443-
$this->registry->getTool('discovered_tool');
444-
$this->registry->getResource('test://discovered');
445-
$this->registry->getPrompt('discovered_prompt');
446-
$this->registry->getResourceTemplate('discovered://{id}');
436+
// Import discovered elements via setDiscoveryState
437+
$this->registry->setDiscoveryState(new DiscoveryState(
438+
tools: ['discovered_tool' => new ToolReference($discoveredTool, static fn () => 'discovered')],
439+
resources: ['test://discovered' => new ResourceReference($discoveredResource, static fn () => 'discovered')],
440+
prompts: ['discovered_prompt' => new PromptReference($discoveredPrompt, static fn () => [])],
441+
resourceTemplates: ['discovered://{id}' => new ResourceTemplateReference($discoveredTemplate, static fn () => 'discovered')],
442+
));
443+
444+
// All elements exist before clear
445+
$this->assertInstanceOf(ToolReference::class, $this->registry->getTool('discovered_tool'));
447446

448447
$this->registry->clear();
449448

450-
// Manual elements should still exist
451-
$this->registry->getTool('manual_tool');
452-
$this->registry->getResource('test://manual');
453-
$this->registry->getPrompt('manual_prompt');
454-
$this->registry->getResourceTemplate('manual://{id}');
449+
// Manual elements survive
450+
$this->assertInstanceOf(ToolReference::class, $this->registry->getTool('manual_tool'));
451+
$this->assertInstanceOf(ResourceReference::class, $this->registry->getResource('test://manual'));
452+
$this->assertInstanceOf(PromptReference::class, $this->registry->getPrompt('manual_prompt'));
453+
$this->assertInstanceOf(ResourceTemplateReference::class, $this->registry->getResourceTemplate('manual://{id}'));
455454

456-
// Test that all discovered elements throw exceptions
455+
// Discovered elements are gone
457456
$this->expectException(ToolNotFoundException::class);
458457
$this->registry->getTool('discovered_tool');
458+
}
459+
460+
public function testClearRemovesDiscoveredResources(): void
461+
{
462+
$discoveredResource = $this->createValidResource('test://discovered');
463+
$this->registry->setDiscoveryState(new DiscoveryState(
464+
resources: ['test://discovered' => new ResourceReference($discoveredResource, static fn () => 'discovered')],
465+
));
466+
467+
$this->registry->clear();
459468

460469
$this->expectException(ResourceNotFoundException::class);
461470
$this->registry->getResource('test://discovered');
471+
}
472+
473+
public function testClearRemovesDiscoveredPrompts(): void
474+
{
475+
$discoveredPrompt = $this->createValidPrompt('discovered_prompt');
476+
$this->registry->setDiscoveryState(new DiscoveryState(
477+
prompts: ['discovered_prompt' => new PromptReference($discoveredPrompt, static fn () => [])],
478+
));
479+
480+
$this->registry->clear();
462481

463482
$this->expectException(PromptNotFoundException::class);
464483
$this->registry->getPrompt('discovered_prompt');
484+
}
485+
486+
public function testClearRemovesDiscoveredResourceTemplates(): void
487+
{
488+
$discoveredTemplate = $this->createValidResourceTemplate('discovered://{id}');
489+
$this->registry->setDiscoveryState(new DiscoveryState(
490+
resourceTemplates: ['discovered://{id}' => new ResourceTemplateReference($discoveredTemplate, static fn () => 'discovered')],
491+
));
492+
493+
$this->registry->clear();
465494

466495
$this->expectException(ResourceNotFoundException::class);
467496
$this->registry->getResourceTemplate('discovered://{id}');
@@ -611,6 +640,127 @@ public function testExtractStructuredContentReturnsArrayDirectlyForArrayOutputSc
611640
], $structuredContent);
612641
}
613642

643+
public function testClearPreservesDynamicallyRegisteredElements(): void
644+
{
645+
// 1. Register a manual tool
646+
$manualTool = $this->createValidTool('manual_tool');
647+
$this->registry->registerTool($manualTool, static fn () => 'manual', true);
648+
649+
// 2. Import discovered tools via setDiscoveryState
650+
$discoveredTool = $this->createValidTool('discovered_tool');
651+
$this->registry->setDiscoveryState(new DiscoveryState(
652+
tools: ['discovered_tool' => new ToolReference($discoveredTool, static fn () => 'discovered')],
653+
));
654+
655+
// 3. Register a dynamic tool (not manual, not discovered)
656+
$dynamicTool = $this->createValidTool('dynamic_tool');
657+
$this->registry->registerTool($dynamicTool, static fn () => 'dynamic');
658+
659+
// 4. Restore discovery state again (simulates next HTTP request)
660+
$discoveredTool2 = $this->createValidTool('discovered_tool_v2');
661+
$this->registry->setDiscoveryState(new DiscoveryState(
662+
tools: ['discovered_tool_v2' => new ToolReference($discoveredTool2, static fn () => 'discovered_v2')],
663+
));
664+
665+
// Manual tool survives
666+
$this->assertInstanceOf(ToolReference::class, $this->registry->getTool('manual_tool'));
667+
// Dynamic tool survives
668+
$this->assertInstanceOf(ToolReference::class, $this->registry->getTool('dynamic_tool'));
669+
// New discovered tool is present
670+
$this->assertInstanceOf(ToolReference::class, $this->registry->getTool('discovered_tool_v2'));
671+
// Old discovered tool is gone
672+
$this->expectException(ToolNotFoundException::class);
673+
$this->registry->getTool('discovered_tool');
674+
}
675+
676+
public function testGetDiscoveryStateExcludesDynamicTools(): void
677+
{
678+
// Import discovered tool via setDiscoveryState
679+
$discoveredTool = $this->createValidTool('discovered_tool');
680+
$this->registry->setDiscoveryState(new DiscoveryState(
681+
tools: ['discovered_tool' => new ToolReference($discoveredTool, static fn () => 'discovered')],
682+
));
683+
684+
// Register a dynamic tool
685+
$dynamicTool = $this->createValidTool('dynamic_tool');
686+
$this->registry->registerTool($dynamicTool, static fn () => 'dynamic');
687+
688+
// Register a manual tool
689+
$manualTool = $this->createValidTool('manual_tool');
690+
$this->registry->registerTool($manualTool, static fn () => 'manual', true);
691+
692+
$state = $this->registry->getDiscoveryState();
693+
694+
$this->assertArrayHasKey('discovered_tool', $state->getTools());
695+
$this->assertArrayNotHasKey('dynamic_tool', $state->getTools());
696+
$this->assertArrayNotHasKey('manual_tool', $state->getTools());
697+
}
698+
699+
public function testSetDiscoveryStateRoundTrip(): void
700+
{
701+
// Import initial discovered state
702+
$tool = $this->createValidTool('round_trip_tool');
703+
$resource = $this->createValidResource('test://round-trip');
704+
$prompt = $this->createValidPrompt('round_trip_prompt');
705+
$template = $this->createValidResourceTemplate('round-trip://{id}');
706+
707+
$initialState = new DiscoveryState(
708+
tools: ['round_trip_tool' => new ToolReference($tool, static fn () => 'result')],
709+
resources: ['test://round-trip' => new ResourceReference($resource, static fn () => 'content')],
710+
prompts: ['round_trip_prompt' => new PromptReference($prompt, static fn () => [])],
711+
resourceTemplates: ['round-trip://{id}' => new ResourceTemplateReference($template, static fn () => 'tpl')],
712+
);
713+
714+
$this->registry->setDiscoveryState($initialState);
715+
716+
// Round-trip: get and set again
717+
$exportedState = $this->registry->getDiscoveryState();
718+
$this->registry->setDiscoveryState($exportedState);
719+
720+
// All elements still present
721+
$this->assertInstanceOf(ToolReference::class, $this->registry->getTool('round_trip_tool'));
722+
$this->assertInstanceOf(ResourceReference::class, $this->registry->getResource('test://round-trip'));
723+
$this->assertInstanceOf(PromptReference::class, $this->registry->getPrompt('round_trip_prompt'));
724+
$this->assertInstanceOf(ResourceTemplateReference::class, $this->registry->getResourceTemplate('round-trip://{id}'));
725+
726+
// Exported state matches
727+
$reExportedState = $this->registry->getDiscoveryState();
728+
$this->assertCount(\count($exportedState->getTools()), $reExportedState->getTools());
729+
$this->assertCount(\count($exportedState->getResources()), $reExportedState->getResources());
730+
$this->assertCount(\count($exportedState->getPrompts()), $reExportedState->getPrompts());
731+
$this->assertCount(\count($exportedState->getResourceTemplates()), $reExportedState->getResourceTemplates());
732+
}
733+
734+
public function testSetDiscoveryStateDoesNotOverwriteManualOrDynamicTools(): void
735+
{
736+
// Register a manual tool
737+
$manualTool = $this->createValidTool('conflict_tool');
738+
$this->registry->registerTool($manualTool, static fn () => 'manual_result', true);
739+
740+
// Register a dynamic tool
741+
$dynamicTool = $this->createValidTool('dynamic_conflict');
742+
$this->registry->registerTool($dynamicTool, static fn () => 'dynamic_result');
743+
744+
// Try to import discovered tools with same names
745+
$discoveredConflict = $this->createValidTool('conflict_tool');
746+
$discoveredDynConflict = $this->createValidTool('dynamic_conflict');
747+
$this->registry->setDiscoveryState(new DiscoveryState(
748+
tools: [
749+
'conflict_tool' => new ToolReference($discoveredConflict, static fn () => 'discovered_result'),
750+
'dynamic_conflict' => new ToolReference($discoveredDynConflict, static fn () => 'discovered_dyn_result'),
751+
],
752+
));
753+
754+
// Manual tool preserved with original handler
755+
$manualRef = $this->registry->getTool('conflict_tool');
756+
$this->assertTrue($manualRef->isManual);
757+
$this->assertEquals('manual_result', ($manualRef->handler)());
758+
759+
// Dynamic tool preserved with original handler
760+
$dynamicRef = $this->registry->getTool('dynamic_conflict');
761+
$this->assertEquals('dynamic_result', ($dynamicRef->handler)());
762+
}
763+
614764
private function createValidTool(string $name, ?array $outputSchema = null): Tool
615765
{
616766
return new Tool(

0 commit comments

Comments
 (0)