From 4831e9c2c47b97f85fd771521f247a017d1f43e1 Mon Sep 17 00:00:00 2001
From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com>
Date: Sun, 13 Jul 2025 01:59:17 -0700
Subject: [PATCH] Refactor previous pr. (#8893)

---
 execution.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/execution.py b/execution.py
index bd638afba..c3a62f1cb 100644
--- a/execution.py
+++ b/execution.py
@@ -123,6 +123,8 @@ class CacheSet:
         }
         return result
 
+SENSITIVE_EXTRA_DATA_KEYS = ("auth_token_comfy_org", "api_key_comfy_org")
+
 def get_input_data(inputs, class_def, unique_id, outputs=None, dynprompt=None, extra_data={}):
     valid_inputs = class_def.INPUT_TYPES()
     input_data_all = {}
@@ -1045,11 +1047,10 @@ class PromptQueue:
             if status is not None:
                 status_dict = copy.deepcopy(status._asdict())
 
-            # Remove auth tokens from extra_data before storing in history
-            if "auth_token_comfy_org" in prompt[3]:
-                del prompt[3]["auth_token_comfy_org"]
-            if "api_key_comfy_org" in prompt[3]:
-                del prompt[3]["api_key_comfy_org"]
+            # Remove sensitive data from extra_data before storing in history
+            for sensitive_val in SENSITIVE_EXTRA_DATA_KEYS:
+                if sensitive_val in prompt[3]:
+                    prompt[3].pop(sensitive_val)
 
             self.history[prompt[1]] = {
                 "prompt": prompt,