diff --git a/tests/test_file_prompt.py b/tests/test_file_prompt.py index 79c944f0..bceea11c 100644 --- a/tests/test_file_prompt.py +++ b/tests/test_file_prompt.py @@ -82,6 +82,43 @@ def _gated(graph): return graph +def _prompt_manager_branches(prepend="PRE", append="APP"): + """Two KSampler branches, each fed by a PromptManagerText string.""" + graph = sampled_graph() + del graph["1"] + sampler = graph.pop("3") + for index, text in enumerate(("a cat", "a dog"), start=1): + base = index * 10 + graph[str(base)] = { + "class_type": "PromptManagerText", + "inputs": {"text": text, "prepend_text": prepend, "append_text": append}, + } + graph[str(base + 1)] = { + "class_type": "CLIPTextEncode", + "inputs": {"text": [str(base), 0]}, + } + graph[str(base + 2)] = { + **sampler, + "inputs": {**sampler["inputs"], "positive": [str(base + 1), 0]}, + } + return graph + + +def _numeric_first_branches(): + """Three branches; the legacy parser's answer (the first) is a number.""" + graph = sampled_graph("2048") + for index, text in ((5, "a cat"), (6, "a dog")): + graph[str(index * 10)] = { + "class_type": "CLIPTextEncode", + "inputs": {"text": text}, + } + graph[str(index * 10 + 1)] = { + **graph["3"], + "inputs": {**graph["3"]["inputs"], "positive": [str(index * 10), 0]}, + } + return graph + + class TestStepsBeforeResolving(unittest.TestCase): def test_step_1_no_metadata_tries_the_recorded_run(self): for metadata in ({}, None): @@ -114,6 +151,19 @@ class TestStepsBeforeResolving(unittest.TestCase): ) self.assertNotIn("file_info", _decide(_graph(sampled_graph())).canonical) + def test_a_file_size_reader_is_called_only_for_prompt_metadata(self): + reads = [] + + def size(): + reads.append(1) + return 9 + + for metadata in ({}, {"workflow": "{}"}): + _decide(metadata, file_size=size) + self.assertEqual(reads, []) + decision = _decide(_graph(sampled_graph()), file_size=size) + self.assertEqual((reads, decision.canonical["file_info"]), ([1], {"size": 9})) + class TestResolvedSteps(unittest.TestCase): def test_step_6_a_blank_positive_counts_as_unresolved(self): @@ -184,8 +234,10 @@ class TestUndecidedGraph(unittest.TestCase): self.assertEqual(decision.source, "legacy-parser") def test_step_8_1_the_file_name_anchor_alone_never_decides(self): - """#390 ruling: a file only its saver's file name resolves is queued - for Needs review by the scan, never linked; the monitor leaves it.""" + """#390: the file-name anchor never decides a link. The scan's legacy + floor may still link the file (as before #247), and a file nothing + links is queued with the anchor's choice per #390; the monitor + stores nothing.""" name = "s45_branch_2_00001_.png" scan = _decide(_graph(shapes.branch_bench()), lookup=_Lookup(), file_name=name) self.assertEqual((scan.kind, scan.source), ("positive", "legacy-parser")) @@ -231,8 +283,26 @@ class TestUndecidedGraph(unittest.TestCase): ) self.assertEqual(decision.kind, "never") self.assertEqual(decision.unlinked_status, "ambiguous") + self.assertFalse(decision.try_recorded_run) self.assertEqual(lookup.calls, 1) + def test_step_8_4_joined_prompt_manager_text_is_never_the_legacy_answer(self): + """The legacy parser returns prepend + text + append; the stored text + of a PromptManager node cannot be told from it.""" + decision = _decide(_graph(_prompt_manager_branches()), lookup=_Lookup()) + self.assertEqual(decision.kind, "never") + self.assertEqual(decision.unlinked_status, "ambiguous") + self.assertEqual(decision.review, "ambiguous") + self.assertTrue(decision.try_recorded_run) + linked = _prompt_manager_branches(prepend=["99", 0], append="") + self.assertEqual(_decide(_graph(linked), lookup=_Lookup()).kind, "never") + + def test_step_8_4_prompt_manager_text_without_parts_keeps_the_floor(self): + """Regression guard (passes on the base): no prepend or append.""" + graph = _prompt_manager_branches(prepend="", append=" ") + decision = _decide(_graph(graph), lookup=_Lookup()) + self.assertEqual((decision.kind, decision.text), ("positive", "a cat")) + def test_step_8_4_an_ambiguous_file_gets_the_legacy_answer_in_the_scan(self): decision = _decide(_graph(shapes.shared_prefix()), lookup=_Lookup()) self.assertEqual(decision.kind, "positive") @@ -358,6 +428,7 @@ class TestJunkGate(unittest.TestCase): "nan", "a portrait saved as final.safetensors", "4k", + "C:\\My Models\\x.safetensors", ): with self.subTest(text=text): self.assertIsNone(junk_text(text)) @@ -381,6 +452,16 @@ class TestJunkGate(unittest.TestCase): decision = _decide(_graph(_gated(sampled_graph(LORA))), lookup=_Lookup()) self.assertEqual(decision.source, "junk:lora-only") self.assertEqual(decision.kind, "never") + self.assertEqual(decision.unlinked_status, "unresolved") + + def test_step_13_junk_at_step_8_keeps_that_steps_status_and_review(self): + # An ambiguous file whose legacy answer is a number is still a + # Needs review case: its status stays "ambiguous". + decision = _decide(_graph(_numeric_first_branches()), lookup=_Lookup()) + self.assertEqual((decision.kind, decision.source), ("never", "junk:number")) + self.assertEqual(decision.unlinked_status, "ambiguous") + self.assertEqual(decision.review, "ambiguous") + self.assertTrue(decision.try_recorded_run) class TestParseComfyuiPrompt(unittest.TestCase): diff --git a/utils/file_prompt.py b/utils/file_prompt.py index b0b6b81d..ae6bb4f2 100644 --- a/utils/file_prompt.py +++ b/utils/file_prompt.py @@ -20,7 +20,7 @@ import json import re from dataclasses import dataclass, field from enum import Enum -from typing import Any, Callable, Mapping, Optional +from typing import Any, Callable, Mapping, Optional, Union from .generation_params import ( PLACEHOLDERS, @@ -95,7 +95,7 @@ def decide_file_prompt( metadata: Any, *, file_name: Optional[str], - file_size: Optional[int], + file_size: Union[None, int, Callable[[], Optional[int]]], in_output: bool, lookup: Optional[Callable[[], Any]], policy: Policy, @@ -106,6 +106,8 @@ def decide_file_prompt( (the monitor). ``lookup`` returns the file's ``RecordedOutput`` or None; it is None when history must not be asked yet. It is called at most once, and only for a graph the default profile leaves undecided. + ``file_size`` may be a callable: it is read only for a file with prompt + metadata whose ``file_info`` has no size. """ if not metadata: return Decision(NEVER, try_recorded_run=True) @@ -127,7 +129,7 @@ def decide_file_prompt( lookup, policy, ) - return _gated(_decide(found)) + return _gated(_decide(found), found) def _decide(found: _File) -> Decision: @@ -165,7 +167,11 @@ def _undecided_graph(found: _File) -> Decision: if found.lookup is None and found.in_output: return _decision(found, HISTORY, None, "none") if not inspected.malformed and inspected.resolution["state"] == "ambiguous": - return _never(found, _unlinked(found), try_recorded_run=True) + return _never(found, _unlinked(found)) + if _joins_prompt_manager_parts(found.canonical["prompt"]): + # The legacy parser joins prepend + text + append: the text a + # PromptManager node stores cannot be told from it (#343). + return _never(found, _unlinked(found), try_recorded_run=True) text = graph_prompt_text(found.metadata) if text is not None: return _decision(found, POSITIVE, text, "legacy-parser") @@ -187,8 +193,35 @@ def _sampler_less(found: _File, graph: Any) -> Decision: return _never(found, status, try_recorded_run=True) -def _gated(decision: Decision) -> Decision: - """Junk text is never stored; the file reports ``found`` (step 13).""" +_UNDECIDED_SOURCES = ("inspector", "recorded-run", "legacy-parser") + + +def _joins_prompt_manager_parts(graph: Any) -> bool: + """Whether a PromptManager-family node has a prepend or append part.""" + if not isinstance(graph, dict): + return False + for node in graph.values(): + if not isinstance(node, dict): + continue + name, inputs = node.get("class_type"), node.get("inputs") + if not isinstance(name, str) or "promptmanager" not in name.lower(): + continue + if not isinstance(inputs, dict): + continue + for key in ("prepend_text", "append_text"): + part = inputs.get(key) + if part is not None and not (isinstance(part, str) and not part.strip()): + return True + return False + + +def _gated(decision: Decision, found: _File) -> Decision: + """Junk text is never stored (step 13). + + A junk answer of the undecided-graph step keeps that step's unlinked + status, so an ambiguous file is still a Needs review case; any other + step reports ``found``. + """ if decision.kind not in (PAIR, POSITIVE): return decision junk = junk_text(decision.text) @@ -201,7 +234,9 @@ def _gated(decision: Decision) -> Decision: source=f"junk:{junk}", review=decision.review, try_recorded_run=True, - unlinked_status="found", + unlinked_status=( + _unlinked(found) if decision.source in _UNDECIDED_SOURCES else "found" + ), ) @@ -263,7 +298,7 @@ def _parsed(metadata: Mapping) -> dict: return canonical -def _with_size(canonical: dict, file_size: Optional[int]) -> None: +def _with_size(canonical: dict, file_size: Any) -> None: """Record the file's size in ``file_info`` unless the metadata carries one. Stored as ``generated_images.file_size``; the "new" scan scope skips a @@ -271,8 +306,10 @@ def _with_size(canonical: dict, file_size: Optional[int]) -> None: """ info = canonical.get("file_info") info = dict(info) if isinstance(info, dict) else {} - if info.get("size") is None and file_size is not None: - info["size"] = file_size + if info.get("size") is None: + size = file_size() if callable(file_size) else file_size + if size is not None: + info["size"] = size if info: canonical["file_info"] = info