Skip to content

Context propagation primitives for async task execution - #342

Open
Mishenevd wants to merge 1 commit into
mainfrom
feat/context-propagation-primitives
Open

Context propagation primitives for async task execution#342
Mishenevd wants to merge 1 commit into
mainfrom
feat/context-propagation-primitives

Conversation

@Mishenevd

@Mishenevd Mishenevd commented Aug 17, 2026

Copy link
Copy Markdown
Collaborator

When a request is handled, its context lives in a ThreadLocal. The moment work is handed to another thread that's gone, so any detector (SQLi, SSRF, …) running there sees no context. The building block for fixing that is a small wrapper that snapshots the current context when a task is created and restores it around the task's run.

ContextPropagation.wrap(...) takes a Runnable/Callable and returns a wrapper that:

  • captures the current Context at wrap time,
  • makes it current while the task runs, and
  • restores whatever the worker had before, even if the task throws.

It stays a no-op where it should: null, already-wrapped, and no-active-context tasks are returned untouched, so wrapping twice is harmless.

Each wrapped task gets its own copy of the context (ContextObject.copyForPropagation), not a shared reference. So when one request fans out several tasks in parallel, the workers don't race on the request's mutable working state (the extracted-input cache, redirect-tracking nodes). This mirrors how OpenTelemetry keeps its context immutable across the async boundary — same safety, done as a per-task snapshot since our context is mutable.

No executor is touched here — that's the next PR. Covered by ContextPropagationTest (capture, restore, restore-previous, clear-when-none, exception safety, idempotency, delegate execution, and snapshot isolation).

@codecov

codecov Bot commented Aug 17, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 96.82540% with 2 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
...va/dev/aikido/agent_api/context/ContextObject.java 91.30% 2 Missing ⚠️

📢 Thoughts on this report? Let us know!

Comment on lines +9 to +10
if (task == null || task instanceof ContextPropagatingRunnable) {
return task;

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Medium - Re-wrapping a retained task keeps the first request's context

ContextPropagation.wrap(...) returns an existing ContextPropagatingRunnable/Callable unchanged instead of rebinding it to the current request. If an application keeps a task instance and calls wrap() before submitting it for later requests, the wrapper continues to install the original ContextObject, so later async work inherits stale route, user, IP, and forcedProtectionOff state. That can misattribute attack reports to the wrong request and, when the first request had protection forced off, skip vulnerability scanning for subsequent requests.

Show fix

Do not treat already-wrapped tasks as safe to reuse across requests. Either always create a fresh wrapper for each wrap() call, or unwrap and rebind the delegate when the current ContextObject differs from the one captured previously so each submission propagates the caller's current request context.

More info - Reply on this comment to give feedback or ignore the issue.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

the wrapper is created inside the agent at submit time and never handed back to the app, so a task is only "already wrapped" within the same submit (same request/context). An app can't retain and resubmit a wrapped task across requests, so the stale-context path is unreachable.

@Mishenevd
Mishenevd force-pushed the feat/context-propagation-primitives branch from c0d0a2a to 75593a3 Compare August 17, 2026 10:54
@Mishenevd
Mishenevd force-pushed the feat/context-propagation-primitives branch 2 times, most recently from 0f414fa to 6f88295 Compare August 17, 2026 16:33
Comment thread agent_api/src/main/java/dev/aikido/agent_api/context/ContextObject.java Outdated
@Mishenevd
Mishenevd force-pushed the feat/context-propagation-primitives branch 3 times, most recently from bceeaf3 to 776ac96 Compare August 17, 2026 23:54
@AikidoSec AikidoSec deleted a comment from aikido-pr-checks Bot Aug 18, 2026
@AikidoSec AikidoSec deleted a comment from aikido-pr-checks Bot Aug 18, 2026
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.
@Mishenevd
Mishenevd force-pushed the feat/context-propagation-primitives branch from 776ac96 to 1fdff6b Compare August 18, 2026 10:22
@AikidoSec AikidoSec deleted a comment from aikido-pr-checks Bot Aug 18, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant