From c3b70f92d61f1035ea944559cfae47b67dfa331b Mon Sep 17 00:00:00 2001
From: KATO Kanryu <k.kanryu@gmail.com>
Date: Mon, 25 May 2026 17:04:25 +0900
Subject: [PATCH] Feat: Refine Prompt Token Counter node description and add
 dynamic token count display widget

---
 nodes/text_nodes_token_counter.py |  3 --
 web/prompt_token_counter.js       | 78 +++++++++++++++++++++++++++++++
 2 files changed, 78 insertions(+), 3 deletions(-)
 create mode 100644 web/prompt_token_counter.js

diff --git a/nodes/text_nodes_token_counter.py b/nodes/text_nodes_token_counter.py
index 05aec09..2de4976 100644
--- a/nodes/text_nodes_token_counter.py
+++ b/nodes/text_nodes_token_counter.py
@@ -12,7 +12,6 @@ class PromptTokenCounterTextDriven(io.ComfyNode):
             node_id="TextDrivenPromptTokenCounter",
             display_name="Prompt Token Counter (Text-Driven)",
             category="Text-Driven Workflows",
-            # FIXED: Refined description based on developer specifications
             description=(
                 "A prompt token counter node designed for Stable Diffusion 1.5 and SDXL. "
                 "Executing this node or running the project calculates an approximate token count "
@@ -27,7 +26,6 @@ class PromptTokenCounterTextDriven(io.ComfyNode):
             ),
             is_output_node=True,
             inputs=[
-                # Keeps the exact static V3 compliant connection link-pin
                 io.String.Input("prompt", force_input=True),
             ],
             hidden=[
@@ -38,7 +36,6 @@ class PromptTokenCounterTextDriven(io.ComfyNode):
 
     @classmethod
     def execute(cls, prompt) -> tuple:
-        # LOGIC REMAINS ABSOLUTELY UNTOUCHED
         unique_id = cls.hidden.unique_id
         
         if not prompt:
diff --git a/web/prompt_token_counter.js b/web/prompt_token_counter.js
new file mode 100644
index 0000000..e12edd8
--- /dev/null
+++ b/web/prompt_token_counter.js
@@ -0,0 +1,78 @@
+import { app } from "../../scripts/app.js";
+import { api } from "../../scripts/api.js";
+
+// Helper function to determine the dynamic background and text colors safely
+function getTokenStyles(count) {
+    if (count <= 65) {
+        return { bgColor: "#0d3a19", textColor: "#00ffaa" }; // Safe zone: Deep Green
+    } else if (count <= 75) {
+        return { bgColor: "#4a420d", textColor: "#ffee00" }; // Warning zone: Deep Yellow
+    } else {
+        return { bgColor: "#4a0d0d", textColor: "#ff5555" }; // Danger zone: Deep Red
+    }
+}
+
+// Global listener to capture backend execution signals and dispatch to the specific target node
+api.addEventListener("text_driven_token_count_update", (event) => {
+    const detail = event.detail;
+    if (!detail || !detail.node_id) return;
+
+    const targetNode = app.graph.getNodeById(parseInt(detail.node_id));
+    if (targetNode) {
+        const displayWidget = targetNode.widgets.find(w => w.name === "token_count_display");
+        if (displayWidget) {
+            displayWidget.value = detail.count;
+            targetNode.setDirtyCanvas(true, true); // Force UI refresh to repaint colors instantly
+        }
+    }
+});
+
+app.registerExtension({
+    name: "ComfyUI-Text-Driven-Workflows.PromptTokenCounter",
+    async beforeRegisterNodeDef(nodeType, nodeData, app) {
+        if (nodeData.name === "TextDrivenPromptTokenCounter") {
+            const onNodeCreated = nodeType.prototype.onNodeCreated;
+            
+            nodeType.prototype.onNodeCreated = function () {
+                const r = onNodeCreated ? onNodeCreated.apply(this, arguments) : undefined;
+                const self = this;
+                
+                // Append the specialized read-only counter widget
+                const displayWidget = self.addWidget("text", "token_count_display", "2", null, { serialize: false });
+                
+                if (displayWidget) {
+                    displayWidget.label = "Token"; 
+                    displayWidget.value = "2";
+                    
+                    // Directly paint the container and traffic-light colors onto the node canvas area
+                    displayWidget.draw = function(ctx, node, widget_width, y, margin) {
+                        const count = parseInt(this.value) || 0;
+                        const styles = getTokenStyles(count);
+
+                        // Paint the container background box
+                        ctx.fillStyle = styles.bgColor;
+                        ctx.fillRect(margin, y, widget_width - margin * 2, LiteGraph.NODE_WIDGET_HEIGHT);
+
+                        // Render the clean English label text
+                        ctx.fillStyle = "#ffffff";
+                        ctx.font = "12px Arial";
+                        ctx.fillText(this.label || this.name, margin + 10, y + LiteGraph.NODE_WIDGET_HEIGHT * 0.7);
+
+                        // Render the right-aligned dynamic counter token integer
+                        ctx.fillStyle = styles.textColor;
+                        ctx.font = "bold 14px monospace";
+                        ctx.textAlign = "right";
+                        ctx.fillText(this.value, widget_width - margin - 10, y + LiteGraph.NODE_WIDGET_HEIGHT * 0.7);
+                        ctx.textAlign = "left"; // Restore canvas context behavior
+                    };
+                }
+                
+                if (self.computeSize) {
+                    self.setSize(self.computeSize());
+                }
+                
+                return r;
+            };
+        }
+    }
+});
\ No newline at end of file