6161
6262from __future__ import annotations
6363
64+ import hashlib
6465import json
6566import os
6667import tempfile
7172from pathlib import Path
7273from typing import Any
7374
75+ from comfy_cli .file_utils import atomic_write_text , cache_dir
7476from comfy_cli .http import plain_urlopen , read_capped
7577
7678# Public source of truth. Both files are at the repo root on ``main``.
101103# miss, so an old-format file is re-fetched rather than misread.
102104_CACHE_SCHEMA = 1
103105
106+ # Bump when parse_supported_nodes / parse_disable_config change their output shape.
107+ _PARSED_SCHEMA = 1
108+
104109# These are two small YAML documents (~32 KB bundled). The shared 64 MiB default
105110# is a ceiling for ``/object_info``-sized payloads; a much tighter cap here means
106111# a misbehaving or hostile upstream can't stream unbounded data into the memory
@@ -119,9 +124,17 @@ def network_disabled() -> bool:
119124 return os .environ .get ("COMFY_CLI_NO_REMOTE_REFRESH" , "" ).strip ().lower () not in _FALSEY
120125
121126
127+ def index_cache_disabled () -> bool :
128+ """True when ``COMFY_NO_CACHE`` opts out of the parsed-annotation cache.
129+
130+ Scope is :func:`parsed_annotations` only: the raw YAML pair cache and the
131+ object_info cache are unaffected.
132+ """
133+ return os .environ .get ("COMFY_NO_CACHE" , "" ).strip ().lower () not in _FALSEY
134+
135+
122136def _cache_dir () -> Path :
123- base = os .environ .get ("XDG_CACHE_HOME" ) or os .path .expanduser ("~/.cache" )
124- return Path (base ) / "comfy-cli" / "comfy-complete"
137+ return cache_dir () / "comfy-complete"
125138
126139
127140def _cache_file () -> Path :
@@ -328,10 +341,10 @@ def _drop_legacy_cache() -> None:
328341 Harmless if left behind — nothing reads it any more — but it's dead bytes in
329342 the user's cache dir, so clean it up the first time we write the new one.
330343 """
331- cache_dir = _cache_dir ()
344+ directory = _cache_dir ()
332345 for filename in _FILES :
333346 try :
334- (cache_dir / filename ).unlink ()
347+ (directory / filename ).unlink ()
335348 except OSError :
336349 pass
337350
@@ -351,18 +364,33 @@ def _read_cached_pair(*, require_fresh: bool) -> dict[str, bytes] | None:
351364 """The cached pair, if it is complete, valid, and (optionally) fresh.
352365
353366 Returns ``None`` on anything short of that — missing file, unreadable,
354- non-JSON, wrong schema, an absent entry, or a body that fails its validator.
355- A cache written by an older comfy-cli (per-file, pre-validation) lands in
356- that bucket too, so it is simply re-fetched rather than trusted. Every
357- ``None`` hands the decision back to the caller's next fallback, which
358- eventually reaches the bundled snapshot — never a silent blank annotation.
367+ oversized, non-JSON, wrong schema, an absent entry, or (for a pair not
368+ already vouched for below) a body that fails its validator. A cache written by an
369+ older comfy-cli (per-file, pre-validation) lands in that bucket too, so it
370+ is simply re-fetched rather than trusted. Every ``None`` hands the
371+ decision back to the caller's next fallback, which eventually reaches the
372+ bundled snapshot — never a silent blank annotation.
373+
374+ ``_persist_pair`` only ever writes a pair that passed ``_VALIDATORS`` in
375+ ``_fetch_one``, so a *freshly written* file is valid by construction. But
376+ ``COMFY_CACHE_DIR`` makes this file easy to point at or hand-edit, so that
377+ guarantee doesn't hold forever. The full YAML validation this used to run
378+ on every read is skipped only once ``parsed_annotations`` has already
379+ parsed this exact pair successfully — its cache file existing on disk is
380+ proof the bytes are good — so the hot path stays cheap while a poisoned
381+ cache is still caught the first time it's seen.
359382 """
360383 cache = _cache_file ()
361384 if require_fresh and not _is_fresh (cache ):
362385 return None
363386 try :
387+ # Same reasoning as ``_read_parsed``'s size cap: this path isn't even
388+ # digest-derived, so anything able to write the cache dir can plant an
389+ # arbitrarily large or deeply-nested file here.
390+ if cache .stat ().st_size > _MAX_ANNOTATION_BYTES :
391+ return None
364392 payload = json .loads (cache .read_bytes ())
365- except (OSError , ValueError ):
393+ except (OSError , ValueError , RecursionError ):
366394 return None
367395 if not isinstance (payload , dict ) or payload .get ("schema" ) != _CACHE_SCHEMA :
368396 return None
@@ -375,13 +403,125 @@ def _read_cached_pair(*, require_fresh: bool) -> dict[str, bytes] | None:
375403 text = files .get (filename )
376404 if not isinstance (text , str ):
377405 return None
378- data = text .encode ("utf-8" )
379- if not _VALIDATORS [filename ](data ):
380- return None
381- out [filename ] = data
406+ out [filename ] = text .encode ("utf-8" )
407+
408+ if not _parsed_cache_path (out [_SUPPORTED_NODES ], out [_CLOUD_DISABLE ]).exists ():
409+ for filename in _FILES :
410+ if not _VALIDATORS [filename ](out [filename ]):
411+ return None
382412 return out
383413
384414
415+ # ---------------------------------------------------------------------------
416+ # Parsed-annotation cache
417+ # ---------------------------------------------------------------------------
418+
419+ _PARSED_PREFIX = "annotations-parsed-"
420+ _PARSED_GLOB = f"{ _PARSED_PREFIX } *.json"
421+
422+ ParsedAnnotations = tuple [dict [str , str ], dict [str , list [str ]], set [str ]]
423+
424+
425+ def _parsed_cache_path (sup : bytes , dis : bytes ) -> Path :
426+ key = hashlib .sha256 (sup + b"\0 " + dis ).hexdigest ()[:16 ]
427+ return _cache_dir () / f"{ _PARSED_PREFIX } v{ _PARSED_SCHEMA } -{ key } .json"
428+
429+
430+ def _read_parsed (path : Path ) -> ParsedAnnotations | None :
431+ try :
432+ # The filename is a predictable digest over public upstream YAML, so
433+ # anything able to write the cache dir can plant a huge file at the
434+ # exact path this will read; cap it the same way network reads are
435+ # capped rather than handing an unbounded read to ``json.loads``.
436+ if path .stat ().st_size > _MAX_ANNOTATION_BYTES :
437+ return None
438+ payload = json .loads (path .read_bytes ())
439+ node_pack = payload ["node_pack" ]
440+ node_labels = payload ["node_labels" ]
441+ disable_labels = payload ["disable_labels" ]
442+ if not (
443+ isinstance (node_pack , dict )
444+ and isinstance (node_labels , dict )
445+ and isinstance (disable_labels , list )
446+ and all (isinstance (pack , str ) for pack in node_pack .values ())
447+ and all (
448+ isinstance (labels , list ) and all (isinstance (label , str ) for label in labels )
449+ for labels in node_labels .values ()
450+ )
451+ and all (isinstance (label , str ) for label in disable_labels )
452+ ):
453+ return None
454+ return node_pack , node_labels , set (disable_labels )
455+ except (OSError , ValueError , KeyError , TypeError , RecursionError ):
456+ return None
457+
458+
459+ def _write_parsed (path : Path , parsed : ParsedAnnotations ) -> None :
460+ """Persist ``parsed`` at ``path``, evicting every other parsed file after.
461+
462+ Newest-only keeps the cache dir from accumulating one file per annotation
463+ generation. Writing first and evicting siblings afterward means a failed
464+ write leaves the previous good entry in place instead of an empty dir.
465+ Best-effort throughout.
466+ """
467+ node_pack , node_labels , disable_labels = parsed
468+ # ``json.dumps`` coerces a non-str dict key (int, float, bool, None) to a
469+ # string instead of raising, which would silently change what a later
470+ # cache hit returns compared to this cache miss. Treat it like any other
471+ # unserialisable result: skip the write, keep re-parsing.
472+ if not (all (isinstance (k , str ) for k in node_pack ) and all (isinstance (k , str ) for k in node_labels )):
473+ return
474+ try :
475+ blob = json .dumps (
476+ {"node_pack" : node_pack , "node_labels" : node_labels , "disable_labels" : sorted (disable_labels )}
477+ )
478+ atomic_write_text (path , blob )
479+ except (OSError , TypeError , ValueError ):
480+ return
481+ for sibling in path .parent .glob (_PARSED_GLOB ):
482+ if sibling != path :
483+ try :
484+ sibling .unlink ()
485+ except OSError :
486+ pass
487+
488+
489+ def parsed_annotations (
490+ supported_nodes_yaml : bytes | None ,
491+ cloud_disable_yaml : bytes | None ,
492+ ) -> ParsedAnnotations :
493+ """``(node_pack, node_labels, disable_labels)`` for the given annotation bytes.
494+
495+ The parse result is stored on disk as JSON keyed by a digest of the input
496+ bytes, so the YAML is parsed once per upstream generation rather than on
497+ every call. A new pair gets a new key and the previous file is evicted on
498+ write. ``COMFY_NO_CACHE`` skips the disk entirely. Never raises: a missing,
499+ unreadable or malformed cache file is a miss and the documents are parsed
500+ again.
501+ """
502+ sup = supported_nodes_yaml or b""
503+ dis = cloud_disable_yaml or b""
504+ if not sup and not dis :
505+ return {}, {}, set ()
506+ use_cache = not index_cache_disabled ()
507+ path = _parsed_cache_path (sup , dis ) if use_cache else None
508+
509+ if path is not None :
510+ cached = _read_parsed (path )
511+ if cached is not None :
512+ return cached
513+
514+ from comfy_cli .cql .engine import parse_disable_config , parse_supported_nodes
515+
516+ node_pack , node_labels = parse_supported_nodes (sup ) if sup else ({}, {})
517+ disable_labels = parse_disable_config (dis ) if dis else set ()
518+ parsed : ParsedAnnotations = (node_pack , node_labels , disable_labels )
519+
520+ if path is not None :
521+ _write_parsed (path , parsed )
522+ return parsed
523+
524+
385525def _failure_stamp () -> Path :
386526 return _cache_dir () / _FAILURE_STAMP
387527
0 commit comments