diff --git a/src/main/java/com/powsybl/ws/commons/computation/service/AbstractWorkerService.java b/src/main/java/com/powsybl/ws/commons/computation/service/AbstractWorkerService.java
index 3bd1155..4787ad4 100644
--- a/src/main/java/com/powsybl/ws/commons/computation/service/AbstractWorkerService.java
+++ b/src/main/java/com/powsybl/ws/commons/computation/service/AbstractWorkerService.java
@@ -153,12 +153,20 @@ public Consumer<Message<String>> consumeRun() {
                     this.handleNonCancellationException(resultContext, e, rootReporter);
                 }
             } finally {
-                futures.remove(resultContext.getResultUuid());
-                cancelComputationRequests.remove(resultContext.getResultUuid());
+                clean(resultContext);
             }
         };
     }
 
+    /**
+     * Perform cleaning
+     * @param resultContext The context of the computation
+     */
+    protected void clean(AbstractResultContext<C> resultContext) {
+        futures.remove(resultContext.getResultUuid());
+        cancelComputationRequests.remove(resultContext.getResultUuid());
+    }
+
     /**
      * Handle exception in consumeRun that is not a CancellationException
      * @param resultContext The context of the computation