From 911e9e014bfcc9928a879fa2d8d842dc22d0d74d Mon Sep 17 00:00:00 2001
From: Stefan <stefan@gorules.io>
Date: Mon, 10 Feb 2025 10:15:33 +0100
Subject: [PATCH] fix: cleanup dt reference

---
 core/engine/src/handler/graph.rs     |  1 -
 core/engine/src/handler/table/zen.rs | 10 ++++++++--
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/core/engine/src/handler/graph.rs b/core/engine/src/handler/graph.rs
index c19fae58..eae6e9dd 100644
--- a/core/engine/src/handler/graph.rs
+++ b/core/engine/src/handler/graph.rs
@@ -401,7 +401,6 @@ impl<L: DecisionLoader + 'static, A: CustomNodeAdapter + 'static> DecisionGraph<
 
                     node_request.input.dot_remove("$nodes");
                     res.output.dot_remove("$nodes");
-                    res.output.dot_remove("$");
 
                     trace!({
                         input: node_request.input,
diff --git a/core/engine/src/handler/table/zen.rs b/core/engine/src/handler/table/zen.rs
index 216ea745..a97239f8 100644
--- a/core/engine/src/handler/table/zen.rs
+++ b/core/engine/src/handler/table/zen.rs
@@ -67,10 +67,16 @@ impl<'a> DecisionTableHandlerInner<'a> {
 
                     self_ref.isolate.clear_references();
                     self_ref.isolate.set_environment(input);
-                    match &content.hit_policy {
+                    let output = match &content.hit_policy {
                         DecisionTableHitPolicy::First => self_ref.handle_first_hit(&content).await,
                         DecisionTableHitPolicy::Collect => self_ref.handle_collect(&content).await,
-                    }
+                    };
+
+                    if let Ok(ok) = &output {
+                        ok.output.dot_remove("$");
+                    };
+
+                    output
                 }
             })
             .await