Skip to content

Commit c89c99c

Browse files
author
klapaudius
committed
Fix progress token cleanup and improve error response handling
 - Wrap tool execution in try-finally block to ensure progress tokens are always unregistered, preventing resource leaks when tools throw exceptions - Update DataUtil to correctly identify both result and error responses in message parsing - Add comprehensive test coverage for MCPProtocol response handling, MCPServer registration methods, and ToolsCallHandler edge cases including exception scenarios, deprecated schemas, and special token values
1 parent bb6ca8e commit c89c99c

File tree

5 files changed

+882
-4
lines changed

5 files changed

+882
-4
lines changed

‎src/Server/Request/ToolsCallHandler.php‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -102,9 +102,11 @@ public function execute(string $method, string|int $clientId, string|int $messag
102102
$tool->setSamplingClient($this->samplingClient);
103103
}
104104

105-
$result = $tool->execute($arguments);
106-
107-
$this->progressNotifierRepository->unregisterToken($progressToken);
105+
try {
106+
$result = $tool->execute($arguments);
107+
} finally {
108+
$this->progressNotifierRepository->unregisterToken($progressToken);
109+
}
108110

109111
if ($method === 'tools/call') {
110112
if (! $result instanceof ToolResultInterface) {

‎src/Utils/DataUtil.php‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ public static function makeRequestData(string $clientId, array $message): Reques
1717
} else {
1818
$data = NotificationData::fromArray(data: array_merge(['clientId' => $clientId], $message));
1919
}
20-
} elseif (isset($message['result'])) {
20+
} elseif (isset($message['result']) || isset($message['error'])) {
2121
if (isset($message['id'])) {
2222
$data = ResponseData::fromArray(data: $message);
2323
} else {

‎tests/Protocol/MCPProtocolTest.php‎

Lines changed: 261 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -523,4 +523,265 @@ public function test_get_response_result_skips_null_messages(): void
523523

524524
$this->assertEquals($decodedMessages, $result);
525525
}
526+
527+
/**
528+
* Test that handleResponseProcess executes matching handler successfully
529+
*/
530+
public function test_handle_response_process_with_matching_handler(): void
531+
{
532+
$clientId = 'response_client';
533+
$messageId = 'resp_123';
534+
$result = ['data' => 'response'];
535+
$responseMessage = [
536+
'jsonrpc' => '2.0',
537+
'id' => $messageId,
538+
'result' => $result,
539+
];
540+
541+
$mockHandler = $this->createMock(\KLP\KlpMcpServer\Protocol\Handlers\ResponseHandler::class);
542+
$mockHandler->expects($this->once())
543+
->method('isHandle')
544+
->with($messageId)
545+
->willReturn(true);
546+
547+
$mockHandler->expects($this->once())
548+
->method('execute')
549+
->with($clientId, $messageId, $result, null);
550+
551+
$this->mcpProtocol->registerResponseHandler($mockHandler);
552+
553+
$this->mockTransport
554+
->expects($this->never())
555+
->method('pushMessage');
556+
557+
$this->mcpProtocol->handleMessage($clientId, $responseMessage);
558+
}
559+
560+
/**
561+
* Test that handleResponseProcess silently ignores when no handler matches
562+
*/
563+
public function test_handle_response_process_with_no_matching_handler(): void
564+
{
565+
$clientId = 'response_client';
566+
$messageId = 'resp_456';
567+
$responseMessage = [
568+
'jsonrpc' => '2.0',
569+
'id' => $messageId,
570+
'result' => ['data' => 'test'],
571+
];
572+
573+
$mockHandler = $this->createMock(\KLP\KlpMcpServer\Protocol\Handlers\ResponseHandler::class);
574+
$mockHandler->expects($this->once())
575+
->method('isHandle')
576+
->with($messageId)
577+
->willReturn(false);
578+
579+
$mockHandler->expects($this->never())
580+
->method('execute');
581+
582+
$this->mcpProtocol->registerResponseHandler($mockHandler);
583+
584+
$this->mockTransport
585+
->expects($this->never())
586+
->method('pushMessage');
587+
588+
$this->mcpProtocol->handleMessage($clientId, $responseMessage);
589+
}
590+
591+
/**
592+
* Test that handleResponseProcess pushes error when handler throws JsonRpcErrorException
593+
*/
594+
public function test_handle_response_process_handler_throws_json_rpc_error(): void
595+
{
596+
$clientId = 'response_client';
597+
$messageId = 'resp_789';
598+
$responseMessage = [
599+
'jsonrpc' => '2.0',
600+
'id' => $messageId,
601+
'result' => ['data' => 'test'],
602+
];
603+
604+
$mockHandler = $this->createMock(\KLP\KlpMcpServer\Protocol\Handlers\ResponseHandler::class);
605+
$mockHandler->expects($this->once())
606+
->method('isHandle')
607+
->with($messageId)
608+
->willReturn(true);
609+
610+
$expectedException = new \KLP\KlpMcpServer\Exceptions\JsonRpcErrorException(
611+
'Handler error',
612+
\KLP\KlpMcpServer\Exceptions\Enums\JsonRpcErrorCode::INTERNAL_ERROR
613+
);
614+
615+
$mockHandler->expects($this->once())
616+
->method('execute')
617+
->with($clientId, $messageId, ['data' => 'test'], null)
618+
->willThrowException($expectedException);
619+
620+
$this->mcpProtocol->registerResponseHandler($mockHandler);
621+
622+
$this->mockTransport
623+
->expects($this->once())
624+
->method('pushMessage')
625+
->with($this->callback(function (...$args) use ($clientId, $messageId) {
626+
$data = $args[1];
627+
$this->assertEquals($clientId, $args[0]);
628+
$this->assertEquals('2.0', $data['jsonrpc']);
629+
$this->assertEquals($messageId, $data['id']);
630+
$this->assertEquals(-32603, $data['error']['code']);
631+
$this->assertEquals('Handler error', $data['error']['message']);
632+
633+
return true;
634+
}));
635+
636+
$this->mcpProtocol->handleMessage($clientId, $responseMessage);
637+
}
638+
639+
/**
640+
* Test that handleResponseProcess wraps generic Exception as INTERNAL_ERROR
641+
*/
642+
public function test_handle_response_process_handler_throws_generic_exception(): void
643+
{
644+
$clientId = 'response_client';
645+
$messageId = 'resp_101';
646+
$responseMessage = [
647+
'jsonrpc' => '2.0',
648+
'id' => $messageId,
649+
'result' => ['data' => 'test'],
650+
];
651+
652+
$mockHandler = $this->createMock(\KLP\KlpMcpServer\Protocol\Handlers\ResponseHandler::class);
653+
$mockHandler->expects($this->once())
654+
->method('isHandle')
655+
->with($messageId)
656+
->willReturn(true);
657+
658+
$mockHandler->expects($this->once())
659+
->method('execute')
660+
->with($clientId, $messageId, ['data' => 'test'], null)
661+
->willThrowException(new \RuntimeException('Unexpected error'));
662+
663+
$this->mcpProtocol->registerResponseHandler($mockHandler);
664+
665+
$this->mockTransport
666+
->expects($this->once())
667+
->method('pushMessage')
668+
->with($this->callback(function (...$args) use ($clientId, $messageId) {
669+
$data = $args[1];
670+
$this->assertEquals($clientId, $args[0]);
671+
$this->assertEquals('2.0', $data['jsonrpc']);
672+
$this->assertEquals($messageId, $data['id']);
673+
$this->assertEquals(-32603, $data['error']['code']);
674+
$this->assertEquals('Unexpected error', $data['error']['message']);
675+
676+
return true;
677+
}));
678+
679+
$this->mcpProtocol->handleMessage($clientId, $responseMessage);
680+
}
681+
682+
/**
683+
* Test that handleResponseProcess returns early on first matching handler
684+
*/
685+
public function test_handle_response_process_early_return_first_match(): void
686+
{
687+
$clientId = 'response_client';
688+
$messageId = 'resp_202';
689+
$responseMessage = [
690+
'jsonrpc' => '2.0',
691+
'id' => $messageId,
692+
'result' => ['data' => 'test'],
693+
];
694+
695+
$firstHandler = $this->createMock(\KLP\KlpMcpServer\Protocol\Handlers\ResponseHandler::class);
696+
$firstHandler->expects($this->once())
697+
->method('isHandle')
698+
->with($messageId)
699+
->willReturn(true);
700+
701+
$firstHandler->expects($this->once())
702+
->method('execute')
703+
->with($clientId, $messageId, ['data' => 'test'], null);
704+
705+
$secondHandler = $this->createMock(\KLP\KlpMcpServer\Protocol\Handlers\ResponseHandler::class);
706+
$secondHandler->expects($this->never())
707+
->method('isHandle');
708+
709+
$secondHandler->expects($this->never())
710+
->method('execute');
711+
712+
$this->mcpProtocol->registerResponseHandler($firstHandler);
713+
$this->mcpProtocol->registerResponseHandler($secondHandler);
714+
715+
$this->mockTransport
716+
->expects($this->never())
717+
->method('pushMessage');
718+
719+
$this->mcpProtocol->handleMessage($clientId, $responseMessage);
720+
}
721+
722+
/**
723+
* Test that handleResponseProcess passes result data correctly to handler
724+
*/
725+
public function test_handle_response_process_with_result_data(): void
726+
{
727+
$clientId = 'response_client';
728+
$messageId = 'resp_303';
729+
$result = ['key1' => 'value1', 'key2' => ['nested' => 'data']];
730+
$responseMessage = [
731+
'jsonrpc' => '2.0',
732+
'id' => $messageId,
733+
'result' => $result,
734+
];
735+
736+
$mockHandler = $this->createMock(\KLP\KlpMcpServer\Protocol\Handlers\ResponseHandler::class);
737+
$mockHandler->expects($this->once())
738+
->method('isHandle')
739+
->with($messageId)
740+
->willReturn(true);
741+
742+
$mockHandler->expects($this->once())
743+
->method('execute')
744+
->with($clientId, $messageId, $result, null);
745+
746+
$this->mcpProtocol->registerResponseHandler($mockHandler);
747+
748+
$this->mockTransport
749+
->expects($this->never())
750+
->method('pushMessage');
751+
752+
$this->mcpProtocol->handleMessage($clientId, $responseMessage);
753+
}
754+
755+
/**
756+
* Test that handleResponseProcess passes error data correctly to handler
757+
*/
758+
public function test_handle_response_process_with_error_data(): void
759+
{
760+
$clientId = 'response_client';
761+
$messageId = 'resp_404';
762+
$error = ['code' => -32600, 'message' => 'Invalid Request'];
763+
$responseMessage = [
764+
'jsonrpc' => '2.0',
765+
'id' => $messageId,
766+
'error' => $error,
767+
];
768+
769+
$mockHandler = $this->createMock(\KLP\KlpMcpServer\Protocol\Handlers\ResponseHandler::class);
770+
$mockHandler->expects($this->once())
771+
->method('isHandle')
772+
->with($messageId)
773+
->willReturn(true);
774+
775+
$mockHandler->expects($this->once())
776+
->method('execute')
777+
->with($clientId, $messageId, null, $error);
778+
779+
$this->mcpProtocol->registerResponseHandler($mockHandler);
780+
781+
$this->mockTransport
782+
->expects($this->never())
783+
->method('pushMessage');
784+
785+
$this->mcpProtocol->handleMessage($clientId, $responseMessage);
786+
}
526787
}

0 commit comments

Comments
 (0)