diff --git a/src/Fluss.HotChocolate/AddExtensionMiddleware.cs b/src/Fluss.HotChocolate/AddExtensionMiddleware.cs
index 484c70f..30eca4f 100644
--- a/src/Fluss.HotChocolate/AddExtensionMiddleware.cs
+++ b/src/Fluss.HotChocolate/AddExtensionMiddleware.cs
@@ -21,7 +21,7 @@ public async ValueTask InvokeAsync(IRequestContext context)
     {
         await _next.Invoke(context);
 
-        if (!context.ContextData.TryGetValue(nameof(UnitOfWork), out var unitOfWork))
+        if (!context.ContextData.TryGetValue(nameof(IUnitOfWork), out var unitOfWork))
         {
             return;
         }
@@ -38,7 +38,7 @@ public async ValueTask InvokeAsync(IRequestContext context)
         {
             if (context.Result is QueryResult subsequentQueryResult)
             {
-                context.Result = QueryResultBuilder.FromResult(subsequentQueryResult).AddContextData(nameof(UnitOfWork),
+                context.Result = QueryResultBuilder.FromResult(subsequentQueryResult).AddContextData(nameof(IUnitOfWork),
 unitOfWork).Create();
             }
 
@@ -84,7 +84,7 @@ private async IAsyncEnumerable<IQueryResult> LiveResults(IReadOnlyDictionary<str
 
         while (true)
         {
-            if (contextData == null || !contextData.TryGetValue(nameof(UnitOfWork), out var value))
+            if (contextData == null || !contextData.TryGetValue(nameof(IUnitOfWork), out var value))
             {
                 break;
             }
diff --git a/src/Fluss.HotChocolate/UnitOfWorkParameterExpressionBuilder.cs b/src/Fluss.HotChocolate/UnitOfWorkParameterExpressionBuilder.cs
index 5a23aa7..59381f0 100644
--- a/src/Fluss.HotChocolate/UnitOfWorkParameterExpressionBuilder.cs
+++ b/src/Fluss.HotChocolate/UnitOfWorkParameterExpressionBuilder.cs
@@ -12,7 +12,7 @@ public class UnitOfWorkParameterExpressionBuilder : IParameterExpressionBuilder
     private static readonly MethodInfo GetOrSetGlobalStateUnitOfWorkMethod =
         typeof(ResolverContextExtensions).GetMethods()
             .First(m => m.Name == nameof(ResolverContextExtensions.GetOrSetGlobalState))
-            .MakeGenericMethod(typeof(UnitOfWork));
+            .MakeGenericMethod(typeof(IUnitOfWork));
 
     private static readonly MethodInfo GetGlobalStateOrDefaultLongMethod =
         typeof(ResolverContextExtensions).GetMethods()
@@ -21,22 +21,22 @@ public class UnitOfWorkParameterExpressionBuilder : IParameterExpressionBuilder
 
     private static readonly MethodInfo ServiceUnitOfWorkMethod =
         typeof(IPureResolverContext).GetMethods().First(
-            method => method is { Name: nameof(IPureResolverContext.Service), IsGenericMethod: true })
-            .MakeGenericMethod(typeof(UnitOfWork));
+                method => method is { Name: nameof(IPureResolverContext.Service), IsGenericMethod: true })
+            .MakeGenericMethod(typeof(IUnitOfWork));
 
     private static readonly MethodInfo WithPrefilledVersionMethod =
-        typeof(UnitOfWork).GetMethods(BindingFlags.Instance | BindingFlags.Public)
-            .First(m => m.Name == nameof(UnitOfWork.WithPrefilledVersion));
+        typeof(IUnitOfWork).GetMethods(BindingFlags.Instance | BindingFlags.Public)
+            .First(m => m.Name == nameof(IUnitOfWork.WithPrefilledVersion));
 
-    public bool CanHandle(ParameterInfo parameter) => typeof(UnitOfWork) == parameter.ParameterType
-                                                      || typeof(IUnitOfWork) == parameter.ParameterType;
+    public bool CanHandle(ParameterInfo parameter) =>
+        typeof(IUnitOfWork) == parameter.ParameterType;
 
     /*
      * Produces something like this: context.GetOrSetGlobalState(
-     *      nameof(UnitOfWork.UnitOfWork),
+     *      nameof(IUnitOfWork),
      *      _ =>
      *          context
-     *              .Service<UnitOfWork.UnitOfWork>()
+     *              .Service<IUnitOfWork>()
      *              .WithPrefilledVersion(
      *                  context.GetGlobalState<long>(PrefillUnitOfWorkVersion)
      *              ))!;
@@ -53,13 +53,14 @@ public Expression Build(ParameterExpressionBuilderContext builderContext)
                 context,
                 Expression.Constant(PrefillUnitOfWorkVersion)));
 
-        return Expression.Call(null, GetOrSetGlobalStateUnitOfWorkMethod, context, Expression.Constant(nameof(UnitOfWork)),
-            Expression.Lambda<Func<string, UnitOfWork>>(
+        return Expression.Call(null, GetOrSetGlobalStateUnitOfWorkMethod, context,
+            Expression.Constant(nameof(IUnitOfWork)),
+            Expression.Lambda<Func<string, IUnitOfWork>>(
                 getNewUnitOfWork,
                 Expression.Parameter(typeof(string))));
     }
 
     public ArgumentKind Kind => ArgumentKind.Custom;
-    public bool IsPure => true;
+    public bool IsPure => false;
     public bool IsDefaultHandler => false;
-}
+}
\ No newline at end of file
diff --git a/src/Fluss.UnitTest/HotChocolate/UnitOfWorkParameterExpressionBuilderTest.cs b/src/Fluss.UnitTest/HotChocolate/UnitOfWorkParameterExpressionBuilderTest.cs
index 433de92..7976a2e 100644
--- a/src/Fluss.UnitTest/HotChocolate/UnitOfWorkParameterExpressionBuilderTest.cs
+++ b/src/Fluss.UnitTest/HotChocolate/UnitOfWorkParameterExpressionBuilderTest.cs
@@ -12,10 +12,10 @@ public class UnitOfWorkParameterExpressionBuilderTest
     private readonly UnitOfWorkParameterExpressionBuilder _builder = new();
 
     [Fact]
-    public void CanHandle_ShouldReturnTrueForUnitOfWorkParameter()
+    public void CanHandle_ShouldReturnFalseForUnitOfWorkParameter()
     {
         var parameter = typeof(TestClass).GetMethod(nameof(TestClass.MethodWithUnitOfWork))!.GetParameters()[0];
-        Assert.True(_builder.CanHandle(parameter));
+        Assert.False(_builder.CanHandle(parameter));
     }
 
     [Fact]
@@ -39,9 +39,9 @@ public void Kind_ShouldReturnCustom()
     }
 
     [Fact]
-    public void IsPure_ShouldReturnTrue()
+    public void IsPure_ShouldReturnFalse()
     {
-        Assert.True(_builder.IsPure);
+        Assert.False(_builder.IsPure);
     }
 
     [Fact]