using messenger HandleTrait as QueryBus with appropriate result typing

README.md

@@ -21,6 +21,8 @@ This extension provides following features:
 * Provides correct return type for `InputBag::get()` method based on the `$default` parameter.
 * Provides correct return type for `InputBag::all()` method based on the `$key` parameter.
 * Provides correct return types for `TreeBuilder` and `NodeDefinition` objects.
+* Provides correct return type for Messenger `HandleTrait::handle()` method based on the message type.
+* Provides configurable return type resolution for methods that internally use Messenger `HandleTrait`.
 * Notifies you when you try to get an unregistered service from the container.
 * Notifies you when you try to get a private service from the container.
 * Optionally correct return types for `InputInterface::getArgument()`, `::getOption`, `::hasArgument`, and `::hasOption`.
@@ -168,3 +170,94 @@ Call the new env in your `console-application.php`:
 ```php
 $kernel = new \App\Kernel('phpstan_env', (bool) $_SERVER['APP_DEBUG']);
 ```
+
+## Messenger HandleTrait Wrappers
+
+The extension provides advanced type inference for methods that internally use Symfony Messenger's `HandleTrait`. This feature is particularly useful for query bus implementations (in CQRS pattern) that use/wrap the `HandleTrait::handle()` method.
+
+### Configuration
+
+```neon
+parameters:
+    symfony:
+        messenger:
+            handleTraitWrappers:
+                - App\Bus\QueryBus::dispatch
+                - App\Bus\QueryBus::execute
+                - App\Bus\QueryBusInterface::dispatch
+```
+
+### Message Handlers
+
+```php
+use Symfony\Component\Messenger\Attribute\AsMessageHandler;
+
+// Product handler that returns Product
+#[AsMessageHandler]
+class GetProductQueryHandler
+{
+    public function __invoke(GetProductQuery $query): Product
+    {
+        return $this->productRepository->get($query->productId);
+    }
+}
+```
+
+### PHP Examples
+
+```php
+use Symfony\Component\Messenger\HandleTrait;
+use Symfony\Component\Messenger\MessageBusInterface;
+
+// Basic query bus implementation
+class QueryBus
+{
+    use HandleTrait;
+
+    public function __construct(MessageBusInterface $messageBus)
+    {
+        $this->messageBus = $messageBus;
+    }
+
+    public function dispatch(object $query): mixed
+    {
+        return $this->handle($query); // Return type will be inferred
+    }
+
+    // Multiple methods per class example
+    public function execute(object $message): mixed
+    {
+        return $this->handle($message);
+    }
+}
+
+// Interface-based configuration example
+interface QueryBusInterface
+{
+    public function dispatch(object $query): mixed;
+}
+
+class QueryBusWithInterface implements QueryBusInterface
+{
+    use HandleTrait;
+
+    public function __construct(MessageBusInterface $queryBus)
+    {
+        $this->messageBus = $queryBus;
+    }
+
+    public function dispatch(object $query): mixed
+    {
+        return $this->handle($query);
+    }
+}
+
+// Usage examples with proper type inference
+$query = new GetProductQuery($productId);
+$queryBus = new QueryBus($messageBus);
+$queryBusWithInterface = new QueryBusWithInterface($messageBus);
+
+$product = $queryBus->dispatch($query); // Returns: Product
+$product2 = $queryBus->execute($query); // Returns: Product
+$product3 = $queryBusWithInterface->dispatch($query); // Returns: Product
+```

extension.neon

@@ -11,6 +11,8 @@ parameters:
 		constantHassers: true
 		console_application_loader: null
 		consoleApplicationLoader: null
+		messenger:
+			handleTraitWrappers: []
 	featureToggles:
 		skipCheckGenericClasses:
 			- Symfony\Component\Form\AbstractType
@@ -115,6 +117,9 @@ parametersSchema:
 		constantHassers: bool()
 		console_application_loader: schema(string(), nullable())
 		consoleApplicationLoader: schema(string(), nullable())
+		messenger: structure([
+			handleTraitWrappers: listOf(string())
+		])
 	])
 
 services:
@@ -214,6 +219,11 @@ services:
 		class: PHPStan\Type\Symfony\MessengerHandleTraitReturnTypeExtension
 		tags: [phpstan.broker.expressionTypeResolverExtension]
 
+	# Messenger HandleTrait wrappers return type
+	-
+		class: PHPStan\Type\Symfony\MessengerHandleTraitWrapperReturnTypeExtension
+		tags: [phpstan.broker.expressionTypeResolverExtension]
+
 	# InputInterface::getArgument() return type
 	-
 		factory: PHPStan\Type\Symfony\InputInterfaceGetArgumentDynamicReturnTypeExtension

src/Symfony/Configuration.php

@@ -31,4 +31,12 @@ public function getConsoleApplicationLoader(): ?string
 		return $this->parameters['consoleApplicationLoader'] ?? $this->parameters['console_application_loader'] ?? null;
 	}
 
+	/**
+	 * @return array<string>
+	 */
+	public function getMessengerHandleTraitWrappers(): array
+	{
+		return $this->parameters['messenger']['handleTraitWrappers'] ?? [];
+	}
+
 }

src/Type/Symfony/MessengerHandleTraitWrapperReturnTypeExtension.php

@@ -0,0 +1,125 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Type\Symfony;
+
+use PhpParser\Node\Expr;
+use PhpParser\Node\Expr\MethodCall;
+use PhpParser\Node\Identifier;
+use PHPStan\Analyser\Scope;
+use PHPStan\Reflection\ReflectionProvider;
+use PHPStan\Symfony\Configuration;
+use PHPStan\Symfony\MessageMap;
+use PHPStan\Symfony\MessageMapFactory;
+use PHPStan\Type\ExpressionTypeResolverExtension;
+use PHPStan\Type\Type;
+use function count;
+use function in_array;
+use function is_null;
+
+/**
+ * Configurable extension for resolving return types of methods that internally use HandleTrait.
+ *
+ * Configured via PHPStan parameters under symfony.messenger.handleTraitWrappers with
+ * "class::method" patterns, e.g.:
+ * - App\Bus\QueryBus::dispatch
+ * - App\Bus\QueryBus::query
+ * - App\Bus\CommandBus::execute
+ * - App\Bus\CommandBus::handle
+ */
+final class MessengerHandleTraitWrapperReturnTypeExtension implements ExpressionTypeResolverExtension
+{
+
+	/** @var MessageMapFactory */
+	private $messageMapFactory;
+
+	/** @var MessageMap|null */
+	private $messageMap;
+
+	/** @var array<string> */
+	private $wrappers;
+
+	/** @var ReflectionProvider */
+	private $reflectionProvider;
+
+	public function __construct(MessageMapFactory $messageMapFactory, Configuration $configuration, ReflectionProvider $reflectionProvider)
+	{
+		$this->messageMapFactory = $messageMapFactory;
+		$this->wrappers = $configuration->getMessengerHandleTraitWrappers();
+		$this->reflectionProvider = $reflectionProvider;
+	}
+
+	public function getType(Expr $expr, Scope $scope): ?Type
+	{
+		if (!$this->isSupported($expr, $scope)) {
+			return null;
+		}
+
+		$args = $expr->getArgs();
+		if (count($args) !== 1) {
+			return null;
+		}
+
+		$arg = $args[0]->value;
+		$argClassNames = $scope->getType($arg)->getObjectClassNames();
+
+		if (count($argClassNames) === 1) {
+			$messageMap = $this->getMessageMap();
+			$returnType = $messageMap->getTypeForClass($argClassNames[0]);
+
+			if (!is_null($returnType)) {
+				return $returnType;
+			}
+		}
+
+		return null;
+	}
+
+	/**
+	 * @phpstan-assert-if-true =MethodCall $expr
+	 */
+	private function isSupported(Expr $expr, Scope $scope): bool
+	{
+		if (!($expr instanceof MethodCall) || !($expr->name instanceof Identifier)) {
+			return false;
+		}
+
+		$methodName = $expr->name->name;
+		$varType = $scope->getType($expr->var);
+		$classNames = $varType->getObjectClassNames();
+
+		if (count($classNames) !== 1) {
+			return false;
+		}
+
+		$className = $classNames[0];
+		$classMethodCombination = $className . '::' . $methodName;
+
+		// Check if this exact class::method combination is configured
+		if (in_array($classMethodCombination, $this->wrappers, true)) {
+			return true;
+		}
+
+		// Check if any interface implemented by this class::method is configured
+		if ($this->reflectionProvider->hasClass($className)) {
+			$classReflection = $this->reflectionProvider->getClass($className);
+			foreach ($classReflection->getInterfaces() as $interface) {
+				$interfaceMethodCombination = $interface->getName() . '::' . $methodName;
+				if (in_array($interfaceMethodCombination, $this->wrappers, true)) {
+					return true;
+				}
+			}
+		}
+
+		return false;
+	}
+
+	private function getMessageMap(): MessageMap
+	{
+		if ($this->messageMap === null) {
+			$this->messageMap = $this->messageMapFactory->create();
+		}
+
+		return $this->messageMap;
+	}
+
+}

tests/Type/Symfony/data/messenger_handle_trait.php

@@ -111,3 +111,56 @@ public function __invoke()
         assertType('mixed', $this->handle(new MultiHandlersForTheSameMessageQuery()));
     }
 }
+
+class QueryBus {
+    use HandleTrait;
+
+    public function dispatch(object $query): mixed
+    {
+        return $this->handle($query);
+    }
+
+	public function dispatch2(object $query): mixed
+	{
+		return $this->handle($query);
+	}
+}
+
+interface QueryBusInterface {
+	public function dispatch(object $query): mixed;
+}
+
+class QueryBusWithInterface implements QueryBusInterface {
+	use HandleTrait;
+
+	public function dispatch(object $query): mixed
+	{
+		return $this->handle($query);
+	}
+}
+
+class Controller {
+    public function action()
+    {
+        $queryBus = new QueryBus();
+
+        assertType(RegularQueryResult::class, $queryBus->dispatch(new RegularQuery()));
+
+        assertType('bool', $queryBus->dispatch(new BooleanQuery()));
+        assertType('int', $queryBus->dispatch(new IntQuery()));
+        assertType('float', $queryBus->dispatch(new FloatQuery()));
+        assertType('string', $queryBus->dispatch(new StringQuery()));
+
+        assertType(TaggedResult::class, $queryBus->dispatch(new TaggedQuery()));
+
+		assertType(RegularQueryResult::class, $queryBus->dispatch2(new RegularQuery()));
+
+		$queryBusWithInterface = new QueryBusWithInterface();
+
+		assertType(RegularQueryResult::class, $queryBusWithInterface->dispatch(new RegularQuery()));
+
+        // HandleTrait will throw exception in fact due to multiple handle methods/handlers per single query
+        assertType('mixed', $queryBus->dispatch(new MultiHandlesForInTheSameHandlerQuery()));
+        assertType('mixed', $queryBus->dispatch(new MultiHandlersForTheSameMessageQuery()));
+    }
+}

tests/Type/Symfony/extension-test.neon

@@ -2,3 +2,8 @@ parameters:
 	symfony:
 		console_application_loader: console_application_loader.php
 		container_xml_path: container.xml
+		messenger:
+			handleTraitWrappers:
+				- MessengerHandleTrait\QueryBus::dispatch
+				- MessengerHandleTrait\QueryBus::dispatch2
+				- MessengerHandleTrait\QueryBusInterface::dispatch