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;
             }