Ability to find an envelope instance from the given stack

src/Message/EnvelopeInterface.php

@@ -17,4 +17,17 @@ public static function fromMessage(MessageInterface $message): self;
     public function getMessage(): MessageInterface;
 
     public function withMessage(MessageInterface $message): self;
+
+    /**
+     * Finds an envelope in the current envelope stack or creates a new one from the message.
+     *
+     * @template T of EnvelopeInterface
+     *
+     * @psalm-param class-string<T> $className
+     * @throws NotEnvelopInterfaceException Implementation MUST throw this exception if the given class does not
+     *         implement {@see EnvelopeInterface}.
+     *
+     * @psalm-return T
+     */
+    public function getEnvelopeFromStack(string $className): self;
 }

src/Message/EnvelopeTrait.php

@@ -45,6 +45,42 @@
         );
     }
 
+    /**
+     * Finds an envelope in the current envelope stack or creates a new one from the message.
+     *
+     * @template T of EnvelopeInterface
+     *
+     * @psalm-param class-string<T> $className
+     * @throws NotEnvelopInterfaceException is thrown if the given class does not implement {@see EnvelopeInterface}.
+     *
+     * @psalm-return T
+     */
+    public function getEnvelopeFromStack(string $className): EnvelopeInterface
+    {
+        if (!is_a($className, EnvelopeInterface::class, true)) {
+            throw new NotEnvelopInterfaceException($className);
+        }
+
+        if (static::class === $className) {
+            return $this;
+        }
+
+        if ($this->message instanceof EnvelopeInterface) {
+            return $this->message->getEnvelopeFromStack($className);
+        }
+
+        return $className::fromMessage($this->message);
+    }
+
+    public static function getEnvelopeFromMessage(MessageInterface $message): EnvelopeInterface
+    {
+        if ($message instanceof EnvelopeInterface) {
+            return $message->getEnvelopeFromStack(static::class);
+        }
+
+        return static::fromMessage($message);
+    }
+
     public function getEnvelopeMetadata(): array
     {
         return [];

src/Message/NotEnvelopInterfaceException.php

@@ -0,0 +1,22 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Yiisoft\Queue\Message;
+
+use InvalidArgumentException;
+use Throwable;
+
+final class NotEnvelopInterfaceException extends InvalidArgumentException
+{
+    public function __construct(string $className = '', int $code = 0, ?Throwable $previous = null)
+    {
+        $message = sprintf(
+            'The given class "%s" does not implement "%s".',
+            $className,
+            EnvelopeInterface::class
+        );
+
+        parent::__construct($message, $code, $previous);
+    }
+}

tests/Unit/Message/EnvelopeTraitTest.php

@@ -0,0 +1,65 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Yiisoft\Queue\Tests\Unit\Message;
+
+use PHPUnit\Framework\TestCase;
+use Yiisoft\Queue\Message\EnvelopeInterface;
+use Yiisoft\Queue\Message\IdEnvelope;
+use Yiisoft\Queue\Message\Message;
+use Yiisoft\Queue\Message\NotEnvelopInterfaceException;
+use Yiisoft\Queue\Middleware\FailureHandling\FailureEnvelope;
+
+class EnvelopeTraitTest extends TestCase
+{
+    public function testGetEnvelopeFromStack(): void
+    {
+        $message = new Message('handler', 'data');
+        $envelope = new FailureEnvelope(new IdEnvelope($message, 'id value'), ['fail' => true]);
+        $result = $envelope->getEnvelopeFromStack(IdEnvelope::class);
+
+        $this->assertInstanceOf(IdEnvelope::class, $result);
+        $this->assertEquals('id value', $result->getId());
+    }
+
+    public function testGetEnvelopeFromStackNotInStack(): void
+    {
+        $message = new Message('handler', 'data', [IdEnvelope::MESSAGE_ID_KEY => 'id value']);
+        $envelope = new FailureEnvelope($message, ['fail' => true]);
+        $result = $envelope->getEnvelopeFromStack(IdEnvelope::class);
+
+        $this->assertInstanceOf(IdEnvelope::class, $result);
+        $this->assertEquals('id value', $result->getId());
+    }
+
+    public function testGetEnvelopeFromStackNotEnvelope(): void
+    {
+        $this->expectException(NotEnvelopInterfaceException::class);
+        $interface = EnvelopeInterface::class;
+        $this->expectExceptionMessage("The given class \"foo\" does not implement \"$interface\".");
+
+        $message = new Message('handler', 'data', [IdEnvelope::MESSAGE_ID_KEY => 'id value']);
+        $envelope = new FailureEnvelope($message, ['fail' => true]);
+        $envelope->getEnvelopeFromStack('foo');
+    }
+
+    public function testGetEnvelopeFromMessage(): void
+    {
+        $message = new Message('handler', 'data');
+        $envelope = new FailureEnvelope(new IdEnvelope($message, 'id value'), ['fail' => true]);
+        $result = IdEnvelope::getEnvelopeFromMessage($envelope);
+
+        $this->assertInstanceOf(IdEnvelope::class, $result);
+        $this->assertEquals('id value', $result->getId());
+    }
+
+    public function testGetEnvelopeFromMessageNotInStack(): void
+    {
+        $message = new Message('handler', 'data', [IdEnvelope::MESSAGE_ID_KEY => 'id value']);
+        $result = IdEnvelope::getEnvelopeFromMessage($message);
+
+        $this->assertInstanceOf(IdEnvelope::class, $result);
+        $this->assertEquals('id value', $result->getId());
+    }
+}