package io.quarkus.arc.impl;

import io.quarkus.arc.Arc;
import io.quarkus.arc.ManagedContext;
import io.smallrye.mutiny.Uni;
import javax.annotation.Priority;
import javax.enterprise.context.control.ActivateRequestContext;
import javax.interceptor.AroundInvoke;
import javax.interceptor.Interceptor;
import javax.interceptor.InvocationContext;

@Interceptor
@ActivateRequestContext
@Priority(Interceptor.Priority.PLATFORM_BEFORE + 100)
public class ActivateRequestContextInterceptor {

    @AroundInvoke
    Object aroundInvoke(InvocationContext ctx) throws Exception {
        if (ctx.getMethod().getReturnType().equals(Uni.class)) {
            return invokeUni(ctx);
        }

        return invoke(ctx);
    }

    private Uni<?> invokeUni(InvocationContext ctx) {
        return Uni.createFrom().deferred(() -> {
            ManagedContext requestContext = Arc.container().requestContext();
            if (requestContext.isActive()) {
                return proceedWithUni(ctx);
            }

            return Uni.createFrom().deferred(() -> {
                requestContext.activate();
                return proceedWithUni(ctx);
            }).eventually(requestContext::terminate);
        });
    }

    private Uni<?> proceedWithUni(InvocationContext ctx) {
        try {
            return (Uni<?>) ctx.proceed();
        } catch (Throwable t) {
            return Uni.createFrom().failure(t);
        }
    }

    private Object invoke(InvocationContext ctx) throws Exception {
        ManagedContext requestContext = Arc.container().requestContext();
        if (requestContext.isActive()) {
            return ctx.proceed();
        }

        try {
            requestContext.activate();
            return ctx.proceed();
        } finally {
            requestContext.terminate();
        }
    }
}

@@ -119,16 +119,108 @@
 ny.Uni;%0A
+import java.util.concurrent.CompletableFuture;%0Aimport java.util.concurrent.CompletionStage;%0A
 import j
@@ -245,16 +245,16 @@
 iority;%0A
-
 import j
@@ -779,27 +779,1122 @@
     
-return invoke(ctx);
+if (ctx.getMethod().getReturnType().equals(CompletionStage.class)) %7B%0A            return invokeStage(ctx);%0A        %7D%0A%0A        return invoke(ctx);%0A    %7D%0A%0A    private CompletionStage%3C?%3E invokeStage(InvocationContext ctx) %7B%0A        ManagedContext requestContext = Arc.container().requestContext();%0A        if (requestContext.isActive()) %7B%0A            return proceedWithStage(ctx);%0A        %7D%0A%0A        return activate(requestContext)%0A                .thenCompose(v -%3E proceedWithStage(ctx))%0A                .whenComplete((r, t) -%3E requestContext.terminate());%0A    %7D%0A%0A    private static CompletionStage%3CManagedContext%3E activate(ManagedContext requestContext) %7B%0A        try %7B%0A            requestContext.activate();%0A            return CompletableFuture.completedStage(requestContext);%0A        %7D catch (Throwable t) %7B%0A            return CompletableFuture.failedStage(t);%0A        %7D%0A    %7D%0A%0A    private CompletionStage%3C?%3E proceedWithStage(InvocationContext ctx) %7B%0A        try %7B%0A            return (CompletionStage%3C?%3E) ctx.proceed();%0A        %7D catch (Throwable t) %7B%0A            return CompletableFuture.failedStage(t);%0A        %7D
 %0A   

Add support for @ActivateRequestContext with CompletionStage