diff --git a/CHANGELOG.md b/CHANGELOG.md
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -6,6 +6,11 @@
 and this project adheres to the
 [Haskell Package Versioning Policy](https://pvp.haskell.org/).
 
+## [1.3.3.0] - 2026-01-04
+### Added
+- **Type Families Support**: `deriveMockInstances` now supports type classes containing Associated Type Families.
+- **Improved `withMockIO`**: `withMockIO` now enables direct execution of IO actions within mock contexts without `liftIO`.
+
 ## [1.3.2.0] - 2026-01-04
 ### Added
 - **New Derivation Macros**:
diff --git a/README-ja.md b/README-ja.md
--- a/README-ja.md
+++ b/README-ja.md
@@ -235,7 +235,17 @@
 
   -- 実行
   f "arg"
+
+#### `withMockIO`: IO テストの簡略化
+`withMockIO` は `withMock` を IO に特化させたバージョンです。`liftIO` を使わずにモックコンテキスト内で直接 IO アクションを実行できます。
+
+```haskell
+it "IO test" $ withMockIO do
+  f <- mock (any ~> pure "result")
+  res <- someIOCall f
+  res `shouldBe` "result"
 ```
+```
 
 > [!IMPORTANT]
 > `expects`（宣言的検証）を使用する場合、モック定義部分は必ず **括弧 `(...)`** で囲んでください。
@@ -339,8 +349,6 @@
 deriveMockInstances [t|MonadLogger|]
 ```
 これにより、`lift . logInfo` を呼び出す `MockT m` のインスタンスが自動生成されます。
-> [!NOTE]
-> 現在、`deriveMockInstances` は Type Families を持つ型クラスをサポートしていません。
 
 ##### 明示的な No-op インスタンス (`deriveNoopInstance`)
 メソッド（特に `m ()` を返すもの）に対して、明示的なスタブ定義やベース実装を用意することなく、「何もしない」モックを作成したい場合があります。
diff --git a/README.md b/README.md
--- a/README.md
+++ b/README.md
@@ -237,7 +237,17 @@
 
   -- Execution
   f "arg"
+
+#### `withMockIO`: Simplified IO Testing
+`withMockIO` is an IO-specialized version of `withMock`. It allows you to run IO actions directly within the mock context without needing `liftIO`.
+
+```haskell
+it "IO test" $ withMockIO do
+  f <- mock (any ~> pure "result")
+  res <- someIOCall f
+  res `shouldBe` "result"
 ```
+```
 
 > [!IMPORTANT]
 > When using `expects` (declarative verification), you MUST wrap the mock definition in **parentheses `(...)`**.
@@ -341,8 +351,6 @@
 deriveMockInstances [t|MonadLogger|]
 ```
 This generates an instance for `MockT m` that calls `lift . logInfo`.
-> [!NOTE]
-> `deriveMockInstances` currently does not support type classes with Type Families.
 
 ##### Explicit No-op Instances (`deriveNoopInstance`)
 Sometimes you want a mock to do nothing for certain methods (especially those returning `m ()`) without having to define explicit stubs or provide a base implementation.
@@ -517,9 +525,9 @@
 |-----|-------|----|
 | 9.2.8 | 3.10.3.0 / 3.12.1.0 | Ubuntu, macOS, Windows |
 | 9.4.8 | 3.10.3.0 / 3.12.1.0 | Ubuntu, macOS, Windows |
-| 9.6.3 | 3.10.3.0 / 3.12.1.0 | Ubuntu, macOS, Windows |
-| 9.8.2 | 3.10.3.0 / 3.12.1.0 | Ubuntu, macOS, Windows |
-| 9.10.1 | 3.10.3.0 / 3.12.1.0 | Ubuntu, macOS, Windows |
+| 9.6.7 | 3.12.1.0 | Ubuntu, macOS, Windows |
+| 9.8.4 | 3.12.1.0 | Ubuntu, macOS, Windows |
+| 9.10.3 | 3.12.1.0 | Ubuntu, macOS, Windows |
 | 9.12.2 | 3.12.1.0 | Ubuntu, macOS, Windows |
 
 _Happy Mocking!_ 🐱
diff --git a/mockcat.cabal b/mockcat.cabal
--- a/mockcat.cabal
+++ b/mockcat.cabal
@@ -5,7 +5,7 @@
 -- see: https://github.com/sol/hpack
 
 name:           mockcat
-version:        1.3.2.0
+version:        1.3.3.0
 synopsis:       Declarative mocking with a single arrow `~>`.
 description:    Mockcat is a minimal, architecture-agnostic mocking library for Haskell.
                 It enables declarative verification and intent-driven matching, allowing you to define function behavior and expectations without specific architectural dependencies.
@@ -123,8 +123,10 @@
       Test.MockCat.TypeClassMinimalSpec
       Test.MockCat.TypeClassSpec
       Test.MockCat.TypeClassTHSpec
+      Test.MockCat.TypeFamilySpec
       Test.MockCat.UnsafeCheck
       Test.MockCat.WithMockErrorDiffSpec
+      Test.MockCat.WithMockIOSpec
       Test.MockCat.WithMockSpec
       Paths_mockcat
   hs-source-dirs:
diff --git a/src/Test/MockCat/Internal/Registry/Core.hs b/src/Test/MockCat/Internal/Registry/Core.hs
--- a/src/Test/MockCat/Internal/Registry/Core.hs
+++ b/src/Test/MockCat/Internal/Registry/Core.hs
@@ -22,6 +22,9 @@
   , getLastRecorder
   , getLastRecorderRaw
   , resetMockHistory
+  , getThreadWithMockContext
+  , setThreadWithMockContext
+  , clearThreadWithMockContext
   ) where
 
 import Control.Concurrent.STM
@@ -44,7 +47,7 @@
 import qualified Data.IntMap.Strict as IntMap
 import qualified Data.Map.Strict as Map
 import System.IO.Unsafe (unsafePerformIO)
-import Test.MockCat.Internal.Types (MockName, InvocationRecorder(..))
+import Test.MockCat.Internal.Types (MockName, InvocationRecorder(..), WithMockContext(..))
 import Test.MockCat.Internal.GHC.StableName (StableName, eqStableName, hashStableName, makeStableName)
 
 data SomeStableName = forall a. SomeStableName (StableName a)
@@ -123,7 +126,32 @@
 resetMockHistory :: IO ()
 resetMockHistory = do
   tid <- myThreadId
-  atomically $ modifyTVar' threadMockHistory (Map.delete tid)
+  atomically $ do
+    modifyTVar' threadMockHistory (Map.delete tid)
+    modifyTVar' threadWithMockStore (Map.delete tid)
+
+-- | Thread-local storage for WithMockContext
+threadWithMockStore :: TVar (Map.Map ThreadId WithMockContext)
+threadWithMockStore = unsafePerformIO $ newTVarIO Map.empty
+{-# NOINLINE threadWithMockStore #-}
+
+-- | Get the WithMockContext for the current thread.
+getThreadWithMockContext :: IO (Maybe WithMockContext)
+getThreadWithMockContext = do
+  tid <- myThreadId
+  atomically $ Map.lookup tid <$> readTVar threadWithMockStore
+
+-- | Set the WithMockContext for the current thread.
+setThreadWithMockContext :: WithMockContext -> IO ()
+setThreadWithMockContext ctx = do
+  tid <- myThreadId
+  atomically $ modifyTVar' threadWithMockStore (Map.insert tid ctx)
+
+-- | Clear the WithMockContext for the current thread.
+clearThreadWithMockContext :: IO ()
+clearThreadWithMockContext = do
+  tid <- myThreadId
+  atomically $ modifyTVar' threadWithMockStore (Map.delete tid)
 
 
 
diff --git a/src/Test/MockCat/Internal/Types.hs b/src/Test/MockCat/Internal/Types.hs
--- a/src/Test/MockCat/Internal/Types.hs
+++ b/src/Test/MockCat/Internal/Types.hs
@@ -18,7 +18,6 @@
 import Test.MockCat.AssociationList (AssociationList)
 import Prelude hiding (lookup)
 import Control.Monad.State ( State, MonadState, execState, modify )
-import Control.Monad.Reader (MonadReader, ask)
 
 type MockName = String
 
@@ -113,12 +112,6 @@
 --   of the `withMock` block. Storing `IO ()` avoids forcing concrete param
 --   types at registration time.
 newtype WithMockContext = WithMockContext (TVar [IO ()])
-
-class MonadWithMockContext m where
-  askWithMockContext :: m WithMockContext
-
-instance {-# OVERLAPPABLE #-} (MonadReader WithMockContext m) => MonadWithMockContext m where
-  askWithMockContext = ask
 
 -- | Expectation specification
 data Expectation params where
diff --git a/src/Test/MockCat/Mock.hs b/src/Test/MockCat/Mock.hs
--- a/src/Test/MockCat/Mock.hs
+++ b/src/Test/MockCat/Mock.hs
@@ -64,6 +64,7 @@
 import Test.MockCat.Internal.Verify (verifyExpectationDirect)
 import qualified Test.MockCat.Internal.MockRegistry as MockRegistry ( register )
 import Test.MockCat.Internal.Types
+import Test.MockCat.WithMock (askWithMockContext)
 import Test.MockCat.Param
 import Test.MockCat.Verify
 import Test.MockCat.Cons (Head(..), (:>)(..))
@@ -172,15 +173,9 @@
     _ <- liftIO $ MockRegistry.register (Just name) recorder fn
     pure fn
 
-
-
-    
-
-
 -- Specific instance for MockSpec (flag ~ 'True)
 instance
   ( MonadIO m
-  , MonadWithMockContext m
   , CreateMock params
   , MockBuilder (ToMockParams params) fn verifyParams
   , Show verifyParams
@@ -194,7 +189,7 @@
     BuiltMock { builtMockFn = fn, builtMockRecorder = recorder } <- buildMock (Just name) (toParams params) :: m (BuiltMock fn verifyParams)
     _ <- liftIO $ MockRegistry.register (Just name) recorder fn
     
-    WithMockContext ctxRef <- askWithMockContext
+    WithMockContext ctxRef <- liftIO askWithMockContext
     let resolved = ResolvedMock (Just name) recorder
     let verifyAction = mapM_ (verifyExpectationDirect resolved) exps
     liftIO $ atomically $ modifyTVar' ctxRef (++ [verifyAction])
@@ -232,7 +227,6 @@
 --   > f <- mock $ any ~> True `expects` do called once
 instance {-# OVERLAPPING #-}
   ( MonadIO m
-  , MonadWithMockContext m
   , CreateMock params
   , MockBuilder (ToMockParams params) fn verifyParams
   , Show verifyParams
@@ -246,7 +240,7 @@
     BuiltMock { builtMockFn = fn, builtMockRecorder = recorder } <- buildMock Nothing (toParams params) :: m (BuiltMock fn verifyParams)
     _ <- liftIO $ MockRegistry.register Nothing recorder fn
     
-    WithMockContext ctxRef <- askWithMockContext
+    WithMockContext ctxRef <- liftIO askWithMockContext
     let resolved = ResolvedMock Nothing recorder
     let verifyAction = mapM_ (verifyExpectationDirect resolved) exps
     liftIO $ atomically $ modifyTVar' ctxRef (++ [verifyAction])
diff --git a/src/Test/MockCat/MockT.hs b/src/Test/MockCat/MockT.hs
--- a/src/Test/MockCat/MockT.hs
+++ b/src/Test/MockCat/MockT.hs
@@ -24,7 +24,7 @@
   )
 import Control.Monad.IO.Class (MonadIO(..))
 import Control.Monad.Trans.Class (MonadTrans(..))
-import Control.Monad.Reader (ReaderT(..), runReaderT, asks, MonadReader(..))
+import Control.Monad.Reader (ReaderT(..), runReaderT, MonadReader(..))
 import Control.Monad.Error.Class (MonadError(..))
 import Control.Monad.State.Class (MonadState(..))
 import Control.Monad.Writer.Class (MonadWriter(..))
@@ -33,7 +33,7 @@
 import Data.IORef (newIORef, IORef)
 import Data.Dynamic (Dynamic)
 import UnliftIO (MonadUnliftIO(..))
-import Test.MockCat.Internal.Types (InvocationRecorder, WithMockContext(..), MonadWithMockContext(..))
+import Test.MockCat.Internal.Types (InvocationRecorder, WithMockContext(..))
 import Test.MockCat.Verify (ResolvableParamsOf)
 import Control.Concurrent.MVar (MVar)
 import qualified Data.Map.Strict as Map
@@ -74,8 +74,6 @@
   withRunInIO inner = MockT $ ReaderT $ \env ->
     withRunInIO $ \run -> inner (\(MockT r) -> run (runReaderT r env))
 
-instance {-# OVERLAPPING #-} Monad m => MonadWithMockContext (MockT m) where
-  askWithMockContext = MockT $ asks envWithMockContext
 
 instance {-# OVERLAPPABLE #-} MonadReader r m => MonadReader r (MockT m) where
   ask = lift ask
@@ -154,22 +152,25 @@
 runMockT :: MonadIO m => MockT m a -> m a
 runMockT (MockT r) = do
   liftIO Registry.resetMockHistory
-  defsVar <- liftIO $ newTVarIO []
   expectsVar <- liftIO $ newTVarIO []
+  let withMockCtx = WithMockContext expectsVar
+  defsVar <- liftIO $ newTVarIO []
   fwdRef <- liftIO $ newIORef Map.empty
   let env =
         MockTEnv
           { envDefinitions = defsVar
-          , envWithMockContext = WithMockContext expectsVar
+          , envWithMockContext = withMockCtx
           , envNameForwarders = fwdRef
           }
   -- Run user code with a per-run overlay registry active so registry writes/read
   -- during this MockT invocation are isolated to this run.
   overlay <- liftIO Registry.createOverlay
   liftIO $ Registry.installOverlay overlay
+  liftIO $ Registry.setThreadWithMockContext withMockCtx
   a <- runReaderT r env
   actions <- liftIO $ readTVarIO expectsVar
   liftIO $ sequence_ actions
+  liftIO Registry.clearThreadWithMockContext
   liftIO Registry.clearOverlay
   pure a
 
diff --git a/src/Test/MockCat/TH.hs b/src/Test/MockCat/TH.hs
--- a/src/Test/MockCat/TH.hs
+++ b/src/Test/MockCat/TH.hs
@@ -626,6 +626,7 @@
   
   let isSupportedDec (SigD _ _) = True
       isSupportedDec (PragmaD _) = True
+      isSupportedDec (OpenTypeFamilyD _) = True
       isSupportedDec _ = False
   let unsupportedDecs = filter (not . isSupportedDec) (cmDecs classMetadata)
   
@@ -633,14 +634,21 @@
     case unsupportedDecs of
       (x:_) -> pure $ Left $ "deriveMockInstances: Unsupported declaration in class: " <> pprint x <> 
                         ". This error is reported at the usage site, but the cause is the macro definition for `" <> show className <> "`."
-      [] -> sequence <$> mapM (createLiftInstanceFnDec monadVarName) sigDecs
+      [] -> do
+        let typeFamilyHeads = [head | OpenTypeFamilyD head <- cmDecs classMetadata]
+        typeInstDecs <- sequence <$> mapM (\h -> Right <$> createTypeInstanceDec monadVarName h) typeFamilyHeads
+        sigInstDecs <- sequence <$> mapM (createLiftInstanceFnDec monadVarName) sigDecs
+        case (typeInstDecs, sigInstDecs) of
+          (Right tDecs, Right sDecs) -> pure $ Right (tDecs ++ sDecs)
+          (Left err, _) -> pure $ Left err
+          (_, Left err) -> pure $ Left err
 
   case instanceBodyDecsResult of
     Right decs -> do
       instanceHead <- createInstanceType ty monadVarName newTypeVars
       let instanceConstraint = foldl AppT ty (map (VarT . getTypeVarName) newTypeVars)
       instanceDec <- instanceD
-        (pure (instanceConstraint : cmContext classMetadata))
+        (pure [instanceConstraint])
         (pure instanceHead)
         (map pure decs)
       pure [instanceDec]
diff --git a/src/Test/MockCat/WithMock.hs b/src/Test/MockCat/WithMock.hs
--- a/src/Test/MockCat/WithMock.hs
+++ b/src/Test/MockCat/WithMock.hs
@@ -17,6 +17,8 @@
 -}
 module Test.MockCat.WithMock
   ( withMock
+  , withMockIO
+  , askWithMockContext
   , expects
   , MockResult(..)
   , called
@@ -32,7 +34,6 @@
   , lessThan
   , anything
   , WithMockContext(..)
-  , MonadWithMockContext(..)
   , Expectation(..)
   , Expectations(..)
   , verifyExpectationDirect
@@ -40,6 +41,7 @@
 
 import Control.Monad.IO.Class (MonadIO, liftIO)
 import Control.Monad.Reader (ReaderT(..), runReaderT)
+import Control.Exception (bracket_)
 import Control.Concurrent.STM (newTVarIO, readTVarIO, atomically, modifyTVar')
 import Control.Monad.State (get, put, modify)
 import Test.MockCat.Verify (TimesSpec(..), times, once, never, atLeast, atMost, greaterThan, lessThan, anything, ResolvableMock, ResolvableParamsOf)
@@ -47,7 +49,6 @@
 import Test.MockCat.Internal.Types
   ( VerifyOrderMethod(..)
   , WithMockContext(..)
-  , MonadWithMockContext(..)
   , Expectation(..)
   , Expectations(..)
   , runExpectations
@@ -56,6 +57,7 @@
   , ResolvedMock(..)
   )
 import qualified Test.MockCat.Internal.Registry.Core as MockRegistry
+import Test.MockCat.Internal.Registry.Core (getThreadWithMockContext, setThreadWithMockContext, clearThreadWithMockContext)
 import Unsafe.Coerce (unsafeCoerce)
 
 import Test.MockCat.Param (Param(..), param, EqParams(..))
@@ -74,12 +76,34 @@
 withMock action = do
   ctxVar <- newTVarIO []
   let ctx = WithMockContext ctxVar
-  result <- runReaderT action ctx
-  -- Verify all registered verification actions
-  actions <- readTVarIO ctxVar
-  sequence_ actions
-  pure result
+  bracket_ (setThreadWithMockContext ctx) clearThreadWithMockContext $ do
+    result <- runReaderT action ctx
+    -- Verify all registered verification actions
+    actions <- readTVarIO ctxVar
+    sequence_ actions
+    pure result
 
+-- | IO version of withMock
+withMockIO :: IO a -> IO a
+withMockIO action = do
+  ctxVar <- newTVarIO []
+  let ctx = WithMockContext ctxVar
+  bracket_ (setThreadWithMockContext ctx) clearThreadWithMockContext $ do
+    result <- action
+    -- Verify all registered verification actions
+    actions <- readTVarIO ctxVar
+    sequence_ actions
+    pure result
+
+-- | Retrieve the current mock context from thread-local storage.
+--   Throws an error if no context is found.
+askWithMockContext :: IO WithMockContext
+askWithMockContext = do
+  mCtx <- getThreadWithMockContext
+  case mCtx of
+    Just ctx -> pure ctx
+    Nothing -> errorWithoutStackTrace "askWithMockContext: No WithMockContext found in current thread. Use withMock or withMockIO."
+
 -- | Attach expectations to a mock function
 --   Supports both single expectation and multiple expectations in a do block
 infixl 0 `expects`
@@ -147,7 +171,6 @@
 --   Strict matching of params
 instance
   ( MonadIO m
-  , MonadWithMockContext m
   , ResolvableMock fn
   , ResolvableParamsOf fn ~ params
   , ExtractParams exp
@@ -159,7 +182,7 @@
   ExpectsDispatchImpl 'False fn exp m
   where
   expectsDispatchImpl mockFnM exp = do
-    (WithMockContext ctxVar) <- askWithMockContext
+    WithMockContext ctxVar <- liftIO askWithMockContext
     -- Try to help type inference by using exp first
     let _ = extractParams exp :: Proxy params
     mockFn <- mockFnM
@@ -180,7 +203,6 @@
 --   Dynamic resolution using expectation params
 instance
   ( MonadIO m
-  , MonadWithMockContext m
   , BuildExpectations (MockResult params) (Expectations params ()) params
   , Show params
   , EqParams params
@@ -188,7 +210,7 @@
   ExpectsDispatchImpl 'True (MockResult params) (Expectations params ()) m
   where
   expectsDispatchImpl mockFnM exp = do
-    (WithMockContext ctxVar) <- askWithMockContext
+    WithMockContext ctxVar <- liftIO askWithMockContext
     _ <- mockFnM
     (mockName, mRecorder) <- liftIO MockRegistry.getLastRecorderRaw
     resolved <- case mRecorder of
@@ -206,7 +228,6 @@
 --   Dynamic resolution using expectation params
 instance
   ( MonadIO m
-  , MonadWithMockContext m
   , BuildExpectations () (Expectations params ()) params
   , Show params
   , EqParams params
@@ -214,7 +235,7 @@
   ExpectsDispatchImpl 'True () (Expectations params ()) m
   where
   expectsDispatchImpl mockFnM exp = do
-    (WithMockContext ctxVar) <- askWithMockContext
+    WithMockContext ctxVar <- liftIO askWithMockContext
     _ <- mockFnM
     (mockName, mRecorder) <- liftIO MockRegistry.getLastRecorderRaw
     resolved <- case mRecorder of
diff --git a/test/Spec.hs b/test/Spec.hs
--- a/test/Spec.hs
+++ b/test/Spec.hs
@@ -22,12 +22,14 @@
 import Test.MockCat.ShouldBeCalledSpec as ShouldBeCalled
 import Test.MockCat.ShouldBeCalledMockMSpec as ShouldBeCalledMockM
 import Test.MockCat.WithMockSpec as WithMock
+import Test.MockCat.WithMockIOSpec as WithMockIO
 import Test.MockCat.ShouldBeCalledErrorDiffSpec as ShouldBeCalledErrorDiff
 import Test.MockCat.WithMockErrorDiffSpec as WithMockErrorDiff
 import Test.MockCat.THCompareSpec as THCompare
 import ReadmeVerifySpec as ReadmeVerify
 import qualified Test.MockCat.HPCFallbackSpec as HPCFallback
 import qualified Test.MockCat.MultipleMocksSpec as MultipleMocks
+import qualified Test.MockCat.TypeFamilySpec as TypeFamily
 import Test.MockCat.UnsafeCheck ()
 import Test.QuickCheck (property)
 import qualified Property.ConcurrentCountProp as ConcurrencyProp
@@ -65,11 +67,13 @@
     ShouldBeCalled.spec
     ShouldBeCalledMockM.spec
     WithMock.spec
+    WithMockIO.spec
     ShouldBeCalledErrorDiff.spec
     WithMockErrorDiff.spec
     ReadmeVerify.spec
     HPCFallback.spec
     MultipleMocks.spec
+    TypeFamily.spec
     describe "Property Concurrency" $ do
       it "total apply count is preserved across threads" $ property ConcurrencyProp.prop_concurrent_total_apply_count
     describe "Property Lazy Evaluation" $ do
diff --git a/test/Test/MockCat/TypeFamilySpec.hs b/test/Test/MockCat/TypeFamilySpec.hs
new file mode 100644
--- /dev/null
+++ b/test/Test/MockCat/TypeFamilySpec.hs
@@ -0,0 +1,31 @@
+{-# LANGUAGE TemplateHaskell #-}
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE DataKinds #-}
+module Test.MockCat.TypeFamilySpec (spec) where
+
+import Test.Hspec
+import Test.MockCat
+import Data.Kind (Type)
+
+class Monad m => MonadKeyValue m where
+  type Key m :: Type
+  getValue :: Key m -> m String
+
+deriveMockInstances [t|MonadKeyValue|]
+
+spec :: Spec
+spec = do
+  describe "deriveMockInstances with Type Families" $ do
+    it "can lift MonadKeyValue to MockT" $ do
+      withMock $ do
+        runMockT $ do
+          -- This just verifies it compiles and the instance is valid
+          pure ()
+
+instance MonadKeyValue IO where
+  type Key IO = String
+  getValue k = pure $ "Value for " ++ k
diff --git a/test/Test/MockCat/WithMockIOSpec.hs b/test/Test/MockCat/WithMockIOSpec.hs
new file mode 100644
--- /dev/null
+++ b/test/Test/MockCat/WithMockIOSpec.hs
@@ -0,0 +1,58 @@
+{-# LANGUAGE BlockArguments #-}
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE TypeApplications #-}
+module Test.MockCat.WithMockIOSpec (spec) where
+
+import Test.Hspec
+import Test.MockCat
+import Control.Exception (try, ErrorCall(..), SomeException)
+import Control.Concurrent (forkIO)
+import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)
+import Control.Monad (void)
+
+spec :: Spec
+spec = do
+  describe "withMockIO" $ do
+    it "can run IO actions directly without liftIO" $ do
+      withMockIO $ do
+        f <- mock $ "hello" ~> "world"
+        f "hello" `shouldBe` "world"
+        f `shouldBeCalled` "hello"
+
+    it "cleans up context even if an exception occurs" $ do
+      -- Raise an exception in the first test
+      void $ try @ErrorCall $ withMockIO $ do
+        void $ mock $ "a" ~> "b"
+        error "force fail"
+
+      -- Verify that no remnants (like expectations) from the previous test remain in the second test
+      withMockIO $ do
+        f <- mock $ "x" ~> "y"
+        f "x" `shouldBe` "y"
+        -- Verify that the expectation "a" ~> "b" from the previous test is not verified here
+
+    it "supports nested withMockIO" $ do
+      withMockIO $ do
+        f1 <- mock $ "outer" ~> "ok"
+        withMockIO $ do
+          f2 <- mock $ "inner" ~> "ok"
+          f2 "inner" `shouldBe` "ok"
+          f2 `shouldBeCalled` "inner"
+        f1 "outer" `shouldBe` "ok"
+        f1 `shouldBeCalled` "outer"
+
+    it "isolates context between parent and child threads" $ do
+      mvar <- newEmptyMVar
+      withMockIO $ do
+        f <- mock $ "parent" ~> "ok"
+        void $ forkIO $ do
+          res <- try @SomeException (mock ("child" ~> "ok") `expects` called once)
+          putMVar mvar res
+
+        f "parent" `shouldBe` "ok"
+        f `shouldBeCalled` "parent"
+
+      childRes <- takeMVar mvar
+      case childRes of
+        Left e -> show e `shouldContain` "No WithMockContext found"
+        Right _ -> fail "Child thread should not have accessed parent context"