Merge pull request #9 from EdouardCourty/feat/add-unsafe-prompt-argument

feat(prompt): support unsafe arguments

Author: EdouardCourty

src/MethodHandler/PromptsGetMethodHandler.php

@@ -17,7 +17,6 @@
 use Ecourty\McpServerBundle\Service\InputSanitizer;
 use Ecourty\McpServerBundle\Service\PromptRegistry;
 use Psr\EventDispatcher\EventDispatcherInterface;
-use Symfony\Component\Console\Attribute\Argument;
 use Symfony\Component\HttpFoundation\Response;
 
 /**
@@ -61,21 +60,20 @@ public function handle(JsonRpcRequest $request): array
 
         try {
             $arguments = $request->params['arguments'] ?? [];
-            $sanitizedArguments = $this->inputSanitizer->sanitize($arguments);
 
-            $this->validatePromptArguments($sanitizedArguments, $promptDefinition);
-            $argumentCollection = ArgumentCollection::fromArray($sanitizedArguments);
+            $this->validatePromptArguments($arguments, $promptDefinition);
+            $sanitizedArguments = $this->getSanitizedArguments($arguments, $promptDefinition);
 
             if (method_exists($prompt, '__invoke') === false) {
                 throw new \LogicException(\sprintf('Prompt "%s" does not implement the __invoke method.', $promptName));
             }
 
             $this->eventDispatcher?->dispatch(new PromptGetEvent(
                 promptName: $promptName,
-                arguments: $argumentCollection,
+                arguments: $sanitizedArguments,
             ));
 
-            $promptResult = $prompt->__invoke($argumentCollection);
+            $promptResult = $prompt->__invoke($sanitizedArguments);
 
             if ($promptResult instanceof PromptResult === false) {
                 throw new \LogicException(\sprintf(
@@ -85,11 +83,11 @@ public function handle(JsonRpcRequest $request): array
                 ));
             }
 
-            $this->eventDispatcher?->dispatch(new PromptResultEvent($promptName, $argumentCollection, $promptResult));
+            $this->eventDispatcher?->dispatch(new PromptResultEvent($promptName, $sanitizedArguments, $promptResult));
         } catch (\Throwable $exception) {
             $this->eventDispatcher?->dispatch(new PromptExceptionEvent(
                 promptName: $promptName,
-                arguments: $argumentCollection ?? null,
+                arguments: $sanitizedArguments ?? null,
                 exception: $exception,
             ));
 
@@ -103,6 +101,33 @@ public function handle(JsonRpcRequest $request): array
         return $promptResult->toArray();
     }
 
+    /**
+     * @param array<string, mixed> $inputArgument
+     */
+    private function getSanitizedArguments(array $inputArgument, PromptDefinition $promptDefinition): ArgumentCollection
+    {
+        $sanitizedArguments = [];
+        $promptArguments = $promptDefinition->arguments ?? [];
+
+        foreach ($promptArguments as $argument) {
+            $argumentName = $argument->name;
+            $unsafeArgumentValue = $inputArgument[$argumentName] ?? null;
+
+            if ($argument->allowUnsafe === true) {
+                $sanitizedArguments[$argumentName] = $unsafeArgumentValue;
+
+                continue;
+            }
+
+            $sanitizedArguments[$argumentName] = $this->inputSanitizer->sanitize($unsafeArgumentValue);
+        }
+
+        return ArgumentCollection::fromArray($sanitizedArguments);
+    }
+
+    /**
+     * @param array<string, mixed> $arguments
+     */
     private function validatePromptArguments(array $arguments, PromptDefinition $promptDefinition): void
     {
         $promptArguments = $promptDefinition->arguments ?? [];

src/Prompt/Argument.php

@@ -9,22 +9,28 @@
  */
 class Argument
 {
+    /**
+     * @param bool $required Indicates if the argument is required (will throw an error if true and not provided).
+     * @param bool $allowUnsafe Indicates if the argument allows unsafe content (will not be sanitized if true).
+     */
     public function __construct(
         public readonly string $name,
         public readonly string $description,
         public readonly bool $required = true,
+        public readonly bool $allowUnsafe = false,
     ) {
     }
 
     /**
-     * @return array{name: string, description: string, required: bool}
+     * @return array{name: string, description: string, required: bool, allowUnsafe: bool}
      */
     public function toArray(): array
     {
         return [
             'name' => $this->name,
             'description' => $this->description,
             'required' => $this->required,
+            'allowUnsafe' => $this->allowUnsafe,
         ];
     }
 }

src/Service/PromptRegistry.php

@@ -53,7 +53,7 @@ public function getPromptsDefinitions(): array
     /**
      * @internal
      *
-     * @param array<array{name: string, description: string, required: bool}> $argumentDefinitions
+     * @param array<array{name: string, description: string, required: bool, allowUnsafe: bool}> $argumentDefinitions
      */
     public function addPromptDefinition(
         string $name,
@@ -70,6 +70,7 @@ public function addPromptDefinition(
                 name: $argumentDefinition['name'],
                 description: $argumentDefinition['description'],
                 required: $argumentDefinition['required'],
+                allowUnsafe: $argumentDefinition['allowUnsafe'],
             );
         }
 

tests/Controller/EntrypointControllerTest.php

@@ -454,4 +454,37 @@ public function testPromptGetWithNonRequiredParameterLeftEmpty(): void
         $this->assertArrayHasKey('result', $responseContent);
         $this->assertArrayNotHasKey('error', $responseContent);
     }
+
+    public function testPromptGetWithUnsafeParameterAllowed(): void
+    {
+        $unsafeContent = '<script>alert("This is unsafe!");</script>';
+
+        $response = $this->request(
+            method: Request::METHOD_POST,
+            url: '/mcp',
+            body: [
+                'jsonrpc' => '2.0',
+                'id' => 1,
+                'method' => 'prompts/get',
+                'params' => [
+                    'name' => 'generate-git-commit-message',
+                    'arguments' => [
+                        'changes' => $unsafeContent,
+                        'scope' => 'feature',
+                    ],
+                ],
+            ],
+        );
+
+        $responseContent = json_decode((string) $response->getContent(), true);
+        $this->assertNotFalse($responseContent);
+
+        $this->assertSame('2.0', $responseContent['jsonrpc']);
+        $this->assertSame(1, $responseContent['id']);
+        $this->assertArrayHasKey('result', $responseContent);
+
+        // Check if the unsafe content is present in the response
+        $resultContent = $responseContent['result'];
+        $this->assertStringContainsString($unsafeContent, $resultContent['messages'][0]['content']['text']);
+    }
 }

tests/Service/PromptRegistryTest.php

@@ -62,7 +62,7 @@ public static function providePromptDefinitionTestData(): array
                 'name' => 'generate-git-commit-message',
                 'expectedDescription' => 'Generate a git commit message based on the provided changes.',
                 'expectedArguments' => [
-                    new Argument('changes', 'The changed made in the codebase', true),
+                    new Argument('changes', 'The changed made in the codebase', true, true),
                     new Argument('scope', 'The scope of the changes, e.g., feature, bugfix, etc.', true),
                 ],
             ],

tests/TestApp/src/Prompt/GenerateGitCommitMessage.php

@@ -16,7 +16,7 @@
     name: 'generate-git-commit-message',
     description: 'Generate a git commit message based on the provided changes.',
     arguments: [
-        new Argument(name: 'changes', description: 'The changed made in the codebase'),
+        new Argument(name: 'changes', description: 'The changed made in the codebase', allowUnsafe: true),
         new Argument(name: 'scope', description: 'The scope of the changes, e.g., feature, bugfix, etc.'),
     ],
 )]

tests/TestApp/src/Prompt/GreetingPrompt.php

@@ -15,7 +15,7 @@
 #[AsPrompt(
     name: 'greeting',
     arguments: [
-        new Argument('name', description: 'The name of the person to greet.', required: false),
+        new Argument('name', description: 'The name of the person to greet.', required: false, allowUnsafe: true),
     ],
 )]
 class GreetingPrompt