Add context propagation primitives for async task execution

Introduce ContextPropagatingRunnable/Callable and the ContextPropagation factory that captures the current request Context at wrap time and restores it around task execution, restoring the worker's previous context afterwards. Each task gets its own copy of the context (ContextObject.copyForPropagation) so parallel workers sharing one request never race on its mutable state. Wrapping is idempotent and passes through null / already-wrapped / no-context tasks. Unit-tested in isolation without agent weaving.

Author: Mishenevd

agent_api/src/main/java/dev/aikido/agent_api/context/ContextObject.java

@@ -6,7 +6,7 @@
 
 import java.util.*;
 
-public class ContextObject {
+public class ContextObject implements Cloneable {
     protected String method;
     protected String source;
     protected String url;
@@ -25,6 +25,24 @@ public class ContextObject {
     protected transient Map<String, Map<String, String>> cache = new HashMap<>();
     protected transient Optional<Boolean> forcedProtectionOff = Optional.empty();
 
+    // Async tasks get their own copy so parallel workers sharing one request's
+    // context never race on its mutable working state (cache, redirect nodes).
+    public ContextObject copyForPropagation() {
+        try {
+            ContextObject copy = (ContextObject) super.clone();
+            copy.cache = new HashMap<>(this.cache);
+            if (this.redirectStartNodes != null) {
+                copy.redirectStartNodes = new ArrayList<>(this.redirectStartNodes.size());
+                for (RedirectNode starter : this.redirectStartNodes) {
+                    copy.redirectStartNodes.add(starter.copyChain());
+                }
+            }
+            return copy;
+        } catch (CloneNotSupportedException e) {
+            return this;
+        }
+    }
+
     public boolean middlewareExecuted() {return executedMiddleware; }
     public void setExecutedMiddleware(boolean value) { executedMiddleware = value; }
 

agent_api/src/main/java/dev/aikido/agent_api/context/ContextPropagatingCallable.java

@@ -0,0 +1,29 @@
+package dev.aikido.agent_api.context;
+
+import java.util.concurrent.Callable;
+
+public final class ContextPropagatingCallable<T> implements Callable<T> {
+    private final Callable<T> delegate;
+    private final ContextObject context;
+
+    public ContextPropagatingCallable(Callable<T> delegate, ContextObject context) {
+        this.delegate = delegate;
+        this.context = context;
+    }
+
+    @Override
+    public T call() throws Exception {
+        ContextObject previousContext = Context.get();
+
+        try {
+            Context.set(context);
+            return delegate.call();
+        } finally {
+            if (previousContext != null) {
+                Context.set(previousContext);
+            } else {
+                Context.reset();
+            }
+        }
+    }
+}

agent_api/src/main/java/dev/aikido/agent_api/context/ContextPropagatingRunnable.java

@@ -0,0 +1,27 @@
+package dev.aikido.agent_api.context;
+
+public final class ContextPropagatingRunnable implements Runnable {
+    private final Runnable delegate;
+    private final ContextObject context;
+
+    public ContextPropagatingRunnable(Runnable delegate, ContextObject context) {
+        this.delegate = delegate;
+        this.context = context;
+    }
+
+    @Override
+    public void run() {
+        ContextObject previousContext = Context.get();
+
+        try {
+            Context.set(context);
+            delegate.run();
+        } finally {
+            if (previousContext != null) {
+                Context.set(previousContext);
+            } else {
+                Context.reset();
+            }
+        }
+    }
+}

agent_api/src/main/java/dev/aikido/agent_api/context/ContextPropagation.java

@@ -0,0 +1,33 @@
+package dev.aikido.agent_api.context;
+
+import java.util.concurrent.Callable;
+
+public final class ContextPropagation {
+    private ContextPropagation() {}
+
+    public static Runnable wrap(Runnable task) {
+        if (task == null || task instanceof ContextPropagatingRunnable) {
+            return task;
+        }
+
+        ContextObject context = Context.get();
+        if (context == null) {
+            return task;
+        }
+
+        return new ContextPropagatingRunnable(task, context.copyForPropagation());
+    }
+
+    public static <T> Callable<T> wrap(Callable<T> task) {
+        if (task == null || task instanceof ContextPropagatingCallable) {
+            return task;
+        }
+
+        ContextObject context = Context.get();
+        if (context == null) {
+            return task;
+        }
+
+        return new ContextPropagatingCallable<>(task, context.copyForPropagation());
+    }
+}

agent_api/src/main/java/dev/aikido/agent_api/storage/RedirectNode.java

@@ -36,6 +36,19 @@ public void setChild(RedirectNode child) {
         this.child = child;
     }
 
+    // Deep-copies the chain from this node downward so a propagated context gets its own
+    // RedirectNode instances instead of sharing mutable ones (see ContextObject.copyForPropagation).
+    public RedirectNode copyChain() {
+        RedirectNode copy = new RedirectNode(this.url);
+        RedirectNode source = this.child;
+        RedirectNode tail = copy;
+        while (source != null) {
+            tail = new RedirectNode(tail, source.url);
+            source = source.child;
+        }
+        return copy;
+    }
+
     @Override
     public boolean equals(Object obj) {
         if (obj == null || getClass() != obj.getClass()) {