From 6559de8255cf808abeebfbe76fd2ba6ddc0f1319 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Tue, 27 Oct 2026 19:43:48 -0700 Subject: [PATCH] fix(usage): never replay a node twice in one job; dedupe receipts once complete (#347) A repeated execution_cached id or message for the same job now counts the node once. Receipts replace equal older ones only when their job finishes (or the next job starts), so a list still recording [A] of [A, B] no longer drops an older [A] receipt. Node-level tests check the receipt each execution records. --- tests/test_cached_runs.py | 26 ++++++++++++++++++ tests/test_pair_usage_tracking.py | 17 ++++++++++++ tests/test_prompt_usage.py | 44 +++++++++++++++++++++++++++++++ utils/cached_runs.py | 25 +++++++++++++++--- 4 files changed, 108 insertions(+), 4 deletions(-) diff --git a/tests/test_cached_runs.py b/tests/test_cached_runs.py index 668acad3..d6ea4954 100644 --- a/tests/test_cached_runs.py +++ b/tests/test_cached_runs.py @@ -337,6 +337,31 @@ class TestCachedJobsCount(JobCase): self.job(g, cached=["pm"]) self.assertEqual([self.count(t) for t in ("one", "two", "three")], [2, 2, 2]) + def test_a_repeated_id_or_message_never_counts_a_node_twice(self): + g = graph() + self.first(g) + self.job(g, cached=["pm", "pm"]) + self.assertEqual(self.count(LLM), 2) + self.item_id += 1 + self.runs.begin_attempt(((1, "again", g, {}, ["save"], {}), self.item_id)) + for _ in range(2): + self.runs.replay(self.executor, "again", ["pm"], lambda: self.db) + self.assertEqual(self.count(LLM), 3) + + def test_a_shorter_list_keeps_its_receipt_after_a_longer_one(self): + g = graph() + self.first(g, text="alpha") + results = [] + + def run(): + for index, text in enumerate(("alpha", "beta")): + results.append(self.execute(g, text=text, index=index)) + + self.job(g, run=run) + self.cache("pm", (mock.ANY, "alpha")) + self.job(g, cached=["pm"]) + self.assertEqual((self.count("alpha"), self.count("beta")), (3, 1)) + def test_a_job_reusing_a_prompt_id_counts_and_executes_normally(self): g = graph() self.first(g) @@ -632,6 +657,7 @@ class TestReceipts(unittest.TestCase): for item_id, text in ((1, "x"), (2, "y"), (3, "z"), (4, "y")): self.begin(item_id) self.runs.record_item("a", (text,), Facts(text)) + self.runs.end_attempt(item_id) receipts = self.runs._receipts[key] self.assertEqual([r.item_id for r in receipts], [4, 3]) diff --git a/tests/test_pair_usage_tracking.py b/tests/test_pair_usage_tracking.py index 9039f29d..f451cad2 100644 --- a/tests/test_pair_usage_tracking.py +++ b/tests/test_pair_usage_tracking.py @@ -85,6 +85,23 @@ class TestPairUsageTracking(unittest.TestCase): self.db.get_prompt_by_id(pair["negative_prompt_id"])["run_count"], 0 ) + def test_execution_records_a_cached_run_receipt_for_both_editors(self): + from utils.cached_runs import Facts + + node = self.node() + node.cached_runs = Mock() + node.cached_runs.replayed.return_value = False + node.encode_prompt( + Mock(), + "cat", + negative_text="", + prompt=pair_graph(negative=""), + unique_id="1", + ) + node.cached_runs.record_item.assert_called_once_with( + "1", ("cat", ""), Facts(usage_tracking.generate_prompt_hash("cat"), None) + ) + def test_validation_rejection_never_counts_or_saves(self): self.submit(pair_graph()) self.assertEqual(self.db.list_prompt_pairs()["total"], 0) diff --git a/tests/test_prompt_usage.py b/tests/test_prompt_usage.py index c4606641..c1dd27f1 100644 --- a/tests/test_prompt_usage.py +++ b/tests/test_prompt_usage.py @@ -278,5 +278,49 @@ class TestNodeRecordsUse(UsageTestCase): self.assertEqual(self.db.get_prompt_by_id(pid)["run_count"], 2) +class TestNodeRecordsCachedRunReceipt(UsageTestCase): + """Each execution leaves a receipt for cached jobs (#347).""" + + GRAPH = { + "1": {"class_type": "PromptManagerText", "inputs": {"text": "typed text"}}, + } + + def _node(self): + from unittest import mock + + from prompt_manager_text import PromptManagerText + + node = PromptManagerText.__new__(PromptManagerText) + node.db = self.db + node.logger = __import__("logging").getLogger("test.prompt_usage") + node.prompt_tracker = mock.Mock() + node.comfyui_integration = mock.Mock() + node._inject_lora_trigger_words = lambda text: text + node.cached_runs = mock.Mock() + node.cached_runs.replayed.return_value = False + return node + + def test_receipt_holds_the_returned_text_and_the_stored_hash(self): + from utils.cached_runs import Facts + + node = self._node() + node.process_text( + "typed text", prepend_text="pre", prompt=self.GRAPH, unique_id="1" + ) + node.cached_runs.record_item.assert_called_once_with( + "1", ("pre typed text",), Facts(generate_prompt_hash("typed text")) + ) + # Hook-visible text: the node counts nothing itself + row = self.db.get_prompt_by_hash(generate_prompt_hash("typed text")) + self.assertEqual(row["run_count"], 0) + + def test_empty_editor_records_empty_facts(self): + from utils.cached_runs import Facts + + node = self._node() + node.process_text(" ", prompt=self.GRAPH, unique_id="1") + node.cached_runs.record_item.assert_called_once_with("1", (" ",), Facts(None)) + + if __name__ == "__main__": unittest.main() diff --git a/utils/cached_runs.py b/utils/cached_runs.py index 7c5e7ecf..fd73cd2e 100644 --- a/utils/cached_runs.py +++ b/utils/cached_runs.py @@ -274,9 +274,17 @@ class CachedRuns: self._attempt = attempt def end_attempt(self, item_id: Any) -> None: + """The job finished: its receipts are complete and replace equal older ones. + + Deduplicating only complete receipts keeps a list execution that is + still recording ([A] so far of [A, B]) from dropping an older [A]. + """ with self._lock: if self._attempt is not None and self._attempt.item_id == item_id: self._attempt = None + for run_key, receipts in self._receipts.items(): + if receipts and receipts[0].item_id == item_id: + self._receipts[run_key] = _deduplicated(receipts) def _current(self) -> Tuple[Optional[_Attempt], Optional[int]]: """The attempt the executing node belongs to, and its list index.""" @@ -320,13 +328,12 @@ class CachedRuns: if receipts and receipts[0].item_id == item_id: receipt = receipts[0] else: + # The previous job's receipt is complete (its task_done may + # never have fired): it replaces older receipts it equals. + receipts = _deduplicated(receipts) receipt = _Receipt(item_id) receipts.insert(0, receipt) receipt.items[index] = item - outputs = receipt.outputs() - receipts = [receipt] + [ - older for older in receipts[1:] if older.outputs() != outputs - ] self._receipts[run_key] = receipts[: self._max_receipts] while len(self._receipts) > self._max_keys: self._receipts.popitem(last=False) @@ -363,6 +370,8 @@ class CachedRuns: ) continue with self._lock: + if node_id in attempt.replayed: + continue # a repeated message or id never counts twice attempt.replayed.add(node_id) counts: List[Count] = [] role = graph_role(attempt.graph, node_id, node_type) @@ -412,6 +421,14 @@ class CachedRuns: return [matches[0].items[index] for index in range(length)] +def _deduplicated(receipts: List[_Receipt]) -> List[_Receipt]: + """The newest receipt replaces older ones with the same outputs.""" + if not receipts: + return receipts + outputs = receipts[0].outputs() + return [receipts[0]] + [r for r in receipts[1:] if r.outputs() != outputs] + + def _job_nodes(graph: Dict[str, Any], outputs: Any) -> Dict[str, Tuple[str, str]]: """PromptManager-family nodes the job's outputs need, with their run keys.""" starts = (