From e1d4296990ef8ee2958de37f8a27896cb85221ec Mon Sep 17 00:00:00 2001
From: Guillaume Bouchard <guillaum.bouchard@gmail.com>
Date: Sat, 10 May 2025 10:28:07 +0400
Subject: [PATCH 1/8] feat: support for exception context propagation

We specialize the `throwIO` call using a newly implemented `rethrowIO'`
which behaves as `rethrowIO` from base 4.21 when available or like the
previous `throw` implementation.

In short:

- Before `base-4.21`, the code is exactly as before
- After `base-4.21`, the code does not override the backtrace
  annotations and instead uses `rethrowIO`.

Example of usage / changes:

The following code:

```haskell
{-# LANGUAGE DeriveAnyClass #-}
import Control.Concurrent.Async
import Control.Exception
import Control.Exception.Context
import Control.Exception.Annotation
import Data.Typeable
import Data.Traversable
import GHC.Stack

data Ann = Ann String
  deriving (Show, ExceptionAnnotation)

asyncTask :: HasCallStack => IO ()
asyncTask = annotateIO (Ann "bonjour") $ do
  error "yoto"

asyncTask' :: HasCallStack => IO ()
asyncTask' = annotateIO (Ann "bonjour2") $ do
  error "yutu"

main = do
  -- withAsync asyncTask wait
  concurrently asyncTask asyncTask'
  -- race asyncTask asyncTask'
```

When run without this commit leads to:

```
ASyncException.hs: Uncaught exception ghc-internal:GHC.Internal.Exception.ErrorCall:

yoto

HasCallStack backtrace:
  throwIO, called at ./Control/Concurrent/Async/Internal.hs:630:24 in async-2.2.5-50rpfAJ7BEc1o5OswtTMUN:Control.Concurrent.Async.Internal
```

When run with this commit:

```
*** Exception: yoto

Ann "bonjour"
HasCallStack backtrace:
  error, called at /home/guillaume//ASyncException.hs:15:3 in async-2.2.5-inplace:Main
  asyncTask, called at /home/guillaume//ASyncException.hs:23:16 in async-2.2.5-inplace:Main
```
---
 Control/Concurrent/Async/Internal.hs | 40 +++++++++++++++++++++++-----
 1 file changed, 34 insertions(+), 6 deletions(-)

diff --git a/Control/Concurrent/Async/Internal.hs b/Control/Concurrent/Async/Internal.hs
index 00abf7f..499ad34 100644
--- a/Control/Concurrent/Async/Internal.hs
+++ b/Control/Concurrent/Async/Internal.hs
@@ -7,6 +7,7 @@
 {-# LANGUAGE DeriveDataTypeable #-}
 #endif
 {-# OPTIONS -Wall #-}
+{-# LANGUAGE ScopedTypeVariables #-}
 
 -----------------------------------------------------------------------------
 -- |
@@ -219,10 +220,24 @@ withAsyncUsing doFork action inner = do
     let a = Async t (readTMVar var)
     r <- restore (inner a) `catchAll` \e -> do
       uninterruptibleCancel a
-      throwIO e
+      rethrowIO' e
     uninterruptibleCancel a
     return r
 
+
+-- | This function attempts at rethrowing while keeping the context
+-- This is internal and only working with GHC >=9.12, otherwise it fallsback to
+-- standard 'throwIO'
+rethrowIO' :: SomeException -> IO a
+#if MIN_VERSION_base(4,21,0)
+rethrowIO' e = 
+  case fromException e of
+    Just (e' :: ExceptionWithContext SomeException) -> rethrowIO e'
+    Nothing -> throwIO e
+#else
+rethrowIO' = throwIO
+#endif
+
 -- | Wait for an asynchronous action to complete, and return its
 -- value.  If the asynchronous action threw an exception, then the
 -- exception is re-thrown by 'wait'.
@@ -265,7 +280,20 @@ poll = atomically . pollSTM
 waitSTM :: Async a -> STM a
 waitSTM a = do
    r <- waitCatchSTM a
-   either throwSTM return r
+   either (rethrowSTM) return r
+
+-- | This function attempts at rethrowing while keeping the context
+-- This is internal and only working with GHC >=9.12, otherwise it fallsback to
+-- standard 'throwSTM'
+rethrowSTM :: SomeException -> STM a
+#if MIN_VERSION_base(4,21,0)
+rethrowSTM e = 
+  case fromException e of
+    Just (e' :: ExceptionWithContext SomeException) -> throwSTM (NoBacktrace e')
+    Nothing -> throwSTM e
+#else
+rethrowSTM = throwSTM
+#endif
 
 -- | A version of 'waitCatch' that can be used inside an STM transaction.
 --
@@ -664,7 +692,7 @@ race left right = concurrently' left right collect
     collect m = do
         e <- m
         case e of
-            Left ex -> throwIO ex
+            Left ex -> rethrowIO' ex
             Right r -> return r
 
 -- race_ :: IO a -> IO b -> IO ()
@@ -678,7 +706,7 @@ concurrently left right = concurrently' left right (collect [])
     collect xs m = do
         e <- m
         case e of
-            Left ex -> throwIO ex
+            Left ex -> rethrowIO' ex
             Right r -> collect (r:xs) m
 
 -- concurrentlyE :: IO (Either e a) -> IO (Either e b) -> IO (Either e (a, b))
@@ -691,7 +719,7 @@ concurrentlyE left right = concurrently' left right (collect [])
     collect xs m = do
         e <- m
         case e of
-            Left ex -> throwIO ex
+            Left ex -> rethrowIO' ex
             Right r -> collect (r:xs) m
 
 concurrently' ::
@@ -752,7 +780,7 @@ concurrently_ left right = concurrently' left right (collect 0)
     collect i m = do
         e <- m
         case e of
-            Left ex -> throwIO ex
+            Left ex -> rethrowIO' ex
             Right _ -> collect (i + 1 :: Int) m
 
 

From 6dad253853cb691886eced6ddb2da77dcd033a8a Mon Sep 17 00:00:00 2001
From: Guillaume Bouchard <guillaum.bouchard@gmail.com>
Date: Sun, 18 Jan 2026 15:43:00 +0400
Subject: [PATCH 2/8] feat: add exceptions rethrow tests

---
 test/test-async.hs | 89 +++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 88 insertions(+), 1 deletion(-)

diff --git a/test/test-async.hs b/test/test-async.hs
index c0d15e0..8c35820 100644
--- a/test/test-async.hs
+++ b/test/test-async.hs
@@ -1,4 +1,5 @@
 {-# LANGUAGE CPP,ScopedTypeVariables,DeriveDataTypeable #-}
+{-# LANGUAGE DeriveAnyClass #-}
 module Main where
 
 import Test.Framework (defaultMain, testGroup)
@@ -21,6 +22,12 @@ import Data.Foldable (foldMap)
 import Data.Maybe
 
 import Prelude hiding (catch)
+#if MIN_VERSION_base(4,21,0)
+import Control.Exception.Annotation (ExceptionAnnotation(..))
+import Control.Exception.Context (displayExceptionContext, getExceptionAnnotations)
+import Control.Exception.Backtrace (Backtraces, displayBacktraces)
+#endif
+import GHC.Stack (HasCallStack)
 
 main = defaultMain tests
 
@@ -78,6 +85,10 @@ tests = [
       case_mapConcurrentlyBounded_exception
   , testCase "Warden" case_Warden
   , testCase "Warden_spawn_after_shutdown" case_Warden_spawn_after_shutdown
+
+#if MIN_VERSION_base(4,21,0)
+  , testGroup "exception rethrow" exception_rethrow
+#endif
   ]
  ]
 
@@ -547,4 +558,80 @@ case_Warden_spawn_after_shutdown = do
   r <- try $ spawn warden $ return ()
   case r of
     Left (WardenException{}) -> return ()  -- expected
-    Right _ -> assertFailure "Expected WardenException"
\ No newline at end of file
+    Right _ -> assertFailure "Expected WardenException"
+
+#if MIN_VERSION_base(4,21,0)
+-- The following regroups tests of exception context propagation to ensure that
+-- exception rethrown by async keep the initial backtrace.
+
+-- | This is a dummy exception that we can throw
+data Exc = Exc
+  deriving (Show, Exception)
+
+action_wrapper :: HasCallStack => (IO x -> IO y) -> IO y
+action_wrapper op = op action
+
+action :: HasCallStack => IO x
+action = throwIO Exc
+
+
+-- | From an exception, extract two lines of context, ignoring the header and
+-- the remaining lines.
+--
+-- For example, when calling the above 'action_wrapper (\x -> x)', in GHC 9.12, the current callstack looks like:
+--
+--
+-- HasCallStack backtrace:
+--   throwIO, called at test/test-async.hs:485:11 in async-2.2.5-inplace-test-async:Main
+--   action, called at test/test-async.hs:482:10 in async-2.2.5-inplace-test-async:Main
+--   action_wrapper, called at <interactive>:2:1 in interactive:Ghci1
+--
+--   We drop the header (e.g. HasCallStack backtrace:) and only keep the two
+--   lines showing the callstack inside "action".
+--
+--   Note that it does not show where action_wrapper was called, but the idea
+--   is that action_wrapper will do the call to the async primitive (e.g.
+--   'concurrently') and will hence keep the trace of where 'concurrently' was
+--   called.
+extractThrowOrigin :: ExceptionWithContext Exc -> [String]
+extractThrowOrigin (ExceptionWithContext ctx e) =  do
+      let backtraces :: [Backtraces] = getExceptionAnnotations ctx
+      case backtraces of
+        [backtrace] -> take 2 $ drop 1 (lines (displayBacktraces backtrace))
+        _ -> error "more than one backtrace"
+
+-- | Run 'action' through a wrapper (using 'action_wrapper') and with a naive
+-- wrapper and show that the wrapper returns the same callstack when the
+-- exception in 'action' is raised.
+compareTwoExceptions op = do
+  Left direct_exception <- tryWithContext (action_wrapper (\x -> x))
+  let direct_origin = extractThrowOrigin direct_exception
+
+  Left indirect_exception <- tryWithContext (action_wrapper op)
+  let indirect_origin = extractThrowOrigin indirect_exception
+
+  assertEqual "The exception origins" direct_origin indirect_origin
+
+doNothing = pure ()
+doForever = doForever
+
+exception_rethrow = [
+    testCase "concurrentlyL" $ compareTwoExceptions (\action -> concurrently action doNothing),
+    testCase "concurrentlyR" $ compareTwoExceptions (\action -> concurrently doNothing action),
+    testCase "concurrently_L" $ compareTwoExceptions (\action -> concurrently_ action doNothing),
+    testCase "concurrently_R" $ compareTwoExceptions (\action -> concurrently_ doNothing action),
+    testCase "raceL" $ compareTwoExceptions (\action -> race action doForever),
+    testCase "raceR" $ compareTwoExceptions (\action -> race doForever action),
+    testCase "race_L" $ compareTwoExceptions (\action -> race_ action doForever),
+    testCase "race_R" $ compareTwoExceptions (\action -> race_ doForever action),
+    testCase "mapConcurrently" $ compareTwoExceptions (\action -> mapConcurrently (\() -> action) [(), (), ()]),
+    testCase "forConcurrently" $ compareTwoExceptions (\action -> forConcurrently [(), (), ()] (\() -> action)),
+    testCase "withAsync wait" $ compareTwoExceptions $ \action -> do
+      withAsync action $ \a -> do
+        wait a,
+    testCase "withAsync inside" $ compareTwoExceptions $ \action -> do
+      withAsync doForever $ \a -> do
+        action
+    ]
+#endif
+

From fb9bfd35bdcc2bb8cf7a494c657b897f0ff4466d Mon Sep 17 00:00:00 2001
From: Guillaume Bouchard <guillaum.bouchard@gmail.com>
Date: Sat, 31 Jan 2026 16:35:00 +0400
Subject: [PATCH 3/8] feat: `withAsync` does not wrap exception in the inner
 block with WhileWaiting

---
 Control/Concurrent/Async/Internal.hs | 23 ++++++++++++++++++++++-
 1 file changed, 22 insertions(+), 1 deletion(-)

diff --git a/Control/Concurrent/Async/Internal.hs b/Control/Concurrent/Async/Internal.hs
index 499ad34..717c006 100644
--- a/Control/Concurrent/Async/Internal.hs
+++ b/Control/Concurrent/Async/Internal.hs
@@ -207,6 +207,26 @@ withAsyncOnWithUnmask ::
 withAsyncOnWithUnmask cpu actionWith =
   withAsyncUsing (rawForkOn cpu) (actionWith unsafeUnmask)
 
+#if MIN_VERSION_base(4,21,0)
+withAsyncUsing ::
+  CALLSTACK
+  (IO () -> IO ThreadId) -> IO a -> (Async a -> IO b) -> IO b
+-- The bracket version works, but is slow.  We can do better by
+-- hand-coding it:
+withAsyncUsing doFork action inner = do
+  var <- newEmptyTMVarIO
+  mask $ \restore -> do
+    let action_plus = debugLabelMe >> action
+    t <- doFork $ try (restore action_plus) >>= atomically . putTMVar var
+    let a = Async t (readTMVar var)
+    -- Using catch/no/propagate and rethrowIO, we do not wrap the exception
+    -- with a `WhileWaiting`
+    r <- restore (inner a) `catchNoPropagate` \e -> do
+      uninterruptibleCancel a
+      rethrowIO (e :: ExceptionWithContext SomeException)
+    uninterruptibleCancel a
+    return r
+#else
 withAsyncUsing ::
   CALLSTACK
   (IO () -> IO ThreadId) -> IO a -> (Async a -> IO b) -> IO b
@@ -220,9 +240,10 @@ withAsyncUsing doFork action inner = do
     let a = Async t (readTMVar var)
     r <- restore (inner a) `catchAll` \e -> do
       uninterruptibleCancel a
-      rethrowIO' e
+      throwIO e
     uninterruptibleCancel a
     return r
+#endif
 
 
 -- | This function attempts at rethrowing while keeping the context

From e84ad663aa85318f9d0e7242ef83d82e7f1550d9 Mon Sep 17 00:00:00 2001
From: Guillaume Bouchard <guillaum.bouchard@gmail.com>
Date: Sat, 31 Jan 2026 16:44:43 +0400
Subject: [PATCH 4/8] feat: annotate 'waitXXX' function with the callsite for
 'wait'

Previous commits extended the exception annotations so the default
callstack represents the exception location of the original exception in
the async process. This callstack also includes where the async process
was started (e.g. in `withAsync`).

This commits extends the exception context by adding a new
`AsyncWaitLocation` exception annotation which contains the location of
the `wait` call.

Note the usage of `withFrozenCallStack` in order to not expose the
internal of the async library in the callstack.
---
 Control/Concurrent/Async/Internal.hs | 42 ++++++++++++++++++++++++----
 1 file changed, 36 insertions(+), 6 deletions(-)

diff --git a/Control/Concurrent/Async/Internal.hs b/Control/Concurrent/Async/Internal.hs
index 717c006..bd1694a 100644
--- a/Control/Concurrent/Async/Internal.hs
+++ b/Control/Concurrent/Async/Internal.hs
@@ -57,6 +57,12 @@ import Data.IORef
 import GHC.Exts
 import GHC.IO hiding (finally, onException)
 import GHC.Conc (ThreadId(..), labelThread)
+import GHC.Stack (CallStack, callStack, prettyCallStack, withFrozenCallStack)
+
+#if MIN_VERSION_base(4,21,0)
+import Control.Exception.Annotation (ExceptionAnnotation (..))
+#endif
+import GHC.Stack.Types (HasCallStack)
 
 #if defined(__MHS__)
 import Data.Traversable
@@ -259,6 +265,30 @@ rethrowIO' e =
 rethrowIO' = throwIO
 #endif
 
+#if MIN_VERSION_base(4,21,0)
+-- | An exception annotation which stores the callstack of a 'wait',
+-- 'waitBoth', 'waitEither' call.
+data AsyncWaitLocation = AsyncWaitLocation CallStack
+  deriving (Show)
+
+instance ExceptionAnnotation AsyncWaitLocation where
+  displayExceptionAnnotation (AsyncWaitLocation callstack) = "AsyncWaitLocation " <> prettyCallStack callstack
+
+-- | Annotate an exception with the current callstack with GHC >= 9.12
+annotateWithCallSite :: HasCallStack => IO b -> IO b
+annotateWithCallSite action = do
+  resM <- tryWithContext action
+  case resM of
+    Right res -> pure res
+    Left (exc :: ExceptionWithContext SomeException) -> do
+      annotateIO (AsyncWaitLocation callStack) $ rethrowIO exc
+#else
+-- | Do nothing with GHC < 9.12
+annotateWithCallSite :: HasCallStack => IO b -> IO b
+annotateWithCallSite action = action
+#endif
+
+
 -- | Wait for an asynchronous action to complete, and return its
 -- value.  If the asynchronous action threw an exception, then the
 -- exception is re-thrown by 'wait'.
@@ -266,8 +296,8 @@ rethrowIO' = throwIO
 -- > wait = atomically . waitSTM
 --
 {-# INLINE wait #-}
-wait :: Async a -> IO a
-wait = tryAgain . atomically . waitSTM
+wait :: HasCallStack => Async a -> IO a
+wait = withFrozenCallStack $ annotateWithCallSite . tryAgain . atomically . waitSTM
   where
     -- See: https://github.com/simonmar/async/issues/14
     tryAgain f = f `catch` \BlockedIndefinitelyOnSTM -> f
@@ -485,8 +515,8 @@ waitEitherCatchCancel left right =
 -- re-thrown by 'waitEither'.
 --
 {-# INLINE waitEither #-}
-waitEither :: Async a -> Async b -> IO (Either a b)
-waitEither left right = atomically (waitEitherSTM left right)
+waitEither :: HasCallStack => Async a -> Async b -> IO (Either a b)
+waitEither left right = withFrozenCallStack $ annotateWithCallSite $ atomically (waitEitherSTM left right)
 
 -- | A version of 'waitEither' that can be used inside an STM transaction.
 --
@@ -524,8 +554,8 @@ waitEitherCancel left right =
 -- re-thrown by 'waitBoth'.
 --
 {-# INLINE waitBoth #-}
-waitBoth :: Async a -> Async b -> IO (a,b)
-waitBoth left right = tryAgain $ atomically (waitBothSTM left right)
+waitBoth :: HasCallStack => Async a -> Async b -> IO (a,b)
+waitBoth left right = withFrozenCallStack $ annotateWithCallSite $ tryAgain $ atomically (waitBothSTM left right)
   where
     -- See: https://github.com/simonmar/async/issues/14
     tryAgain f = f `catch` \BlockedIndefinitelyOnSTM -> f

From 91a23d50dc6c3f3d2690f01d934afb5c514c66b1 Mon Sep 17 00:00:00 2001
From: Guillaume Bouchard <guillaum.bouchard@gmail.com>
Date: Sat, 31 Jan 2026 18:19:10 +0400
Subject: [PATCH 5/8] feat: add tests for location of `wait` call annotation

---
 Control/Concurrent/Async.hs          |  3 ++
 Control/Concurrent/Async/Internal.hs |  2 +-
 test/test-async.hs                   | 41 +++++++++++++++++++++++++---
 3 files changed, 41 insertions(+), 5 deletions(-)

diff --git a/Control/Concurrent/Async.hs b/Control/Concurrent/Async.hs
index 9f99061..0781fda 100644
--- a/Control/Concurrent/Async.hs
+++ b/Control/Concurrent/Async.hs
@@ -185,6 +185,9 @@ module Control.Concurrent.Async (
     -- ** Linking
     link, linkOnly, link2, link2Only, ExceptionInLinkedThread(..),
 
+   -- ** Exception annotations
+   AsyncWaitLocation(..)
+
   ) where
 
 import Control.Concurrent.Async.Internal
diff --git a/Control/Concurrent/Async/Internal.hs b/Control/Concurrent/Async/Internal.hs
index bd1694a..ce12e34 100644
--- a/Control/Concurrent/Async/Internal.hs
+++ b/Control/Concurrent/Async/Internal.hs
@@ -265,12 +265,12 @@ rethrowIO' e =
 rethrowIO' = throwIO
 #endif
 
-#if MIN_VERSION_base(4,21,0)
 -- | An exception annotation which stores the callstack of a 'wait',
 -- 'waitBoth', 'waitEither' call.
 data AsyncWaitLocation = AsyncWaitLocation CallStack
   deriving (Show)
 
+#if MIN_VERSION_base(4,21,0)
 instance ExceptionAnnotation AsyncWaitLocation where
   displayExceptionAnnotation (AsyncWaitLocation callstack) = "AsyncWaitLocation " <> prettyCallStack callstack
 
diff --git a/test/test-async.hs b/test/test-async.hs
index 8c35820..8de4f55 100644
--- a/test/test-async.hs
+++ b/test/test-async.hs
@@ -23,11 +23,12 @@ import Data.Maybe
 
 import Prelude hiding (catch)
 #if MIN_VERSION_base(4,21,0)
-import Control.Exception.Annotation (ExceptionAnnotation(..))
-import Control.Exception.Context (displayExceptionContext, getExceptionAnnotations)
-import Control.Exception.Backtrace (Backtraces, displayBacktraces)
+import Control.Exception.Annotation
+import Control.Exception.Context
+import Control.Exception.Backtrace
 #endif
 import GHC.Stack (HasCallStack)
+import Debug.Trace
 
 main = defaultMain tests
 
@@ -631,7 +632,39 @@ exception_rethrow = [
         wait a,
     testCase "withAsync inside" $ compareTwoExceptions $ \action -> do
       withAsync doForever $ \a -> do
-        action
+        action,
+
+    testCase "withAsync does not wrap with WhileHandling and contain an asyncWaitLocation" $ do
+      -- This test is fragile. It checks that when calling `wait` on an async,
+      -- we end up with at least two interesting annotations: the backtrace
+      -- which shows the localisation of the exception thrown in the async, and
+      -- an AsyncWaitLocation which shows the location of the wait call.
+      --
+      -- It also checks that no other annotation are provided (for example, a
+      -- "WhileHandling")
+      --
+      -- However, this can change in future GHC version, for example, new
+      -- annotations may be added, or Backtraces may change its type / name.
+      -- Also, depending on the build configuration, maybe there will have
+      -- other backtraces (such as DWARF or IPE, ...)
+      e <- tryWithContext $ do
+        withAsync (throwIO Exc) $ \async -> do
+           wait async
+      case e of
+        Right () -> fail "should have raised an exception"
+        Left (ExceptionWithContext (ExceptionContext annotations) Exc) -> do
+          assertEqual "Only two annotations" (length annotations) 2
+          assertBool "Has AsyncWaitLocation annotation" (any isAsyncWaitLocation annotations)
+          assertBool "Has Backtraces annotation" (any isBacktraces annotations)
     ]
+
+isAsyncWaitLocation (SomeExceptionAnnotation ann) = case cast ann of
+   Just (AsyncWaitLocation _) -> True
+   _ -> traceShow (typeOf ann) False
+
+isBacktraces (SomeExceptionAnnotation ann) = case cast ann of
+   Just (_ :: Backtraces) -> True
+   _ -> False
+
 #endif
 

From 3f7e6de6520bed29bac39e83bfe2f441f6fe1907 Mon Sep 17 00:00:00 2001
From: Guillaume Bouchard <guillaum.bouchard@gmail.com>
Date: Sat, 28 Mar 2026 11:50:07 +0400
Subject: [PATCH 6/8] refactor: alternate implementation for exception context

This implementation works differently than the previous one. It actually
implements all the standard function based on the GHC 9.12 api (e.g.
`tryWithContext`, `catchNoPropagate`, `rethrowIO`) and provides a simple
and localise compatibility layer which reimplements these function for
older `base`.

The implementation should be easier to read with less CPP and subtle
logic scattered everywhere.
---
 Control/Concurrent/Async/Internal.hs | 128 ++++++++++++++++-----------
 test/test-async.hs                   |   2 +-
 2 files changed, 76 insertions(+), 54 deletions(-)

diff --git a/Control/Concurrent/Async/Internal.hs b/Control/Concurrent/Async/Internal.hs
index ce12e34..daca9d2 100644
--- a/Control/Concurrent/Async/Internal.hs
+++ b/Control/Concurrent/Async/Internal.hs
@@ -91,7 +91,7 @@ data Async a = Async
   { asyncThreadId :: {-# UNPACK #-} !ThreadId
                   -- ^ Returns the 'ThreadId' of the thread running
                   -- the given 'Async'.
-  , _asyncWait    :: STM (Either SomeException a)
+  , _asyncWait    :: STM (Either (ExceptionWithContext SomeException) a)
   }
 
 instance Eq (Async a) where
@@ -159,7 +159,7 @@ asyncUsing doFork action = do
    -- t <- forkFinally action (\r -> atomically $ putTMVar var r)
    -- slightly faster:
    t <- mask $ \restore ->
-          doFork $ try (restore action_plus) >>= atomically . putTMVar var
+          doFork $ tryWithContext (restore action_plus) >>= atomically . putTMVar var
    return (Async t (readTMVar var))
 
 
@@ -213,7 +213,6 @@ withAsyncOnWithUnmask ::
 withAsyncOnWithUnmask cpu actionWith =
   withAsyncUsing (rawForkOn cpu) (actionWith unsafeUnmask)
 
-#if MIN_VERSION_base(4,21,0)
 withAsyncUsing ::
   CALLSTACK
   (IO () -> IO ThreadId) -> IO a -> (Async a -> IO b) -> IO b
@@ -223,7 +222,7 @@ withAsyncUsing doFork action inner = do
   var <- newEmptyTMVarIO
   mask $ \restore -> do
     let action_plus = debugLabelMe >> action
-    t <- doFork $ try (restore action_plus) >>= atomically . putTMVar var
+    t <- doFork $ tryWithContext (restore action_plus) >>= atomically . putTMVar var
     let a = Async t (readTMVar var)
     -- Using catch/no/propagate and rethrowIO, we do not wrap the exception
     -- with a `WhileWaiting`
@@ -232,37 +231,55 @@ withAsyncUsing doFork action inner = do
       rethrowIO (e :: ExceptionWithContext SomeException)
     uninterruptibleCancel a
     return r
+
+-- * Compatibilty logic with base 4.21 for exception context. The rational here is that this module is implemented with 'ExceptionWithContext' as the basic building block with the following special cases:
+--
+-- - With base >= 4.21 (GHC 9.12), exception context is propagated correctly using the 'rethrowIO', 'catchNoPropagate', ... functions.
+-- - With base >= 4.20 (GHC 9.10), exception context logic exists, but not the 'rethrow' logic. We reimplemented these function which are basically discarding the context
+-- - With base < 4.20 (GHC 9.8 and older), we just use the old functions which does not know anything about exception context. We implement an alias 'ExceptionWithContext' which is actually bare exception.
+--
+-- For all version we implement 'dropContext' which is able to drop the
+-- context, for all the function such as 'poll' which returns an exception without context.
+
+
+-- | Drop the exception context
+dropContext :: ExceptionWithContext t -> t
+
+-- | Rethrow an exception inside 'STM' context, while preserving the 'ExceptionContext'. See 'rethrowIO' for details.
+rethrowSTM :: Exception e => ExceptionWithContext e -> STM a
+
+#if MIN_VERSION_base(4,21,0)
+-- The 'rethrowIO', 'catchNoPropagate' and 'tryWithContext' are already available in base
 #else
-withAsyncUsing ::
-  CALLSTACK
-  (IO () -> IO ThreadId) -> IO a -> (Async a -> IO b) -> IO b
--- The bracket version works, but is slow.  We can do better by
--- hand-coding it:
-withAsyncUsing doFork action inner = do
-  var <- newEmptyTMVarIO
-  mask $ \restore -> do
-    let action_plus = debugLabelMe >> action
-    t <- doFork $ try (restore action_plus) >>= atomically . putTMVar var
-    let a = Async t (readTMVar var)
-    r <- restore (inner a) `catchAll` \e -> do
-      uninterruptibleCancel a
-      throwIO e
-    uninterruptibleCancel a
-    return r
+-- In older version, we reimplement them
+rethrowIO :: ExceptionWithContext SomeException -> IO a
+catchNoPropagate :: forall e a. Exception e => IO a -> (ExceptionWithContext e -> IO a) -> IO a
+tryWithContext :: IO a -> IO (Either (ExceptionWithContext SomeException) a)
 #endif
 
-
--- | This function attempts at rethrowing while keeping the context
--- This is internal and only working with GHC >=9.12, otherwise it fallsback to
--- standard 'throwIO'
-rethrowIO' :: SomeException -> IO a
 #if MIN_VERSION_base(4,21,0)
-rethrowIO' e = 
-  case fromException e of
-    Just (e' :: ExceptionWithContext SomeException) -> rethrowIO e'
-    Nothing -> throwIO e
+dropContext (ExceptionWithContext _context e) = e
+rethrowSTM e = throwSTM (NoBacktrace e)
+#elif MIN_VERSION_base(4,20,0)
+dropContext (ExceptionWithContext _context e) = e
+
+-- For rethrowSTM and rethrowIO, it is important to drop the context, otherwise
+-- we throw an exception which is actually an "ExceptionWithContext" embedding
+-- an exception (so that's an exception inside an exception) and later
+-- "fromException" won't behave as expected.
+rethrowSTM e = throwSTM (dropContext e)
+
+rethrowIO e = throwIO (dropContext e)
+catchNoPropagate = catch
+tryWithContext = try
 #else
-rethrowIO' = throwIO
+dropContext e = e
+rethrowSTM e = throwSTM e
+
+type ExceptionWithContext e = e
+rethrowIO e = throwIO e
+catchNoPropagate = catch
+tryWithContext = try
 #endif
 
 -- | An exception annotation which stores the callstack of a 'wait',
@@ -330,33 +347,38 @@ poll = atomically . pollSTM
 --
 waitSTM :: Async a -> STM a
 waitSTM a = do
-   r <- waitCatchSTM a
+   r <- waitCatchSTMWithContext a
    either (rethrowSTM) return r
 
--- | This function attempts at rethrowing while keeping the context
--- This is internal and only working with GHC >=9.12, otherwise it fallsback to
--- standard 'throwSTM'
-rethrowSTM :: SomeException -> STM a
-#if MIN_VERSION_base(4,21,0)
-rethrowSTM e = 
-  case fromException e of
-    Just (e' :: ExceptionWithContext SomeException) -> throwSTM (NoBacktrace e')
-    Nothing -> throwSTM e
-#else
-rethrowSTM = throwSTM
-#endif
-
 -- | A version of 'waitCatch' that can be used inside an STM transaction.
 --
 {-# INLINE waitCatchSTM #-}
 waitCatchSTM :: Async a -> STM (Either SomeException a)
-waitCatchSTM (Async _ w) = w
+waitCatchSTM (Async _ w) = either (Left . dropContext) Right <$> w
+
+
+-- | A version of 'waitCatch' that can be used inside an STM transaction.
+-- 
+-- The returned exception keep the 'ExceptionContext'. See 'tryWithContext' for details.
+{-# INLINE waitCatchSTMWithContext #-}
+waitCatchSTMWithContext :: Async a -> STM (Either (ExceptionWithContext SomeException) a)
+waitCatchSTMWithContext (Async _ w) = w
 
 -- | A version of 'poll' that can be used inside an STM transaction.
 --
 {-# INLINE pollSTM #-}
 pollSTM :: Async a -> STM (Maybe (Either SomeException a))
-pollSTM (Async _ w) = (Just <$> w) `orElse` return Nothing
+pollSTM (Async _ w) = (Just . either (Left . dropContext) Right <$> w) `orElse` return Nothing
+
+#if MIN_VERSION_base(4,21,0)
+-- | A version of 'poll' that can be used inside an STM transaction.
+--
+-- It keep the exception context associated with the exception. See 'tryWithContext' for details.
+--
+{-# INLINE pollSTMWithContext #-}
+pollSTMWithContext :: Async a -> STM (Maybe (Either (ExceptionWithContext SomeException) a))
+pollSTMWithContext (Async _ w) = (Just <$> w) `orElse` return Nothing
+#endif
 
 -- | Cancel an asynchronous action by throwing the @AsyncCancelled@
 -- exception to it, and waiting for the `Async` thread to quit.
@@ -743,7 +765,7 @@ race left right = concurrently' left right collect
     collect m = do
         e <- m
         case e of
-            Left ex -> rethrowIO' ex
+            Left ex -> rethrowIO ex
             Right r -> return r
 
 -- race_ :: IO a -> IO b -> IO ()
@@ -757,7 +779,7 @@ concurrently left right = concurrently' left right (collect [])
     collect xs m = do
         e <- m
         case e of
-            Left ex -> rethrowIO' ex
+            Left ex -> rethrowIO ex
             Right r -> collect (r:xs) m
 
 -- concurrentlyE :: IO (Either e a) -> IO (Either e b) -> IO (Either e (a, b))
@@ -770,13 +792,13 @@ concurrentlyE left right = concurrently' left right (collect [])
     collect xs m = do
         e <- m
         case e of
-            Left ex -> rethrowIO' ex
+            Left ex -> rethrowIO ex
             Right r -> collect (r:xs) m
 
 concurrently' ::
   CALLSTACK
   IO a -> IO b
-  -> (IO (Either SomeException (Either a b)) -> IO r)
+  -> (IO (Either (ExceptionWithContext SomeException) (Either a b)) -> IO r)
   -> IO r
 concurrently' left right collect = do
     done <- newEmptyMVar
@@ -787,10 +809,10 @@ concurrently' left right collect = do
         -- the thread to terminate.
         lid <- forkIO $ uninterruptibleMask_ $
           restore (left >>= putMVar done . Right . Left)
-            `catchAll` (putMVar done . Left)
+            `catchNoPropagate` (putMVar done . Left)
         rid <- forkIO $ uninterruptibleMask_ $
           restore (right >>= putMVar done . Right . Right)
-            `catchAll` (putMVar done . Left)
+            `catchNoPropagate` (putMVar done . Left)
 
         count <- newIORef (2 :: Int)
         let takeDone = do
@@ -831,7 +853,7 @@ concurrently_ left right = concurrently' left right (collect 0)
     collect i m = do
         e <- m
         case e of
-            Left ex -> rethrowIO' ex
+            Left ex -> rethrowIO ex
             Right _ -> collect (i + 1 :: Int) m
 
 
diff --git a/test/test-async.hs b/test/test-async.hs
index 8de4f55..dca9be4 100644
--- a/test/test-async.hs
+++ b/test/test-async.hs
@@ -196,7 +196,7 @@ withasync_wait_blocked = do
     Left e ->
         case fromException e of
             Just BlockedIndefinitelyOnMVar -> return ()
-            Nothing -> assertFailure $ show e
+            Nothing -> assertFailure $ show ("what", e)
     Right () -> assertFailure ""
 
 concurrently_success :: Assertion

From 887ebeadf3629f8e4659e16084eb143f55af0e59 Mon Sep 17 00:00:00 2001
From: Guillaume Bouchard <guillaum.bouchard@gmail.com>
Date: Sat, 28 Mar 2026 15:04:01 +0400
Subject: [PATCH 7/8] feat: reexport the compatibility API

For easier usage of the Internal module in a backward compatible way in
other libraries.
---
 Control/Concurrent/Async/Internal.hs | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/Control/Concurrent/Async/Internal.hs b/Control/Concurrent/Async/Internal.hs
index daca9d2..d272ce3 100644
--- a/Control/Concurrent/Async/Internal.hs
+++ b/Control/Concurrent/Async/Internal.hs
@@ -26,7 +26,21 @@
 --
 -----------------------------------------------------------------------------
 
-module Control.Concurrent.Async.Internal where
+module Control.Concurrent.Async.Internal (
+  module Control.Concurrent.Async.Internal,
+
+#if MIN_VERSION_base(4,21,0)
+  -- * Compatibility wrapper for base < 4.20
+  -- These items are defined for base < 4.20 in this module and in
+  -- Control.Exception[.Context] for base >= 4.20. In order to ease usage of
+  -- the internal API, we reexport them here.
+  ExceptionWithContext(..),
+  rethrowIO,
+  catchNoPropagate,
+  tryWithContext
+#endif
+
+)where
 
 import Control.Concurrent.STM
 import Control.Exception

From bcf84dea063506ca089cec8e90d67c7df6b8b067 Mon Sep 17 00:00:00 2001
From: Guillaume Bouchard <guillaum.bouchard@gmail.com>
Date: Sat, 28 Mar 2026 15:04:01 +0400
Subject: [PATCH 8/8] chore: update the changelog with exception context
 propagation details

---
 changelog.md | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/changelog.md b/changelog.md
index c1ceae8..961ca9a 100644
--- a/changelog.md
+++ b/changelog.md
@@ -1,3 +1,14 @@
+- #165 async will preserve the `ExceptionContext` (on base >= 4.21, hence
+  starting with GHC 9.12) of exceptions runs inside the `Async`. Especially, it
+  now returs a callstack pointing inside the executed code instead of an
+  arbitrary location inside `async` library. For call such as `concurrently`,
+  it means that the callstack include the location of the `concurrently` and
+  continues inside the called functions. For call such as `withAsync + wait`,
+  the callstack contains the `withAsync` location as well as the location
+  inside the called function. The exception also contains an additional
+  annotation, `AsyncWaitLocation` which contains the location of the `wait`
+  call.
+
 ## Changes in 2.2.6
 
  - Added Control.Concurrent.Stream for processing streams with a fixed