diff --git a/.github/workflows/python-tests.yml b/.github/workflows/python-tests.yml index cb2847f5..79dcfd79 100644 --- a/.github/workflows/python-tests.yml +++ b/.github/workflows/python-tests.yml @@ -7,6 +7,28 @@ on: branches: [ main, master ] jobs: + quality: + runs-on: ubuntu-22.04 + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Install quality tools + run: | + python -m pip install --upgrade pip + pip install -r requirements.txt + pip install xenon + + - name: Check Python files encoding (no UTF-8 BOM) + run: python scripts/check_no_bom.py mjr_am_backend mjr_am_shared tests __init__.py scripts + + - name: Complexity gate (stage A) + run: python -m xenon -b F -m D -a C mjr_am_backend mjr_am_shared tests __init__.py + test: runs-on: ubuntu-22.04 strategy: diff --git a/.gitignore b/.gitignore index f01877cd..9cd95edc 100644 --- a/.gitignore +++ b/.gitignore @@ -132,3 +132,18 @@ run_tests_comfy_cli.bat scripts/sync_and_test_comfy_cli.bat .codex_filelist.txt .codex_existing_tests.txt +reports/complexity_a_remaining.txt +reports/complexity_ab_remaining.txt +reports/complexity_abcd.json +reports/complexity_abcd.txt +reports/complexity_b_remaining.txt +reports/complexity_bc.json +reports/complexity_bc.txt +reports/complexity_c_only_fresh.txt +reports/complexity_c_only.txt +reports/radon_cc_backend.json +reports/radon_cc_pair.json +reports/radon_cc_parser.json +reports/radon_cc_triplet.json +reports/radon_cc_watcher.json +uv.lock diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 00000000..909a3aad --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,28 @@ +minimum_pre_commit_version: "3.6.0" + +repos: + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v5.0.0 + hooks: + - id: trailing-whitespace + - id: end-of-file-fixer + - id: mixed-line-ending + - id: check-yaml + - id: check-json + - id: check-merge-conflict + - id: check-added-large-files + + - repo: local + hooks: + - id: check-no-bom + name: Check Python files encoding (no UTF-8 BOM) + entry: uv run python scripts/check_no_bom.py mjr_am_backend mjr_am_shared tests __init__.py scripts + language: system + pass_filenames: false + + - id: mypy + name: Mypy (manual / pre-push) + entry: uv run python -m mypy --config-file mypy.ini + language: system + pass_filenames: false + stages: [pre-push, manual] diff --git a/__init__.py b/__init__.py index 2caf667c..1cbf1895 100644 --- a/__init__.py +++ b/__init__.py @@ -1,4 +1,4 @@ -""" +""" ComfyUI-Majoor-AssetsManager Advanced asset browser for ComfyUI with ratings, tags, and metadata management. """ diff --git a/mjr_am_backend/adapters/db/schema.py b/mjr_am_backend/adapters/db/schema.py index ab110bc1..44ac206d 100644 --- a/mjr_am_backend/adapters/db/schema.py +++ b/mjr_am_backend/adapters/db/schema.py @@ -3,7 +3,7 @@ Database schema and migrations. """ import hashlib import re -from typing import List +from typing import List, Tuple from ...shared import Result, get_logger, log_success @@ -292,6 +292,87 @@ async def _fts_has_column(db, table: str, col: str) -> bool: return False +async def _sqlite_object_sql(db, obj_type: str, name: str) -> str: + try: + row = await db.aquery( + "SELECT sql FROM sqlite_master WHERE type=? AND name=? LIMIT 1", + (obj_type, name), + ) + if row.ok and row.data: + return str(row.data[0].get("sql") or "") + except Exception: + return "" + return "" + + +async def _asset_metadata_fts_needs_repair(db) -> Tuple[bool, bool]: + ddl_lower = (await _sqlite_object_sql(db, "table", "asset_metadata_fts")).lower() + trig_lower = (await _sqlite_object_sql(db, "trigger", "asset_metadata_fts_update")).lower() + needs_table_rebuild = ("content_rowid" in ddl_lower and "asset_id" in ddl_lower) + needs_trigger_rebuild = ("update asset_metadata_fts" in trig_lower) + missing_tags_text = not await _fts_has_column(db, "asset_metadata_fts", "tags_text") + if missing_tags_text: + needs_table_rebuild = True + return needs_table_rebuild, needs_trigger_rebuild + + +async def _drop_asset_metadata_fts_objects(db, *, include_table: bool) -> None: + script_parts = [ + "DROP TRIGGER IF EXISTS asset_metadata_fts_insert;", + "DROP TRIGGER IF EXISTS asset_metadata_fts_delete;", + "DROP TRIGGER IF EXISTS asset_metadata_fts_update;", + ] + if include_table: + script_parts.append("DROP TABLE IF EXISTS asset_metadata_fts;") + await db.aexecutescript("\n".join(script_parts)) + + +async def _create_asset_metadata_fts_table_if_needed(db, *, needs_table_rebuild: bool) -> None: + if not needs_table_rebuild: + return + await db.aexecutescript( + """ + CREATE VIRTUAL TABLE IF NOT EXISTS asset_metadata_fts USING fts5( + tags, + tags_text, + content='' + ); + """ + ) + + +async def _create_asset_metadata_fts_triggers(db) -> None: + await db.aexecutescript( + """ + CREATE TRIGGER IF NOT EXISTS asset_metadata_fts_insert AFTER INSERT ON asset_metadata BEGIN + INSERT INTO asset_metadata_fts(rowid, tags, tags_text) + VALUES (new.asset_id, COALESCE(new.tags, ''), COALESCE(new.tags_text, '')); + END; + + CREATE TRIGGER IF NOT EXISTS asset_metadata_fts_delete AFTER DELETE ON asset_metadata BEGIN + INSERT INTO asset_metadata_fts(asset_metadata_fts, rowid) VALUES('delete', old.asset_id); + END; + + CREATE TRIGGER IF NOT EXISTS asset_metadata_fts_update AFTER UPDATE ON asset_metadata BEGIN + INSERT INTO asset_metadata_fts(asset_metadata_fts, rowid) VALUES('delete', old.asset_id); + INSERT INTO asset_metadata_fts(rowid, tags, tags_text) + VALUES (new.asset_id, COALESCE(new.tags, ''), COALESCE(new.tags_text, '')); + END; + """ + ) + + +async def _reindex_asset_metadata_fts(db) -> None: + await db.aexecutescript( + """ + DELETE FROM asset_metadata_fts; + INSERT INTO asset_metadata_fts(rowid, tags, tags_text) + SELECT asset_id, COALESCE(tags, ''), COALESCE(tags_text, '') + FROM asset_metadata; + """ + ) + + async def _repair_asset_metadata_fts(db) -> Result[bool]: """ Repair legacy/incorrect FTS definition for asset metadata. @@ -301,25 +382,7 @@ async def _repair_asset_metadata_fts(db) -> Result[bool]: Also check for missing `tags_text` column in existing FTS tables. """ try: - table_row = await db.aquery("SELECT sql FROM sqlite_master WHERE type='table' AND name='asset_metadata_fts' LIMIT 1") - ddl = "" - if table_row.ok and table_row.data: - ddl = str(table_row.data[0].get("sql") or "") - ddl_lower = ddl.lower() - - trig_row = await db.aquery("SELECT sql FROM sqlite_master WHERE type='trigger' AND name='asset_metadata_fts_update' LIMIT 1") - trig_sql = "" - if trig_row.ok and trig_row.data: - trig_sql = str(trig_row.data[0].get("sql") or "") - trig_lower = trig_sql.lower() - - needs_table_rebuild = ("content_rowid" in ddl_lower and "asset_id" in ddl_lower) - needs_trigger_rebuild = ("update asset_metadata_fts" in trig_lower) - missing_tags_text = not await _fts_has_column(db, "asset_metadata_fts", "tags_text") - - if missing_tags_text: - needs_table_rebuild = True - + needs_table_rebuild, needs_trigger_rebuild = await _asset_metadata_fts_needs_repair(db) if not (needs_table_rebuild or needs_trigger_rebuild): return Result.Ok(True) except Exception: @@ -331,60 +394,10 @@ async def _repair_asset_metadata_fts(db) -> Result[bool]: async with db.atransaction(mode="immediate") as tx: if not tx.ok: return Result.Err("DB_ERROR", tx.error or "Failed to begin transaction") - if needs_table_rebuild: - await db.aexecutescript( - """ - DROP TRIGGER IF EXISTS asset_metadata_fts_insert; - DROP TRIGGER IF EXISTS asset_metadata_fts_delete; - DROP TRIGGER IF EXISTS asset_metadata_fts_update; - DROP TABLE IF EXISTS asset_metadata_fts; - """ - ) - await db.aexecutescript( - """ - CREATE VIRTUAL TABLE IF NOT EXISTS asset_metadata_fts USING fts5( - tags, - tags_text, - content='' - ); - """ - ) - else: - await db.aexecutescript( - """ - DROP TRIGGER IF EXISTS asset_metadata_fts_insert; - DROP TRIGGER IF EXISTS asset_metadata_fts_delete; - DROP TRIGGER IF EXISTS asset_metadata_fts_update; - """ - ) - - await db.aexecutescript( - """ - CREATE TRIGGER IF NOT EXISTS asset_metadata_fts_insert AFTER INSERT ON asset_metadata BEGIN - INSERT INTO asset_metadata_fts(rowid, tags, tags_text) - VALUES (new.asset_id, COALESCE(new.tags, ''), COALESCE(new.tags_text, '')); - END; - - CREATE TRIGGER IF NOT EXISTS asset_metadata_fts_delete AFTER DELETE ON asset_metadata BEGIN - INSERT INTO asset_metadata_fts(asset_metadata_fts, rowid) VALUES('delete', old.asset_id); - END; - - CREATE TRIGGER IF NOT EXISTS asset_metadata_fts_update AFTER UPDATE ON asset_metadata BEGIN - INSERT INTO asset_metadata_fts(asset_metadata_fts, rowid) VALUES('delete', old.asset_id); - INSERT INTO asset_metadata_fts(rowid, tags, tags_text) - VALUES (new.asset_id, COALESCE(new.tags, ''), COALESCE(new.tags_text, '')); - END; - """ - ) - - await db.aexecutescript( - """ - DELETE FROM asset_metadata_fts; - INSERT INTO asset_metadata_fts(rowid, tags, tags_text) - SELECT asset_id, COALESCE(tags, ''), COALESCE(tags_text, '') - FROM asset_metadata; - """ - ) + await _drop_asset_metadata_fts_objects(db, include_table=needs_table_rebuild) + await _create_asset_metadata_fts_table_if_needed(db, needs_table_rebuild=needs_table_rebuild) + await _create_asset_metadata_fts_triggers(db) + await _reindex_asset_metadata_fts(db) if not tx.ok: return Result.Err("DB_ERROR", tx.error or "Commit failed") return Result.Ok(True) @@ -466,22 +479,7 @@ async def _repair_assets_fts(db) -> Result[bool]: triggers were not present or the FTS table got out of sync. """ try: - table_row = await db.aquery("SELECT sql FROM sqlite_master WHERE type='table' AND name='assets_fts' LIMIT 1") - ddl = "" - if table_row.ok and table_row.data: - ddl = str(table_row.data[0].get("sql") or "") - ddl_lower = ddl.lower() - - # Check existing trigger presence - trig_row = await db.aquery("SELECT name FROM sqlite_master WHERE type='trigger' AND name LIKE 'assets_fts_%'") - triggers_exist = trig_row.ok and bool(trig_row.data) - - needs_rebuild = False - # If virtual table ddl looks wrong or triggers missing -> rebuild - if not ddl_lower or "fts5" not in ddl_lower or not triggers_exist: - needs_rebuild = True - - if not needs_rebuild: + if not await _assets_fts_needs_rebuild(db): return Result.Ok(True) except Exception: return Result.Ok(True) @@ -492,52 +490,10 @@ async def _repair_assets_fts(db) -> Result[bool]: if not tx.ok: return Result.Err("DB_ERROR", tx.error or "Failed to begin transaction") - await db.aexecutescript( - """ - DROP TRIGGER IF EXISTS assets_fts_insert; - DROP TRIGGER IF EXISTS assets_fts_delete; - DROP TRIGGER IF EXISTS assets_fts_update; - DROP TABLE IF EXISTS assets_fts; - """ - ) - - await db.aexecutescript( - """ - CREATE VIRTUAL TABLE IF NOT EXISTS assets_fts USING fts5( - filename, - subfolder, - content='assets', - content_rowid='id' - ); - """ - ) - - await db.aexecutescript( - """ - CREATE TRIGGER IF NOT EXISTS assets_fts_insert AFTER INSERT ON assets BEGIN - INSERT INTO assets_fts(rowid, filename, subfolder) - VALUES (new.id, new.filename, new.subfolder); - END; - - CREATE TRIGGER IF NOT EXISTS assets_fts_delete AFTER DELETE ON assets BEGIN - DELETE FROM assets_fts WHERE rowid = old.id; - END; - - CREATE TRIGGER IF NOT EXISTS assets_fts_update AFTER UPDATE ON assets BEGIN - UPDATE assets_fts SET filename = new.filename, subfolder = new.subfolder - WHERE rowid = new.id; - END; - """ - ) - - # Rebuild FTS contents from assets table - await db.aexecutescript( - """ - DELETE FROM assets_fts; - INSERT INTO assets_fts(rowid, filename, subfolder) - SELECT id, COALESCE(filename, ''), COALESCE(subfolder, '') FROM assets; - """ - ) + await _drop_assets_fts_objects(db) + await _create_assets_fts_table(db) + await _create_assets_fts_triggers(db) + await _rebuild_assets_fts_content(db) if not tx.ok: return Result.Err("DB_ERROR", tx.error or "Commit failed") @@ -547,6 +503,67 @@ async def _repair_assets_fts(db) -> Result[bool]: return Result.Err("FTS_REPAIR_FAILED", str(exc)) +async def _assets_fts_needs_rebuild(db) -> bool: + ddl_lower = (await _sqlite_object_sql(db, "table", "assets_fts")).lower() + trig_row = await db.aquery("SELECT name FROM sqlite_master WHERE type='trigger' AND name LIKE 'assets_fts_%'") + triggers_exist = bool(trig_row.ok and trig_row.data) + return bool((not ddl_lower) or ("fts5" not in ddl_lower) or (not triggers_exist)) + + +async def _drop_assets_fts_objects(db) -> None: + await db.aexecutescript( + """ + DROP TRIGGER IF EXISTS assets_fts_insert; + DROP TRIGGER IF EXISTS assets_fts_delete; + DROP TRIGGER IF EXISTS assets_fts_update; + DROP TABLE IF EXISTS assets_fts; + """ + ) + + +async def _create_assets_fts_table(db) -> None: + await db.aexecutescript( + """ + CREATE VIRTUAL TABLE IF NOT EXISTS assets_fts USING fts5( + filename, + subfolder, + content='assets', + content_rowid='id' + ); + """ + ) + + +async def _create_assets_fts_triggers(db) -> None: + await db.aexecutescript( + """ + CREATE TRIGGER IF NOT EXISTS assets_fts_insert AFTER INSERT ON assets BEGIN + INSERT INTO assets_fts(rowid, filename, subfolder) + VALUES (new.id, new.filename, new.subfolder); + END; + + CREATE TRIGGER IF NOT EXISTS assets_fts_delete AFTER DELETE ON assets BEGIN + DELETE FROM assets_fts WHERE rowid = old.id; + END; + + CREATE TRIGGER IF NOT EXISTS assets_fts_update AFTER UPDATE ON assets BEGIN + UPDATE assets_fts SET filename = new.filename, subfolder = new.subfolder + WHERE rowid = new.id; + END; + """ + ) + + +async def _rebuild_assets_fts_content(db) -> None: + await db.aexecutescript( + """ + DELETE FROM assets_fts; + INSERT INTO assets_fts(rowid, filename, subfolder) + SELECT id, COALESCE(filename, ''), COALESCE(subfolder, '') FROM assets; + """ + ) + + async def init_schema(db) -> Result[bool]: """ Initialize the schema (useful for tests or first-time installs). diff --git a/mjr_am_backend/adapters/db/sqlite.py b/mjr_am_backend/adapters/db/sqlite.py index 4e06c39e..8cb727ba 100644 --- a/mjr_am_backend/adapters/db/sqlite.py +++ b/mjr_am_backend/adapters/db/sqlite.py @@ -107,6 +107,12 @@ _KNOWN_COLUMNS = { "m.has_generation_data", } _KNOWN_COLUMNS_LOWER = {c.lower(): c for c in _KNOWN_COLUMNS} +_SCHEMA_TABLE_ALIASES = { + "assets": "a", + "asset_metadata": "m", + "scan_journal": None, + "metadata_cache": None, +} def _is_missing_column_error(exc: Exception) -> bool: @@ -129,32 +135,34 @@ def _populate_known_columns_from_schema(db_path: Path) -> None: try: if not db_path.exists(): return - tables = { - "assets": "a", - "asset_metadata": "m", - "scan_journal": None, - "metadata_cache": None, - } with sqlite3.connect(str(db_path)) as conn: cursor = conn.cursor() - for table, alias in tables.items(): - try: - cursor.execute(f"PRAGMA table_info({table})") - for row in cursor.fetchall() or []: - if not row or len(row) < 2: - continue - col = str(row[1]).strip() - if not col: - continue - _KNOWN_COLUMNS.add(col) - if alias: - _KNOWN_COLUMNS.add(f"{alias}.{col}") - except Exception: - pass + for table, alias in _SCHEMA_TABLE_ALIASES.items(): + _populate_table_columns(cursor, table, alias) _KNOWN_COLUMNS_LOWER = {c.lower(): c for c in _KNOWN_COLUMNS} except Exception as exc: logger.warning("Failed to refresh known columns from schema: %s", exc) + +def _populate_table_columns(cursor: sqlite3.Cursor, table: str, alias: Optional[str]) -> None: + try: + cursor.execute(f"PRAGMA table_info({table})") + for row in cursor.fetchall() or []: + col = _extract_schema_column_name(row) + if not col: + continue + _KNOWN_COLUMNS.add(col) + if alias: + _KNOWN_COLUMNS.add(f"{alias}.{col}") + except Exception: + pass + + +def _extract_schema_column_name(row: Any) -> str: + if not row or len(row) < 2: + return "" + return str(row[1]).strip() + _IN_QUERY_FORBIDDEN = re.compile( r"(--|/\*|\*/|;|\bpragma\b|\battach\b|\bdetach\b|\bvacuum\b|\balter\b|\bdrop\b|\binsert\b|\bupdate\b|\bdelete\b)", re.IGNORECASE, @@ -362,13 +370,13 @@ class Sqlite: self._initialized = False self._lock = threading.Lock() self._async_sem: Optional[asyncio.Semaphore] = None - + # Reset mechanics self._resetting = False self._active_conns: set[aiosqlite.Connection] = set() self._active_conns_idle = threading.Event() self._active_conns_idle.set() - + self._timeout = float(user_config.get("timeout", timeout)) self._query_timeout = float(user_config.get("queryTimeout", DB_QUERY_TIMEOUT if DB_QUERY_TIMEOUT is not None else 0.0)) self._lock_retry_attempts = 6 @@ -425,43 +433,53 @@ class Sqlite: 2) /db_config.json """ try: - cfg_path_raw = os.getenv("MAJOOR_DB_CONFIG_PATH", "").strip() - cfg_path = Path(cfg_path_raw).expanduser() if cfg_path_raw else (self.db_path.parent / "db_config.json") - if not cfg_path.exists() or not cfg_path.is_file(): - return {} - - with cfg_path.open("r", encoding="utf-8") as f: - data = json.load(f) + cfg_path = self._resolve_user_db_config_path() + data = self._read_user_db_config_file(cfg_path) if not isinstance(data, dict): return {} - - out: Dict[str, Any] = {} - try: - if "timeout" in data: - timeout_raw = data.get("timeout") - if timeout_raw is not None: - out["timeout"] = max(1.0, float(timeout_raw)) - except Exception: - pass - try: - if "maxConnections" in data: - max_connections_raw = data.get("maxConnections") - if max_connections_raw is not None: - out["maxConnections"] = max(1, int(max_connections_raw)) - except Exception: - pass - try: - if "queryTimeout" in data: - query_timeout_raw = data.get("queryTimeout") - if query_timeout_raw is not None: - out["queryTimeout"] = max(0.0, float(query_timeout_raw)) - except Exception: - pass - return out + return self._normalize_user_db_config(data) except Exception as exc: logger.debug("Ignoring invalid DB user config: %s", exc) return {} + def _resolve_user_db_config_path(self) -> Path: + cfg_path_raw = os.getenv("MAJOOR_DB_CONFIG_PATH", "").strip() + return Path(cfg_path_raw).expanduser() if cfg_path_raw else (self.db_path.parent / "db_config.json") + + @staticmethod + def _read_user_db_config_file(cfg_path: Path) -> Any: + if not cfg_path.exists() or not cfg_path.is_file(): + return {} + with cfg_path.open("r", encoding="utf-8") as f: + return json.load(f) + + @staticmethod + def _normalize_user_db_config(data: Dict[str, Any]) -> Dict[str, Any]: + out: Dict[str, Any] = {} + Sqlite._maybe_set_config_number(out, data, "timeout", min_value=1.0, cast=float) + Sqlite._maybe_set_config_number(out, data, "maxConnections", min_value=1, cast=int) + Sqlite._maybe_set_config_number(out, data, "queryTimeout", min_value=0.0, cast=float) + return out + + @staticmethod + def _maybe_set_config_number( + out: Dict[str, Any], + data: Dict[str, Any], + key: str, + *, + min_value: float, + cast, + ) -> None: + try: + if key not in data: + return + raw = data.get(key) + if raw is None: + return + out[key] = max(min_value, cast(raw)) + except Exception: + return + def _is_locked_error(self, exc: Exception) -> bool: try: msg = str(exc).lower() @@ -517,17 +535,7 @@ class Sqlite: """ Auto-reset DB on corruption by default; can be disabled with MAJOOR_DB_AUTO_RESET=false. """ - try: - auto_reset = str(os.getenv("MAJOOR_DB_AUTO_RESET", "true")).strip().lower() in ( - "1", - "true", - "yes", - "on", - ) - except Exception: - auto_reset = False - - if not auto_reset: + if not self._is_auto_reset_enabled(): logger.error( "Corruption detected (%s). Auto-reset disabled (MAJOOR_DB_AUTO_RESET=false)", error, @@ -535,33 +543,57 @@ class Sqlite: return now = time.time() - with self._auto_reset_lock: - if now - float(self._auto_reset_last_ts or 0.0) < float(self._auto_reset_cooldown_s): - logger.warning("Corruption auto-reset throttled (cooldown %.1fs)", float(self._auto_reset_cooldown_s)) - return - self._auto_reset_last_ts = now - with self._diag_lock: - self._diag["auto_reset_attempts"] = int(self._diag.get("auto_reset_attempts") or 0) + 1 - self._diag["last_auto_reset_at"] = now - self._diag["last_auto_reset_error"] = None + if self._is_auto_reset_throttled(now): + logger.warning("Corruption auto-reset throttled (cooldown %.1fs)", float(self._auto_reset_cooldown_s)) + return + self._mark_auto_reset_attempt(now) try: reset_res = await self.areset() - if reset_res and reset_res.ok: - logger.warning("Database auto-reset completed after corruption") - with self._diag_lock: - self._diag["auto_reset_successes"] = int(self._diag.get("auto_reset_successes") or 0) + 1 - else: - logger.error("Database auto-reset failed after corruption: %s", getattr(reset_res, "error", "unknown error")) - with self._diag_lock: - self._diag["auto_reset_failures"] = int(self._diag.get("auto_reset_failures") or 0) + 1 - self._diag["last_auto_reset_error"] = str(getattr(reset_res, "error", "unknown error")) + self._record_auto_reset_result(reset_res) except Exception as exc: logger.error("Database auto-reset exception after corruption: %s", exc) with self._diag_lock: self._diag["auto_reset_failures"] = int(self._diag.get("auto_reset_failures") or 0) + 1 self._diag["last_auto_reset_error"] = str(exc) + @staticmethod + def _is_auto_reset_enabled() -> bool: + try: + return str(os.getenv("MAJOOR_DB_AUTO_RESET", "true")).strip().lower() in ( + "1", + "true", + "yes", + "on", + ) + except Exception: + return False + + def _is_auto_reset_throttled(self, now: float) -> bool: + with self._auto_reset_lock: + if now - float(self._auto_reset_last_ts or 0.0) < float(self._auto_reset_cooldown_s): + return True + self._auto_reset_last_ts = now + return False + + def _mark_auto_reset_attempt(self, now: float) -> None: + with self._diag_lock: + self._diag["auto_reset_attempts"] = int(self._diag.get("auto_reset_attempts") or 0) + 1 + self._diag["last_auto_reset_at"] = now + self._diag["last_auto_reset_error"] = None + + def _record_auto_reset_result(self, reset_res: Any) -> None: + if reset_res and reset_res.ok: + logger.warning("Database auto-reset completed after corruption") + with self._diag_lock: + self._diag["auto_reset_successes"] = int(self._diag.get("auto_reset_successes") or 0) + 1 + return + error_text = str(getattr(reset_res, "error", "unknown error")) + logger.error("Database auto-reset failed after corruption: %s", error_text) + with self._diag_lock: + self._diag["auto_reset_failures"] = int(self._diag.get("auto_reset_failures") or 0) + 1 + self._diag["last_auto_reset_error"] = error_text + def get_diagnostics(self) -> Dict[str, Any]: with self._diag_lock: data = dict(self._diag) @@ -789,37 +821,48 @@ class Sqlite: if not self._asset_locks: return cutoff = now - float(ASSET_LOCKS_TTL_S) + self._prune_asset_locks_by_ttl(cutoff) + self._prune_asset_locks_by_cap() + + def _prune_asset_locks_by_ttl(self, cutoff: float) -> None: try: for key, entry in list(self._asset_locks.items()): - try: - lock = entry.get("lock") - last = float(entry.get("last") or 0) - except Exception: + last = self._asset_lock_last(entry) + if not last or last >= cutoff: continue - try: - if lock is not None and getattr(lock, "locked", None) and lock.locked(): - continue - except Exception: - pass - if last and last < cutoff: - self._asset_locks.pop(key, None) + if self._asset_lock_is_locked(entry): + continue + self._asset_locks.pop(key, None) except Exception: - pass + return - if ASSET_LOCKS_MAX > 0 and len(self._asset_locks) > ASSET_LOCKS_MAX: - try: - items = sorted(self._asset_locks.items(), key=lambda kv: float(kv[1].get("last") or 0)) - to_drop = max(0, len(items) - ASSET_LOCKS_MAX) - for key, entry in items[:to_drop]: - try: - lock = entry.get("lock") - if lock is not None and getattr(lock, "locked", None) and lock.locked(): - continue - except Exception: - pass - self._asset_locks.pop(key, None) - except Exception: - pass + def _prune_asset_locks_by_cap(self) -> None: + if ASSET_LOCKS_MAX <= 0 or len(self._asset_locks) <= ASSET_LOCKS_MAX: + return + try: + items = sorted(self._asset_locks.items(), key=lambda kv: self._asset_lock_last(kv[1])) + to_drop = max(0, len(items) - ASSET_LOCKS_MAX) + for key, entry in items[:to_drop]: + if self._asset_lock_is_locked(entry): + continue + self._asset_locks.pop(key, None) + except Exception: + return + + @staticmethod + def _asset_lock_last(entry: Dict[str, Any]) -> float: + try: + return float(entry.get("last") or 0) + except Exception: + return 0.0 + + @staticmethod + def _asset_lock_is_locked(entry: Dict[str, Any]) -> bool: + try: + lock = entry.get("lock") + return bool(lock is not None and getattr(lock, "locked", None) and lock.locked()) + except Exception: + return False def _get_or_create_asset_lock(self, asset_id: Any) -> asyncio.Lock: key = _asset_lock_key(asset_id) @@ -1017,24 +1060,13 @@ class Sqlite: try: for attempt in range(self._lock_retry_attempts + 1): try: - cursor = await conn.execute(query, params or ()) - try: - if fetch: - rows = await cursor.fetchall() - return Result.Ok(self._rows_to_dicts(rows)) - if commit: - await conn.commit() - # aiosqlite cursor has lastrowid/rowcount, but may be None. - last_id = getattr(cursor, "lastrowid", None) - rowcount = getattr(cursor, "rowcount", None) - if last_id: - return Result.Ok(last_id) - return Result.Ok(rowcount if rowcount is not None else 0) - finally: - try: - await cursor.close() - except Exception: - pass + return await self._execute_with_cursor_result( + conn, + query, + params, + fetch=fetch, + commit=commit, + ) except sqlite3.OperationalError as exc: if self._is_locked_error(exc) and attempt < self._lock_retry_attempts: await self._sleep_backoff(attempt) @@ -1044,6 +1076,37 @@ class Sqlite: except Exception: raise + async def _execute_with_cursor_result( + self, + conn: aiosqlite.Connection, + query: str, + params: Optional[tuple], + *, + fetch: bool, + commit: bool, + ) -> Result[Any]: + cursor = await conn.execute(query, params or ()) + try: + if fetch: + rows = await cursor.fetchall() + return Result.Ok(self._rows_to_dicts(rows)) + if commit: + await conn.commit() + return self._cursor_write_result(cursor) + finally: + try: + await cursor.close() + except Exception: + pass + + @staticmethod + def _cursor_write_result(cursor: Any) -> Result[Any]: + last_id = getattr(cursor, "lastrowid", None) + rowcount = getattr(cursor, "rowcount", None) + if last_id: + return Result.Ok(last_id) + return Result.Ok(rowcount if rowcount is not None else 0) + async def aexecute(self, query: str, params: Optional[tuple] = None, fetch: bool = False) -> Result[Any]: """Execute SQL on the DB loop thread (async).""" token = _TX_TOKEN.get() @@ -1277,6 +1340,89 @@ class Sqlite: (str(version),), ) + @staticmethod + def _begin_stmt_for_mode(mode: str) -> str: + begin_stmt = "BEGIN IMMEDIATE" + if isinstance(mode, str) and mode.lower() in ("deferred", "immediate", "exclusive"): + begin_stmt = f"BEGIN {mode.upper()}" + elif isinstance(mode, str) and mode.strip() == "": + begin_stmt = "BEGIN" + return begin_stmt + + async def _begin_tx_with_retry(self, conn: aiosqlite.Connection, begin_stmt: str) -> None: + for attempt in range(self._lock_retry_attempts + 1): + try: + await conn.execute(begin_stmt) + return + except sqlite3.OperationalError as exc: + if self._is_locked_error(exc) and attempt < self._lock_retry_attempts: + await self._sleep_backoff(attempt) + continue + raise + + def _register_tx_token(self, token: str, conn: aiosqlite.Connection, lock: Optional[asyncio.Lock]) -> None: + with self._tx_state_lock: + self._tx_conns[token] = conn + if lock is not None: + self._tx_write_lock_tokens.add(token) + + async def _abort_begin_tx( + self, + conn: aiosqlite.Connection, + lock: Optional[asyncio.Lock], + ) -> None: + try: + await conn.rollback() + except Exception: + pass + await self._release_connection_async(conn) + if lock is not None: + try: + lock.release() + except Exception: + pass + + def _get_tx_state(self, token: str) -> tuple[Optional[aiosqlite.Connection], bool]: + try: + with self._tx_state_lock: + conn = self._tx_conns.get(token) + had_write_lock = token in self._tx_write_lock_tokens + except Exception: + return None, False + return conn, had_write_lock + + async def _commit_with_retry(self, conn: aiosqlite.Connection) -> None: + for attempt in range(self._lock_retry_attempts + 1): + try: + await conn.commit() + return + except sqlite3.OperationalError as exc: + if self._is_locked_error(exc) and attempt < self._lock_retry_attempts: + await self._sleep_backoff(attempt) + continue + raise + + async def _cleanup_tx_state( + self, + token: str, + conn: aiosqlite.Connection, + had_write_lock: bool, + lock: Optional[asyncio.Lock], + ) -> None: + try: + with self._tx_state_lock: + self._tx_conns.pop(token, None) + if had_write_lock: + self._tx_write_lock_tokens.discard(token) + except Exception: + pass + await self._release_connection_async(conn) + if lock is not None and had_write_lock: + try: + lock.release() + except Exception: + pass + async def _begin_tx_async(self, mode: str) -> Result[str]: await self._ensure_initialized_async() lock = self._write_lock @@ -1285,64 +1431,23 @@ class Sqlite: conn = await self._acquire_connection_async() token = f"tx_{uuid.uuid4().hex}" - begin_stmt = "BEGIN IMMEDIATE" - if isinstance(mode, str) and mode.lower() in ("deferred", "immediate", "exclusive"): - begin_stmt = f"BEGIN {mode.upper()}" - elif isinstance(mode, str) and mode.strip() == "": - begin_stmt = "BEGIN" + begin_stmt = self._begin_stmt_for_mode(mode) try: - for attempt in range(self._lock_retry_attempts + 1): - try: - await conn.execute(begin_stmt) - break - except sqlite3.OperationalError as exc: - if self._is_locked_error(exc) and attempt < self._lock_retry_attempts: - await self._sleep_backoff(attempt) - continue - raise - try: - with self._tx_state_lock: - self._tx_conns[token] = conn - if lock is not None: - self._tx_write_lock_tokens.add(token) - except Exception: - raise + await self._begin_tx_with_retry(conn, begin_stmt) + self._register_tx_token(token, conn, lock) return Result.Ok(token) except Exception as exc: - try: - await conn.rollback() - except Exception: - pass - await self._release_connection_async(conn) - if lock is not None: - try: - lock.release() - except Exception: - pass + await self._abort_begin_tx(conn, lock) return Result.Err(ErrorCode.DB_ERROR, str(exc)) async def _commit_tx_async(self, token: str) -> Result[bool]: lock = self._write_lock - try: - with self._tx_state_lock: - conn = self._tx_conns.get(token) - had_write_lock = token in self._tx_write_lock_tokens - except Exception: - conn = None - had_write_lock = False + conn, had_write_lock = self._get_tx_state(token) if not conn: return Result.Err(ErrorCode.DB_ERROR, "Transaction connection missing") try: - for attempt in range(self._lock_retry_attempts + 1): - try: - await conn.commit() - break - except sqlite3.OperationalError as exc: - if self._is_locked_error(exc) and attempt < self._lock_retry_attempts: - await self._sleep_backoff(attempt) - continue - raise + await self._commit_with_retry(conn) return Result.Ok(True) except Exception as exc: # Best-effort rollback to avoid returning a connection with an open transaction to the pool. @@ -1352,19 +1457,7 @@ class Sqlite: pass return Result.Err(ErrorCode.DB_ERROR, str(exc)) finally: - try: - with self._tx_state_lock: - self._tx_conns.pop(token, None) - if had_write_lock: - self._tx_write_lock_tokens.discard(token) - except Exception: - pass - await self._release_connection_async(conn) - if lock is not None and had_write_lock: - try: - lock.release() - except Exception: - pass + await self._cleanup_tx_state(token, conn, had_write_lock, lock) async def _rollback_tx_async(self, token: str) -> Result[bool]: lock = self._write_lock @@ -1420,7 +1513,7 @@ class Sqlite: await conn.close() except Exception: pass - + # 3. FORCE CLOSE any checked-out connections (critical for reset) active_list = list(self._active_conns) for conn in active_list: @@ -1457,7 +1550,7 @@ class Sqlite: drained = await asyncio.to_thread(self._active_conns_idle.wait, 5.0) if not drained and self._active_conns: logger.warning("DB Reset: Forced close with %d active connections", len(self._active_conns)) - + await self._close_all_async() async def _checkpoint_wal_before_reset_async(self) -> None: @@ -1518,61 +1611,12 @@ class Sqlite: try: deleted_files: list[str] = [] renamed_files: list[dict[str, str]] = [] - # 1. Drain and close on loop thread - fut = self._loop_thread.submit(self._drain_for_reset_async()) - await asyncio.wrap_future(fut) - - # 2. Wait a bit for OS to release file locks & FORCE GC - await asyncio.sleep(0.5) - gc.collect() - - # 3. Delete files with RETRY logic - base = str(self.db_path) - for ext in ["", "-wal", "-shm", "-journal"]: - p = Path(base + ext) - if not p.exists(): - continue - - deleted = False - last_error = None - for attempt in range(5): - try: - p.unlink() - deleted = True - deleted_files.append(str(p)) - break - except Exception as e: - last_error = e - await asyncio.sleep(0.5 + (attempt * 0.2)) - gc.collect() # aggressive gc - - if not deleted: - logger.warning(f"Failed to delete {p} after retries: {last_error}") - # Final Hail Mary: Try rename if delete fails (Windows hack). - # We do this for main DB and WAL/SHM to guarantee the old files disappear. - try: - trash = p.with_name(f"{p.name}.trash_{uuid.uuid4().hex}") - p.rename(trash) - renamed_files.append({"from": str(p), "to": str(trash)}) - logger.info(f"Renamed locked DB file to {trash}") - deleted = True - except Exception as rename_err: - # Final fallback (Windows): schedule delete at reboot. - # We still fail the reset to avoid reinitializing over potentially - # stale sidecar files in the current process/session. - scheduled = self._schedule_delete_on_reboot(p) - with self._lock: - self._resetting = False - if scheduled: - return Result.Err( - "DELETE_PENDING_REBOOT", - f"Could not delete {p} now; scheduled for deletion at reboot. " - f"Delete error: {last_error} | Rename failed: {rename_err}", - ) - return Result.Err( - "DELETE_FAILED", - f"Could not delete {p}: {last_error} | Rename failed: {rename_err}", - ) + await self._drain_for_reset() + delete_err = await self._delete_reset_files(deleted_files, renamed_files) + if delete_err is not None: + with self._lock: + self._resetting = False + return delete_err # 4. Reset initialization state with self._lock: @@ -1591,7 +1635,7 @@ class Sqlite: schema_res = await ensure_tables_exist(self) if not schema_res.ok: return schema_res - + indexes_res = await ensure_indexes_and_triggers(self) if not indexes_res.ok: return indexes_res @@ -1603,6 +1647,80 @@ class Sqlite: logger.error(f"DB Reset failed: {exc}") return Result.Err("RESET_FAILED", str(exc)) + async def _drain_for_reset(self) -> None: + # 1. Drain and close on loop thread. + fut = self._loop_thread.submit(self._drain_for_reset_async()) + await asyncio.wrap_future(fut) + # 2. Wait a bit for OS to release file locks & force GC. + await asyncio.sleep(0.5) + gc.collect() + + async def _delete_reset_files( + self, + deleted_files: list[str], + renamed_files: list[dict[str, str]], + ) -> Optional[Result[bool]]: + base = str(self.db_path) + for ext in ["", "-wal", "-shm", "-journal"]: + path = Path(base + ext) + if not path.exists(): + continue + err = await self._delete_or_recover_reset_file(path, deleted_files, renamed_files) + if err is not None: + return err + return None + + async def _delete_or_recover_reset_file( + self, + path: Path, + deleted_files: list[str], + renamed_files: list[dict[str, str]], + ) -> Optional[Result[bool]]: + deleted, last_error = await self._delete_reset_file_with_retries(path, deleted_files) + if deleted: + return None + logger.warning(f"Failed to delete {path} after retries: {last_error}") + recover = self._rename_or_schedule_delete(path, renamed_files, last_error) + return None if recover.ok else recover + + async def _delete_reset_file_with_retries(self, path: Path, deleted_files: list[str]) -> tuple[bool, Optional[Exception]]: + last_error: Optional[Exception] = None + for attempt in range(5): + try: + path.unlink() + deleted_files.append(str(path)) + return True, None + except Exception as exc: + last_error = exc + await asyncio.sleep(0.5 + (attempt * 0.2)) + gc.collect() + return False, last_error + + def _rename_or_schedule_delete( + self, + path: Path, + renamed_files: list[dict[str, str]], + last_error: Optional[Exception], + ) -> Result[bool]: + try: + trash = path.with_name(f"{path.name}.trash_{uuid.uuid4().hex}") + path.rename(trash) + renamed_files.append({"from": str(path), "to": str(trash)}) + logger.info(f"Renamed locked DB file to {trash}") + return Result.Ok(True) + except Exception as rename_err: + scheduled = self._schedule_delete_on_reboot(path) + if scheduled: + return Result.Err( + "DELETE_PENDING_REBOOT", + f"Could not delete {path} now; scheduled for deletion at reboot. " + f"Delete error: {last_error} | Rename failed: {rename_err}", + ) + return Result.Err( + "DELETE_FAILED", + f"Could not delete {path}: {last_error} | Rename failed: {rename_err}", + ) + async def _begin_tx_for_async(self, mode: str) -> Result[str]: fut = self._loop_thread.submit(self._begin_tx_async(mode)) return await asyncio.wrap_future(fut) diff --git a/mjr_am_backend/adapters/fs/list_cache_watcher.py b/mjr_am_backend/adapters/fs/list_cache_watcher.py index cdde6de5..6c1dfb11 100644 --- a/mjr_am_backend/adapters/fs/list_cache_watcher.py +++ b/mjr_am_backend/adapters/fs/list_cache_watcher.py @@ -1,4 +1,4 @@ -""" +""" Best-effort filesystem event watcher for invalidating directory listing caches. Uses watchdog when available; must never raise to callers (anti-crash rule). @@ -133,5 +133,3 @@ def stop_global_fs_list_cache_watcher() -> None: obs.join(timeout=2.0) except Exception: pass - - diff --git a/mjr_am_backend/adapters/tools/exiftool.py b/mjr_am_backend/adapters/tools/exiftool.py index ca4a9596..4d4ca611 100644 --- a/mjr_am_backend/adapters/tools/exiftool.py +++ b/mjr_am_backend/adapters/tools/exiftool.py @@ -186,60 +186,70 @@ class ExifTool: invoking an actual exiftool binary (not an arbitrary command string). """ raw = (bin_name or "").strip() + if not self._is_safe_executable_name(raw): + return None + resolved = self._resolve_executable_path(raw) + if not resolved: + return None + if not self._is_under_trusted_dirs(resolved): + return None + return resolved if self._looks_like_exiftool_name(resolved) else None + + @staticmethod + def _is_safe_executable_name(raw: str) -> bool: if not raw: - return None + return False if "\x00" in raw or "\n" in raw or "\r" in raw: - return None - if any(ch in raw for ch in ("&", "|", ";", ">", "<")): - return None + return False + return not any(ch in raw for ch in ("&", "|", ";", ">", "<")) + @staticmethod + def _resolve_executable_path(raw: str) -> Optional[str]: resolved = shutil.which(raw) - if not resolved: - try: - candidate = Path(raw) - if candidate.is_file(): - resolved = str(candidate.resolve(strict=True)) - except (OSError, RuntimeError, ValueError): - return None - - if not resolved: + if resolved: + return resolved + try: + candidate = Path(raw) + if candidate.is_file(): + return str(candidate.resolve(strict=True)) + except (OSError, RuntimeError, ValueError): return None + return None - # Optional hardening: when MAJOOR_EXIFTOOL_TRUSTED_DIRS is set, require the - # resolved binary to live under one of the configured directories. + @staticmethod + def _is_under_trusted_dirs(resolved: str) -> bool: trusted_dirs_raw = str(os.getenv("MAJOOR_EXIFTOOL_TRUSTED_DIRS", "") or "").strip() - if trusted_dirs_raw: - try: - resolved_path = Path(resolved).resolve(strict=True) - trusted_roots: List[Path] = [] - for item in trusted_dirs_raw.split(os.pathsep): - item = item.strip() - if not item: - continue - try: - trusted_roots.append(Path(item).expanduser().resolve(strict=True)) - except Exception: - continue - if trusted_roots: - ok = any( - resolved_path == root or root in resolved_path.parents - for root in trusted_roots - ) - if not ok: - return None - except Exception: - return None + if not trusted_dirs_raw: + return True + try: + resolved_path = Path(resolved).resolve(strict=True) + trusted_roots = ExifTool._trusted_roots(trusted_dirs_raw) + if not trusted_roots: + return True + return any(resolved_path == root or root in resolved_path.parents for root in trusted_roots) + except Exception: + return False + @staticmethod + def _trusted_roots(trusted_dirs_raw: str) -> List[Path]: + roots: List[Path] = [] + for item in trusted_dirs_raw.split(os.pathsep): + item = item.strip() + if not item: + continue + try: + roots.append(Path(item).expanduser().resolve(strict=True)) + except Exception: + continue + return roots + + @staticmethod + def _looks_like_exiftool_name(resolved: str) -> bool: try: name = Path(resolved).name.lower() except Exception: name = str(resolved).lower() - - # Allow exiftool.exe, exiftool(-k).exe, etc. - if not name.startswith("exiftool"): - return None - - return resolved + return name.startswith("exiftool") def _check_available(self) -> bool: """Check if ExifTool is available in PATH.""" @@ -253,6 +263,370 @@ class ExifTool: """Check if ExifTool is available.""" return self._available + @staticmethod + def _assign_result_for_paths( + paths: List[str], + results: Dict[str, Result[Dict[str, Any]]], + value: Result[Dict[str, Any]], + *, + only_missing: bool = False, + ) -> None: + for path in paths: + if only_missing and path in results: + continue + results[path] = value + + @staticmethod + def _collect_valid_batch_paths( + paths: List[str], + ) -> Tuple[List[str], Dict[str, Result[Dict[str, Any]]]]: + valid_paths: List[str] = [] + results: Dict[str, Result[Dict[str, Any]]] = {} + for path in paths: + if not path or "\x00" in str(path): + results[str(path)] = Result.Err( + ErrorCode.INVALID_INPUT, "Invalid file path", quality="none" + ) + continue + try: + p = Path(str(path)) + if not p.exists() or not p.is_file(): + results[str(path)] = Result.Err( + ErrorCode.NOT_FOUND, f"File not found: {path}", quality="none" + ) + continue + valid_paths.append(str(path)) + except Exception: + results[str(path)] = Result.Err( + ErrorCode.INVALID_INPUT, "Invalid file path", quality="none" + ) + return valid_paths, results + + def _validate_batch_tags( + self, + tags: Optional[List[str]], + valid_paths: List[str], + results: Dict[str, Result[Dict[str, Any]]], + ) -> Optional[List[str]]: + tags_res = _validate_exiftool_tags(tags) + if tags_res.ok: + return tags_res.data or [] + err = Result.Err(tags_res.code, tags_res.error or "Invalid tags", **(tags_res.meta or {})) + self._assign_result_for_paths(valid_paths, results, err) + return None + + def _build_batch_command( + self, + cmd_paths: List[str], + safe_tags: List[str], + ) -> Tuple[List[str], Optional[bytes], float]: + cmd = [self.bin, "-j", "-G1", "-ee", "-a", "-U", "-s"] + if safe_tags: + cmd.extend([f"-{tag}" for tag in safe_tags]) + + stdin_input: Optional[bytes] = None + if os.name == "nt": + cmd.extend(["-charset", "filename=utf8", "-@", "-"]) + stdin_lines = "".join(f"{p}\r\n" for p in cmd_paths) + stdin_input = stdin_lines.encode("utf-8", errors="replace") + else: + cmd.extend(cmd_paths) + + timeout_s = self.timeout * len(cmd_paths) + return cmd, stdin_input, timeout_s + + @staticmethod + def _run_batch_subprocess( + cmd: List[str], + timeout_s: float, + stdin_input: Optional[bytes], + ) -> subprocess.CompletedProcess: + return subprocess.run( + cmd, + capture_output=True, + text=False, + check=False, + timeout=timeout_s, + input=stdin_input, + shell=False, + ) + + def _retry_windows_batch_file_not_found( + self, + process: subprocess.CompletedProcess, + safe_tags: List[str], + cmd_paths: List[str], + timeout_s: float, + ) -> subprocess.CompletedProcess: + if os.name != "nt" or process.returncode == 0: + return process + stderr0, _ = _decode_bytes_best_effort(process.stderr) + if "file not found" not in (stderr0 or "").lower(): + return process + + retry_cmd = [self.bin, "-j", "-G1", "-ee", "-a", "-U", "-s"] + if safe_tags: + retry_cmd.extend([f"-{tag}" for tag in safe_tags]) + retry_cmd.extend(["-charset", "filename=utf8", "-@", "-"]) + stdin_lines = "".join(f"{p}\r\n" for p in cmd_paths) + retry_process = self._run_batch_subprocess( + retry_cmd, + timeout_s, + stdin_lines.encode("utf-8", errors="replace"), + ) + if retry_process.returncode == 0: + return retry_process + return process + + def _parse_batch_output( + self, + process: subprocess.CompletedProcess, + valid_paths: List[str], + results: Dict[str, Result[Dict[str, Any]]], + ) -> Tuple[Optional[List[Dict[str, Any]]], bool]: + stdout, stdout_rep = _decode_bytes_best_effort(process.stdout) + stderr, stderr_rep = _decode_bytes_best_effort(process.stderr) + had_replacements = bool(stdout_rep or stderr_rep) + + if had_replacements: + logger.warning("ExifTool batch output contained decoding replacement characters") + + if process.returncode != 0: + stderr_msg = stderr.strip() + logger.warning(f"ExifTool batch error: {stderr_msg}") + + if not stdout.strip(): + if process.returncode != 0: + err = Result.Err( + ErrorCode.EXIFTOOL_ERROR, + (stderr.strip() or f"ExifTool batch failed with return code {int(process.returncode)}"), + return_code=int(process.returncode), + quality="degraded", + ) + else: + err = Result.Err( + ErrorCode.PARSE_ERROR, + "ExifTool returned empty output", + quality="degraded", + ) + self._assign_result_for_paths(valid_paths, results, err) + return None, had_replacements + + try: + data = json.loads(stdout) + except json.JSONDecodeError as e: + logger.error(f"ExifTool batch JSON parse error: {e}") + err = Result.Err( + ErrorCode.PARSE_ERROR, + f"Failed to parse ExifTool output: {e}", + quality="degraded", + ) + self._assign_result_for_paths(valid_paths, results, err) + return None, had_replacements + + if not isinstance(data, list): + err = Result.Err( + ErrorCode.PARSE_ERROR, "ExifTool batch returned non-list", quality="degraded" + ) + self._assign_result_for_paths(valid_paths, results, err) + return None, had_replacements + return data, had_replacements + + @staticmethod + def _map_batch_results( + data: List[Dict[str, Any]], + key_to_paths: Dict[str, List[str]], + results: Dict[str, Result[Dict[str, Any]]], + *, + had_replacements: bool, + ) -> None: + for file_data in data: + if not isinstance(file_data, dict): + logger.warning(f"ExifTool batch returned non-dict item: {type(file_data)}") + continue + + source_file = file_data.get("SourceFile") + if not source_file: + logger.warning("ExifTool result missing SourceFile field") + continue + + source_key = _normalize_match_path(str(source_file)) + if not source_key: + logger.warning(f"Failed to normalize SourceFile path {source_file}") + continue + + matched_paths = key_to_paths.get(source_key) + if matched_paths: + for path in matched_paths: + results[path] = Result.Ok( + file_data, + quality="full" if not had_replacements else "degraded", + ) + else: + logger.warning(f"ExifTool returned data for unexpected file: {source_file}") + + @staticmethod + def _fill_missing_batch_results( + valid_paths: List[str], + results: Dict[str, Result[Dict[str, Any]]], + ) -> None: + for path in valid_paths: + if path in results: + continue + results[path] = Result.Err( + ErrorCode.PARSE_ERROR, + "No metadata returned by ExifTool", + quality="none", + ) + + def _handle_batch_exception( + self, + exc: Exception, + valid_paths: List[str], + results: Dict[str, Result[Dict[str, Any]]], + safe_tags: List[str], + ) -> Dict[str, Result[Dict[str, Any]]]: + if getattr(exc, "winerror", None) == _WINDOWS_CMDLINE_TOO_LONG: + logger.warning( + "ExifTool batch hit WinError 206; falling back to per-file extraction for %d files", + len(valid_paths), + ) + for path in valid_paths: + if path in results: + continue + results[path] = self.read(path, safe_tags) + return results + + logger.error(f"ExifTool batch unexpected error: {exc}") + err = Result.Err(ErrorCode.EXIFTOOL_ERROR, str(exc), quality="degraded") + self._assign_result_for_paths(valid_paths, results, err, only_missing=True) + return results + + @staticmethod + def _validate_single_read_path(path: str) -> Result[str]: + if not path or "\x00" in str(path): + return Result.Err(ErrorCode.INVALID_INPUT, "Invalid file path", quality="none") + try: + p = Path(str(path)) + except Exception: + return Result.Err(ErrorCode.INVALID_INPUT, "Invalid file path", quality="none") + if not p.exists() or not p.is_file(): + return Result.Err(ErrorCode.NOT_FOUND, f"File not found: {path}", quality="none") + return Result.Ok(str(path)) + + @staticmethod + def _validate_single_read_tags(tags: Optional[List[str]]) -> Result[List[str]]: + tags_res = _validate_exiftool_tags(tags) + if tags_res.ok: + return Result.Ok(tags_res.data or []) + return Result.Err(tags_res.code, tags_res.error or "Invalid tags", **(tags_res.meta or {})) + + @staticmethod + def _run_exiftool_process( + cmd: List[str], + *, + timeout: float, + stdin_input: Optional[str] = None, + ) -> subprocess.CompletedProcess: + stdin_bytes: Optional[bytes] + if stdin_input is None: + stdin_bytes = None + else: + stdin_bytes = str(stdin_input).encode("utf-8", errors="replace") + return subprocess.run( + cmd, + capture_output=True, + text=False, + check=False, + timeout=timeout, + input=stdin_bytes, + shell=False, + ) + + def _build_single_read_command( + self, + path: str, + safe_tags: List[str], + ) -> Tuple[List[str], Optional[str]]: + cmd = [self.bin, "-j", "-G1", "-ee", "-a", "-U", "-s"] + if safe_tags: + cmd.extend([f"-{tag}" for tag in safe_tags]) + if os.name == "nt": + cmd.extend(["-charset", "filename=utf8", str(path)]) + return cmd, None + cmd.append(path) + return cmd, None + + def _retry_windows_single_read_file_not_found( + self, + process: subprocess.CompletedProcess, + path: str, + safe_tags: List[str], + ) -> subprocess.CompletedProcess: + if os.name != "nt" or process.returncode == 0: + return process + stderr_msg, _ = _decode_bytes_best_effort(process.stderr) + stderr_msg = stderr_msg.strip() + if "file not found" not in stderr_msg.lower(): + return process + + retry_cmd = [self.bin, "-j", "-G1", "-ee", "-a", "-U", "-s"] + if safe_tags: + retry_cmd.extend([f"-{tag}" for tag in safe_tags]) + retry_cmd.extend(["-charset", "filename=utf8", "-@", "-"]) + retry_process = self._run_exiftool_process( + retry_cmd, + timeout=self.timeout, + stdin_input=f"{path}\r\n", + ) + if retry_process.returncode == 0: + return retry_process + return process + + @staticmethod + def _parse_single_read_process( + process: subprocess.CompletedProcess, + path: str, + ) -> Result[Dict[str, Any]]: + stdout, stdout_rep = _decode_bytes_best_effort(process.stdout) + stderr, stderr_rep = _decode_bytes_best_effort(process.stderr) + had_replacements = bool(stdout_rep or stderr_rep) + if had_replacements: + logger.warning("ExifTool output contained decoding replacement characters for %s", path) + + if process.returncode != 0: + stderr_msg = stderr.strip() + logger.warning(f"ExifTool error for {path}: {stderr_msg}") + return Result.Err( + ErrorCode.EXIFTOOL_ERROR, + stderr_msg or "ExifTool command failed", + return_code=int(process.returncode), + quality="degraded", + ) + + if not stdout.strip(): + return Result.Err( + ErrorCode.PARSE_ERROR, + "ExifTool returned empty output", + quality="degraded", + ) + try: + data = json.loads(stdout) + except json.JSONDecodeError as exc: + logger.error(f"ExifTool JSON parse error: {exc}") + return Result.Err( + ErrorCode.PARSE_ERROR, + f"Failed to parse ExifTool output: {exc}", + quality="degraded", + ) + if isinstance(data, list) and len(data) > 0: + return Result.Ok(data[0], quality="full" if not had_replacements else "degraded") + return Result.Err( + ErrorCode.PARSE_ERROR, + "No metadata found", + quality="none", + ) + def read(self, path: str, tags: Optional[List[str]] = None) -> Result[Dict[str, Any]]: """ Read metadata from file using ExifTool. @@ -264,132 +638,14 @@ class ExifTool: Returns: Result with metadata dict or error """ - if not self._available: - return Result.Err( - ErrorCode.TOOL_MISSING, - "ExifTool not found in PATH", - quality="none" - ) - - if not path or "\x00" in str(path): - return Result.Err(ErrorCode.INVALID_INPUT, "Invalid file path", quality="none") - try: - p = Path(str(path)) - except Exception: - return Result.Err(ErrorCode.INVALID_INPUT, "Invalid file path", quality="none") - if not p.exists() or not p.is_file(): - return Result.Err(ErrorCode.NOT_FOUND, f"File not found: {path}", quality="none") - - tags_res = _validate_exiftool_tags(tags) - if not tags_res.ok: - return Result.Err(tags_res.code, tags_res.error or "Invalid tags", **(tags_res.meta or {})) - safe_tags = tags_res.data or [] - - def _run(cmd: List[str], stdin_input: Optional[str]) -> subprocess.CompletedProcess: - stdin_bytes: Optional[bytes] - if stdin_input is None: - stdin_bytes = None - else: - stdin_bytes = str(stdin_input).encode("utf-8", errors="replace") - return subprocess.run( - cmd, - capture_output=True, - text=False, - check=False, - timeout=self.timeout, - input=stdin_bytes, - shell=False, - ) + safe_tags_or_error = self._prepare_single_read_inputs(path, tags) + if isinstance(safe_tags_or_error, Result): + return safe_tags_or_error + safe_tags = safe_tags_or_error try: - # Build command with comprehensive extraction flags - # -j: JSON output - # -G1: Group names with family 1 - # -ee: Extract embedded data (critical for ComfyUI workflows) - # -a: Allow duplicate tags - # -U: Extract unknown tags - # -s: Short tag names - cmd = [self.bin, "-j", "-G1", "-ee", "-a", "-U", "-s"] - if safe_tags: - cmd.extend([f"-{tag}" for tag in safe_tags]) - - # Windows Unicode compatibility: - # - Prefer argv (most robust with Win32 wide-char argv). - # - If ExifTool reports "File not found" (common when filename charset handling - # is buggy), retry via stdin argfile with explicit filename charset. - stdin_input: Optional[str] = None - if os.name == "nt": - cmd.extend(["-charset", "filename=utf8", str(path)]) - else: - cmd.append(path) - - # Run command - process = _run(cmd, stdin_input) - - if os.name == "nt" and process.returncode != 0: - stderr_msg, _ = _decode_bytes_best_effort(process.stderr) - stderr_msg = stderr_msg.strip() - # Retry only for the specific "file not found" failure mode; avoid - # masking other errors (invalid tags, parse errors, etc.). - if "file not found" in stderr_msg.lower(): - retry_cmd = [self.bin, "-j", "-G1", "-ee", "-a", "-U", "-s"] - if safe_tags: - retry_cmd.extend([f"-{tag}" for tag in safe_tags]) - retry_cmd.extend(["-charset", "filename=utf8", "-@", "-"]) - # Prefer CRLF on Windows to match typical ExifTool stdin parsing behavior. - retry_process = _run(retry_cmd, f"{path}\r\n") - if retry_process.returncode == 0: - process = retry_process - - stdout, stdout_rep = _decode_bytes_best_effort(process.stdout) - stderr, stderr_rep = _decode_bytes_best_effort(process.stderr) - had_replacements = bool(stdout_rep or stderr_rep) - if had_replacements: - logger.warning("ExifTool output contained decoding replacement characters for %s", path) - - if process.returncode != 0: - stderr_msg = stderr.strip() - logger.warning(f"ExifTool error for {path}: {stderr_msg}") - return Result.Err( - ErrorCode.EXIFTOOL_ERROR, - stderr_msg or "ExifTool command failed", - return_code=int(process.returncode), - quality="degraded" - ) - - # Parse JSON output - if not stdout.strip(): - if process.returncode != 0: - stderr_msg = stderr.strip() - return Result.Err( - ErrorCode.EXIFTOOL_ERROR, - stderr_msg or f"ExifTool failed with return code {int(process.returncode)}", - return_code=int(process.returncode), - quality="degraded", - ) - return Result.Err( - ErrorCode.PARSE_ERROR, - "ExifTool returned empty output", - quality="degraded" - ) - try: - data = json.loads(stdout) - except json.JSONDecodeError as e: - logger.error(f"ExifTool JSON parse error: {e}") - return Result.Err( - ErrorCode.PARSE_ERROR, - f"Failed to parse ExifTool output: {e}", - quality="degraded" - ) - if isinstance(data, list) and len(data) > 0: - return Result.Ok(data[0], quality="full" if not had_replacements else "degraded") - - return Result.Err( - ErrorCode.PARSE_ERROR, - "No metadata found", - quality="none" - ) - + process = self._run_single_read_process(path, safe_tags) + return self._parse_single_read_process(process, path) except subprocess.TimeoutExpired: logger.error(f"ExifTool timeout for {path}") return Result.Err( @@ -412,6 +668,44 @@ class ExifTool: quality="degraded" ) + def _prepare_single_read_inputs( + self, + path: str, + tags: Optional[List[str]], + ) -> Result[Dict[str, Any]] | List[str]: + availability_error = self._single_read_availability_error() + if availability_error is not None: + return availability_error + path_res = self._validate_single_read_path(path) + if not path_res.ok: + return Result.Err(path_res.code, path_res.error or "Invalid file path", **(path_res.meta or {})) + tags_res = self._validate_single_read_tags(tags) + if not tags_res.ok: + return Result.Err(tags_res.code, tags_res.error or "Invalid tags", **(tags_res.meta or {})) + return tags_res.data or [] + + def _single_read_availability_error(self) -> Optional[Result[Dict[str, Any]]]: + if self._available: + return None + return Result.Err( + ErrorCode.TOOL_MISSING, + "ExifTool not found in PATH", + quality="none" + ) + + def _run_single_read_process( + self, + path: str, + safe_tags: List[str], + ) -> subprocess.CompletedProcess: + cmd, stdin_input = self._build_single_read_command(path, safe_tags) + process = self._run_exiftool_process( + cmd, + timeout=self.timeout, + stdin_input=stdin_input, + ) + return self._retry_windows_single_read_file_not_found(process, path, safe_tags) + def read_batch(self, paths: List[str], tags: Optional[List[str]] = None) -> Dict[str, Result[Dict[str, Any]]]: """ Read metadata from multiple files using a single ExifTool invocation. @@ -434,187 +728,34 @@ class ExifTool: if not paths: return {} - # Filter valid paths - valid_paths: List[str] = [] - results: Dict[str, Result[Dict[str, Any]]] = {} - - for path in paths: - if not path or "\x00" in str(path): - results[str(path)] = Result.Err(ErrorCode.INVALID_INPUT, "Invalid file path", quality="none") - continue - try: - p = Path(str(path)) - if not p.exists() or not p.is_file(): - results[str(path)] = Result.Err(ErrorCode.NOT_FOUND, f"File not found: {path}", quality="none") - continue - valid_paths.append(str(path)) - except Exception: - results[str(path)] = Result.Err(ErrorCode.INVALID_INPUT, "Invalid file path", quality="none") - + valid_paths, results = self._collect_valid_batch_paths(paths) if not valid_paths: return results - tags_res = _validate_exiftool_tags(tags) - if not tags_res.ok: - err = Result.Err(tags_res.code, tags_res.error or "Invalid tags", **(tags_res.meta or {})) - for path in valid_paths: - results[path] = err + safe_tags = self._validate_batch_tags(tags, valid_paths, results) + if safe_tags is None: return results - safe_tags = tags_res.data or [] key_to_paths, cmd_paths = _build_match_map(valid_paths) try: - # Build batch command - cmd = [self.bin, "-j", "-G1", "-ee", "-a", "-U", "-s"] - if safe_tags: - cmd.extend([f"-{tag}" for tag in safe_tags]) - - stdin_input: Optional[bytes] = None - if os.name == "nt": - # Avoid WinError 206 by passing file list via stdin argfile. - cmd.extend(["-charset", "filename=utf8", "-@", "-"]) - stdin_lines = "".join(f"{p}\r\n" for p in cmd_paths) - stdin_input = stdin_lines.encode("utf-8", errors="replace") - else: - # Add all paths (de-duplicated by normalized key to avoid repeated work) - cmd.extend(cmd_paths) - - timeout_s = self.timeout * len(cmd_paths) - - # Run command - process = subprocess.run( - cmd, - capture_output=True, - text=False, - check=False, - timeout=timeout_s, # Scale timeout with file count - input=stdin_input, - shell=False, - ) - - # Windows retry path: use stdin argfile in UTF-8 when argv path decoding fails. - if os.name == "nt" and process.returncode != 0: - stderr0, _ = _decode_bytes_best_effort(process.stderr) - if "file not found" in (stderr0 or "").lower(): - retry_cmd = [self.bin, "-j", "-G1", "-ee", "-a", "-U", "-s"] - if safe_tags: - retry_cmd.extend([f"-{tag}" for tag in safe_tags]) - retry_cmd.extend(["-charset", "filename=utf8", "-@", "-"]) - stdin_lines = "".join(f"{p}\r\n" for p in cmd_paths) - retry_process = subprocess.run( - retry_cmd, - capture_output=True, - text=False, - check=False, - timeout=timeout_s, - input=stdin_lines.encode("utf-8", errors="replace"), - shell=False, - ) - if retry_process.returncode == 0: - process = retry_process - - stdout, stdout_rep = _decode_bytes_best_effort(process.stdout) - stderr, stderr_rep = _decode_bytes_best_effort(process.stderr) - had_replacements = bool(stdout_rep or stderr_rep) - - if had_replacements: - logger.warning("ExifTool batch output contained decoding replacement characters") - - if process.returncode != 0: - stderr_msg = stderr.strip() - logger.warning(f"ExifTool batch error: {stderr_msg}") - # Partial failure - some files may have succeeded - - # Parse JSON output - if not stdout.strip(): - if process.returncode != 0: - err = Result.Err( - ErrorCode.EXIFTOOL_ERROR, - (stderr.strip() or f"ExifTool batch failed with return code {int(process.returncode)}"), - return_code=int(process.returncode), - quality="degraded", - ) - else: - err = Result.Err(ErrorCode.PARSE_ERROR, "ExifTool returned empty output", quality="degraded") - for path in valid_paths: - results[path] = err + cmd, stdin_input, timeout_s = self._build_batch_command(cmd_paths, safe_tags) + process = self._run_batch_subprocess(cmd, timeout_s, stdin_input) + process = self._retry_windows_batch_file_not_found(process, safe_tags, cmd_paths, timeout_s) + data, had_replacements = self._parse_batch_output(process, valid_paths, results) + if data is None: return results - - try: - data = json.loads(stdout) - except json.JSONDecodeError as e: - logger.error(f"ExifTool batch JSON parse error: {e}") - err = Result.Err(ErrorCode.PARSE_ERROR, f"Failed to parse ExifTool output: {e}", quality="degraded") - for path in valid_paths: - results[path] = err - return results - - if not isinstance(data, list): - err = Result.Err(ErrorCode.PARSE_ERROR, "ExifTool batch returned non-list", quality="degraded") - for path in valid_paths: - results[path] = err - return results - - # Match results to input paths using SourceFile field. - # ExifTool may skip corrupted files, so we CANNOT rely on array index matching. - for file_data in data: - if not isinstance(file_data, dict): - logger.warning(f"ExifTool batch returned non-dict item: {type(file_data)}") - continue - - # Get the source file from ExifTool result - source_file = file_data.get("SourceFile") - if not source_file: - logger.warning("ExifTool result missing SourceFile field") - continue - - source_key = _normalize_match_path(str(source_file)) - if not source_key: - logger.warning(f"Failed to normalize SourceFile path {source_file}") - continue - - matched_paths = key_to_paths.get(source_key) - if matched_paths: - for path in matched_paths: - results[path] = Result.Ok( - file_data, - quality="full" if not had_replacements else "degraded", - ) - else: - logger.warning(f"ExifTool returned data for unexpected file: {source_file}") - - # Handle files that didn't return results (corrupted, unsupported, etc.) - for path in valid_paths: - if path not in results: - results[path] = Result.Err(ErrorCode.PARSE_ERROR, "No metadata returned by ExifTool", quality="none") - + self._map_batch_results(data, key_to_paths, results, had_replacements=had_replacements) + self._fill_missing_batch_results(valid_paths, results) return results except subprocess.TimeoutExpired: logger.error(f"ExifTool batch timeout for {len(valid_paths)} files") err = Result.Err(ErrorCode.TIMEOUT, "ExifTool batch timeout", quality="degraded") - for path in valid_paths: - if path not in results: - results[path] = err + self._assign_result_for_paths(valid_paths, results, err, only_missing=True) return results except Exception as e: - if getattr(e, "winerror", None) == _WINDOWS_CMDLINE_TOO_LONG: - logger.warning( - "ExifTool batch hit WinError 206; falling back to per-file extraction for %d files", - len(valid_paths), - ) - for path in valid_paths: - if path in results: - continue - results[path] = self.read(path, safe_tags) - return results - logger.error(f"ExifTool batch unexpected error: {e}") - err = Result.Err(ErrorCode.EXIFTOOL_ERROR, str(e), quality="degraded") - for path in valid_paths: - if path not in results: - results[path] = err - return results + return self._handle_batch_exception(e, valid_paths, results, safe_tags) def write(self, path: str, metadata: dict, preserve_workflow: bool = True) -> Result[bool]: """ @@ -628,98 +769,22 @@ class ExifTool: Returns: Result with success boolean """ - if not self._available: + availability_error = self._validate_write_preconditions(path) + if availability_error is not None: + return availability_error + + invalid_keys = self._invalid_write_keys(metadata) + if invalid_keys: return Result.Err( - ErrorCode.TOOL_MISSING, - "ExifTool not found in PATH" + ErrorCode.INVALID_INPUT, + "Invalid ExifTool tag format", + invalid_tags=invalid_keys, ) - if not path or "\x00" in str(path): - return Result.Err(ErrorCode.INVALID_INPUT, "Invalid file path") + cmd, stdin_input = self._build_write_command(path, metadata, preserve_workflow) try: - p = Path(str(path)) - except Exception: - return Result.Err(ErrorCode.INVALID_INPUT, "Invalid file path") - if not p.exists() or not p.is_file(): - return Result.Err(ErrorCode.NOT_FOUND, f"File not found: {path}") - - try: - # Build command - cmd = [self.bin] - - # Preserve existing metadata if requested - if preserve_workflow: - cmd.extend(["-tagsFromFile", "@"]) - - invalid_keys: List[str] = [] - for key in (metadata or {}).keys(): - if not isinstance(key, str): - invalid_keys.append(str(key)) - continue - if not _is_safe_exiftool_tag(key.strip()): - invalid_keys.append(key.strip()) - if invalid_keys: - return Result.Err( - ErrorCode.INVALID_INPUT, - "Invalid ExifTool tag format", - invalid_tags=invalid_keys, - ) - - # Add metadata tags - # - Scalars: `-Tag=Value` - # - Lists: clear then append (`-Tag=` then `-Tag+=Item`) - # This matches ExifTool's array semantics (XMP arrays, IPTC keyword lists, etc.). - for key, value in (metadata or {}).items(): - if value is None: - cmd.append(f"-{key}=") - continue - if isinstance(value, list): - cmd.append(f"-{key}=") - for item in value: - if item is None: - continue - text = str(item).strip() - if text: - cmd.append(f"-{key}+={text}") - continue - cmd.append(f"-{key}={value}") - - stdin_input = None - if os.name == "nt": - cmd.extend(["-charset", "filename=utf8", "-@", "-"]) - stdin_input = f"{path}\n" - cmd.append("-overwrite_original") - else: - # Overwrite original file - cmd.extend(["-overwrite_original", path]) - - # Run command - process = subprocess.run( - cmd, - capture_output=True, - text=False, - check=False, - timeout=self.timeout, - input=(stdin_input.encode("utf-8", errors="replace") if stdin_input is not None else None), - shell=False, - ) - - stdout, stdout_rep = _decode_bytes_best_effort(process.stdout) - stderr, stderr_rep = _decode_bytes_best_effort(process.stderr) - if stdout_rep or stderr_rep: - logger.warning("ExifTool write output contained decoding replacement characters for %s", path) - - if process.returncode != 0: - stderr_msg = stderr.strip() - logger.warning(f"ExifTool write error for {path}: {stderr_msg}") - return Result.Err( - ErrorCode.EXIFTOOL_ERROR, - stderr_msg or "ExifTool write failed", - return_code=int(process.returncode), - ) - - logger.debug(f"Metadata written to {path}") - return Result.Ok(True) + process = self._run_write_command(cmd, stdin_input) + return self._handle_write_process_result(process, path) except subprocess.TimeoutExpired: logger.error(f"ExifTool write timeout for {path}") @@ -734,6 +799,92 @@ class ExifTool: str(e) ) + def _validate_write_preconditions(self, path: str) -> Optional[Result[bool]]: + if not self._available: + return Result.Err(ErrorCode.TOOL_MISSING, "ExifTool not found in PATH") + if not path or "\x00" in str(path): + return Result.Err(ErrorCode.INVALID_INPUT, "Invalid file path") + try: + file_path = Path(str(path)) + except Exception: + return Result.Err(ErrorCode.INVALID_INPUT, "Invalid file path") + if not file_path.exists() or not file_path.is_file(): + return Result.Err(ErrorCode.NOT_FOUND, f"File not found: {path}") + return None + + @staticmethod + def _invalid_write_keys(metadata: dict) -> List[str]: + invalid_keys: List[str] = [] + for key in (metadata or {}).keys(): + if not isinstance(key, str): + invalid_keys.append(str(key)) + continue + if not _is_safe_exiftool_tag(key.strip()): + invalid_keys.append(key.strip()) + return invalid_keys + + def _build_write_command(self, path: str, metadata: dict, preserve_workflow: bool) -> tuple[List[str], Optional[str]]: + cmd = [self.bin] + if preserve_workflow: + cmd.extend(["-tagsFromFile", "@"]) + self._append_metadata_write_args(cmd, metadata) + return self._append_write_target_args(cmd, path) + + @staticmethod + def _append_metadata_write_args(cmd: List[str], metadata: dict) -> None: + for key, value in (metadata or {}).items(): + if value is None: + cmd.append(f"-{key}=") + continue + if isinstance(value, list): + cmd.append(f"-{key}=") + for item in value: + if item is None: + continue + text = str(item).strip() + if text: + cmd.append(f"-{key}+={text}") + continue + cmd.append(f"-{key}={value}") + + @staticmethod + def _append_write_target_args(cmd: List[str], path: str) -> tuple[List[str], Optional[str]]: + stdin_input: Optional[str] = None + if os.name == "nt": + cmd.extend(["-charset", "filename=utf8", "-@", "-"]) + stdin_input = f"{path}\n" + cmd.append("-overwrite_original") + else: + cmd.extend(["-overwrite_original", path]) + return cmd, stdin_input + + def _run_write_command(self, cmd: List[str], stdin_input: Optional[str]): + return subprocess.run( + cmd, + capture_output=True, + text=False, + check=False, + timeout=self.timeout, + input=(stdin_input.encode("utf-8", errors="replace") if stdin_input is not None else None), + shell=False, + ) + + def _handle_write_process_result(self, process: subprocess.CompletedProcess, path: str) -> Result[bool]: + stdout, stdout_rep = _decode_bytes_best_effort(process.stdout) + stderr, stderr_rep = _decode_bytes_best_effort(process.stderr) + if stdout_rep or stderr_rep: + logger.warning("ExifTool write output contained decoding replacement characters for %s", path) + if process.returncode != 0: + stderr_msg = stderr.strip() + logger.warning(f"ExifTool write error for {path}: {stderr_msg}") + return Result.Err( + ErrorCode.EXIFTOOL_ERROR, + stderr_msg or "ExifTool write failed", + return_code=int(process.returncode), + ) + logger.debug(f"Metadata written to {path}") + return Result.Ok(True) + async def aread(self, path: str, tags: Optional[List[str]] = None) -> Result[Dict[str, Any]]: """Async wrapper for read() executed off the event loop thread.""" return await asyncio.to_thread(self.read, path, tags) diff --git a/mjr_am_backend/adapters/tools/ffprobe.py b/mjr_am_backend/adapters/tools/ffprobe.py index 16448d25..09de25b5 100644 --- a/mjr_am_backend/adapters/tools/ffprobe.py +++ b/mjr_am_backend/adapters/tools/ffprobe.py @@ -46,34 +46,43 @@ class FFProbe: actual ffprobe binary (not an arbitrary command string). """ raw = (bin_name or "").strip() + if not self._is_safe_executable_token(raw): + return None + resolved = self._resolve_executable_path(raw) + if not resolved: + return None + return resolved if self._is_ffprobe_name(resolved) else None + + @staticmethod + def _is_safe_executable_token(raw: str) -> bool: if not raw: - return None + return False if "\x00" in raw or "\n" in raw or "\r" in raw: - return None + return False if any(ch in raw for ch in ("&", "|", ";", ">", "<")): - return None + return False + return True + @staticmethod + def _resolve_executable_path(raw: str) -> Optional[str]: resolved = shutil.which(raw) - if not resolved: - try: - candidate = Path(raw) - if candidate.is_file(): - resolved = str(candidate.resolve(strict=True)) - except (OSError, RuntimeError, ValueError): - return None - - if not resolved: + if resolved: + return resolved + try: + candidate = Path(raw) + if candidate.is_file(): + return str(candidate.resolve(strict=True)) + except (OSError, RuntimeError, ValueError): return None + return None + @staticmethod + def _is_ffprobe_name(resolved: str) -> bool: try: name = Path(resolved).name.lower() except Exception: name = str(resolved).lower() - - if not name.startswith("ffprobe"): - return None - - return resolved + return name.startswith("ffprobe") def _check_available(self) -> bool: """Check if ffprobe is available in PATH.""" @@ -105,71 +114,10 @@ class FFProbe: ) try: - # Build command - cmd = [ - self._resolved_bin or self.bin, - "-v", "error", # Hide verbose output - "-print_format", "json", # JSON output - "-show_format", # Show container format - "-show_streams", # Show stream info - path - ] - - # Run command - process = subprocess.run( - cmd, - capture_output=True, - text=True, - check=False, - timeout=self.timeout, - shell=False, - close_fds=os.name != "nt", - ) - - if process.returncode != 0: - stderr = process.stderr.strip() - logger.warning(f"ffprobe error for {path}: {stderr}") - return Result.Err( - ErrorCode.FFPROBE_ERROR, - stderr or "ffprobe command failed", - quality="degraded" - ) - - # Parse JSON output - if not process.stdout or not process.stdout.strip(): - logger.warning(f"ffprobe returned empty output for {path}") - return Result.Err( - ErrorCode.FFPROBE_ERROR, - "No ffprobe output", - quality="degraded" - ) - - data = json.loads(process.stdout) - - if not isinstance(data, dict): - return Result.Err( - ErrorCode.PARSE_ERROR, - "Invalid ffprobe output format", - quality="degraded" - ) - - # Extract useful metadata - result = { - "format": data.get("format", {}), - "streams": data.get("streams", []), - "video_stream": self._find_video_stream(data.get("streams", [])), - "audio_stream": self._find_audio_stream(data.get("streams", [])), - } - - return Result.Ok(result, quality="full") - + process = self._run_ffprobe_cmd(self._build_ffprobe_cmd(path)) + return self._parse_ffprobe_output(process.stdout, process.stderr, process.returncode, path) except subprocess.TimeoutExpired: - logger.error(f"ffprobe timeout for {path}") - return Result.Err( - ErrorCode.TIMEOUT, - f"ffprobe timeout after {self.timeout}s", - quality="degraded" - ) + return self._ffprobe_timeout_error(path) except json.JSONDecodeError as e: logger.error(f"ffprobe JSON parse error: {e}") return Result.Err( @@ -185,6 +133,25 @@ class FFProbe: quality="degraded" ) + def _run_ffprobe_cmd(self, cmd: List[str]) -> subprocess.CompletedProcess[str]: + return subprocess.run( + cmd, + capture_output=True, + text=True, + check=False, + timeout=self.timeout, + shell=False, + close_fds=os.name != "nt", + ) + + def _ffprobe_timeout_error(self, path: str) -> Result[dict]: + logger.error(f"ffprobe timeout for {path}") + return Result.Err( + ErrorCode.TIMEOUT, + f"ffprobe timeout after {self.timeout}s", + quality="degraded", + ) + async def aread(self, path: str) -> Result[dict]: """ Async variant of read() using native asyncio subprocess execution. @@ -197,71 +164,13 @@ class FFProbe: ) try: - cmd = [ - self._resolved_bin or self.bin, - "-v", "error", - "-print_format", "json", - "-show_format", - "-show_streams", - path, - ] - - process = await asyncio.create_subprocess_exec( - *cmd, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - close_fds=os.name != "nt", - ) - - try: - stdout_b, stderr_b = await asyncio.wait_for(process.communicate(), timeout=self.timeout) - except asyncio.TimeoutError: - try: - process.kill() - except Exception: - pass - logger.error(f"ffprobe timeout for {path}") - return Result.Err( - ErrorCode.TIMEOUT, - f"ffprobe timeout after {self.timeout}s", - quality="degraded" - ) - - stdout = (stdout_b or b"").decode("utf-8", errors="replace") - stderr = (stderr_b or b"").decode("utf-8", errors="replace") - - if process.returncode != 0: - stderr_msg = stderr.strip() - logger.warning(f"ffprobe error for {path}: {stderr_msg}") - return Result.Err( - ErrorCode.FFPROBE_ERROR, - stderr_msg or "ffprobe command failed", - quality="degraded" - ) - - if not stdout.strip(): - logger.warning(f"ffprobe returned empty output for {path}") - return Result.Err( - ErrorCode.FFPROBE_ERROR, - "No ffprobe output", - quality="degraded" - ) - - data = json.loads(stdout) - if not isinstance(data, dict): - return Result.Err( - ErrorCode.PARSE_ERROR, - "Invalid ffprobe output format", - quality="degraded" - ) - - result = { - "format": data.get("format", {}), - "streams": data.get("streams", []), - "video_stream": self._find_video_stream(data.get("streams", [])), - "audio_stream": self._find_audio_stream(data.get("streams", [])), - } - return Result.Ok(result, quality="full") + cmd = self._build_ffprobe_cmd(path) + process = await self._spawn_ffprobe_process(cmd) + communicated = await self._communicate_with_timeout(process, path) + if not communicated.ok: + return communicated + stdout, stderr = communicated.data + return self._parse_ffprobe_output(stdout, stderr, process.returncode, path) except json.JSONDecodeError as e: logger.error(f"ffprobe JSON parse error: {e}") return Result.Err( @@ -277,6 +186,83 @@ class FFProbe: quality="degraded" ) + def _build_ffprobe_cmd(self, path: str) -> List[str]: + return [ + self._resolved_bin or self.bin, + "-v", "error", + "-print_format", "json", + "-show_format", + "-show_streams", + path, + ] + + async def _spawn_ffprobe_process(self, cmd: List[str]) -> asyncio.subprocess.Process: + return await asyncio.create_subprocess_exec( + *cmd, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + close_fds=os.name != "nt", + ) + + async def _communicate_with_timeout( + self, + process: asyncio.subprocess.Process, + path: str, + ) -> Result[tuple[str, str]]: + try: + stdout_b, stderr_b = await asyncio.wait_for(process.communicate(), timeout=self.timeout) + except asyncio.TimeoutError: + try: + process.kill() + except Exception: + pass + logger.error(f"ffprobe timeout for {path}") + return Result.Err( + ErrorCode.TIMEOUT, + f"ffprobe timeout after {self.timeout}s", + quality="degraded" + ) + stdout = (stdout_b or b"").decode("utf-8", errors="replace") + stderr = (stderr_b or b"").decode("utf-8", errors="replace") + return Result.Ok((stdout, stderr), quality="full") + + def _parse_ffprobe_output( + self, + stdout: str, + stderr: str, + returncode: Optional[int], + path: str, + ) -> Result[dict]: + if returncode != 0: + stderr_msg = stderr.strip() + logger.warning(f"ffprobe error for {path}: {stderr_msg}") + return Result.Err( + ErrorCode.FFPROBE_ERROR, + stderr_msg or "ffprobe command failed", + quality="degraded" + ) + if not stdout.strip(): + logger.warning(f"ffprobe returned empty output for {path}") + return Result.Err( + ErrorCode.FFPROBE_ERROR, + "No ffprobe output", + quality="degraded" + ) + data = json.loads(stdout) + if not isinstance(data, dict): + return Result.Err( + ErrorCode.PARSE_ERROR, + "Invalid ffprobe output format", + quality="degraded" + ) + result = { + "format": data.get("format", {}), + "streams": data.get("streams", []), + "video_stream": self._find_video_stream(data.get("streams", [])), + "audio_stream": self._find_audio_stream(data.get("streams", [])), + } + return Result.Ok(result, quality="full") + def read_batch(self, paths: List[str]) -> Dict[str, Result[dict]]: """ Read metadata from multiple video files. diff --git a/mjr_am_backend/config.py b/mjr_am_backend/config.py index d8f9f23b..c236f954 100644 --- a/mjr_am_backend/config.py +++ b/mjr_am_backend/config.py @@ -72,80 +72,108 @@ def _env_bool(default: bool, *names: str) -> bool: return default def _resolve_output_root() -> Path: - env_path = _env_raw("MJR_AM_OUTPUT_DIRECTORY", "MAJOOR_OUTPUT_DIRECTORY") - if env_path: - try: - return Path(env_path).expanduser().resolve() - except (OSError, RuntimeError): - logger.warning(f"Failed to resolve MAJOOR_OUTPUT_DIRECTORY: {env_path}, using fallback") + resolved = _resolve_output_root_from_env() + if resolved is not None: + return resolved + resolved = _resolve_output_root_from_comfy_cli() + if resolved is not None: + return resolved + resolved = _resolve_output_root_from_folder_paths() + if resolved is not None: + return resolved + return _resolve_output_root_fallback() - # ComfyUI CLI flag support: --output-directory + +def _resolve_output_root_from_env() -> Path | None: + env_path = _env_raw("MJR_AM_OUTPUT_DIRECTORY", "MAJOOR_OUTPUT_DIRECTORY") + if not env_path: + return None + try: + return Path(env_path).expanduser().resolve() + except (OSError, RuntimeError): + logger.warning("Failed to resolve MAJOOR_OUTPUT_DIRECTORY: %s, using fallback", env_path) + return None + + +def _resolve_output_root_from_comfy_cli() -> Path | None: try: from comfy.cli_args import args as comfy_args # type: ignore + cli_out = getattr(comfy_args, "output_directory", None) if cli_out: return Path(str(cli_out)).expanduser().resolve() except Exception: - pass - - # Try to use already-loaded folder_paths first (ComfyUI runtime), then import. + return None + return None + + +def _resolve_output_root_from_folder_paths() -> Path | None: + output_dir = _folder_paths_output_dir_from_loaded_module() + if output_dir is not None: + return output_dir + try: + import folder_paths + + return Path(folder_paths.get_output_directory()).resolve() + except (ImportError, ModuleNotFoundError): + logger.warning("folder_paths module not available, using fallback path detection") + except (AttributeError, TypeError) as exc: + logger.warning("folder_paths.get_output_directory() failed: %s, using fallback", exc) + except (OSError, RuntimeError) as exc: + logger.warning("Failed to resolve output directory from folder_paths: %s, using fallback", exc) + return None + + +def _folder_paths_output_dir_from_loaded_module() -> Path | None: try: fp_mod = sys.modules.get("folder_paths") if fp_mod and hasattr(fp_mod, "get_output_directory"): - output_dir = fp_mod.get_output_directory() - return Path(output_dir).resolve() + return Path(fp_mod.get_output_directory()).resolve() except Exception: - pass - try: - import folder_paths - output_dir = folder_paths.get_output_directory() - return Path(output_dir).resolve() - except (ImportError, ModuleNotFoundError): - logger.warning("folder_paths module not available, using fallback path detection") - except (AttributeError, TypeError) as e: - logger.warning(f"folder_paths.get_output_directory() failed: {e}, using fallback") - except (OSError, RuntimeError) as e: - logger.warning(f"Failed to resolve output directory from folder_paths: {e}, using fallback") - - # Multiple fallback strategies - try: - current_file = Path(__file__).resolve() + return None + return None - # Strategy 1: detect ComfyUI root by signature files and prefer its output dir. - comfy_root = None - for parent in current_file.parents: - if (parent / "main.py").is_file() and (parent / "folder_paths.py").is_file(): - comfy_root = parent - break - if comfy_root is None and len(current_file.parents) > 3: - # Typical layout: ComfyUI/custom_nodes//mjr_am_backend/config.py - comfy_root = current_file.parents[3] - # Strategy 1b: parse raw argv for --output-directory as extra fallback. - try: - argv = list(sys.argv or []) - for i, tok in enumerate(argv): - if tok == "--output-directory" and i + 1 < len(argv): - return Path(str(argv[i + 1])).expanduser().resolve() - if isinstance(tok, str) and tok.startswith("--output-directory="): - return Path(tok.split("=", 1)[1]).expanduser().resolve() - except Exception: - pass + +def _resolve_output_root_fallback() -> Path: + try: + argv_output = _resolve_output_root_from_argv() + if argv_output is not None: + return argv_output + comfy_root = _detect_comfy_root(Path(__file__).resolve()) if comfy_root is not None: output_dir = comfy_root / "output" if output_dir.is_dir(): return output_dir.resolve() - - # Strategy 2: fallback to CWD/output. logger.warning("Output directory not found from ComfyUI hints, creating fallback in current working directory") fallback_dir = Path.cwd() / "output" fallback_dir.mkdir(exist_ok=True) return fallback_dir.resolve() - - except (OSError, RuntimeError) as e: - logger.error(f"All fallback strategies failed: {e}") - # Last resort: return a path that will be handled gracefully elsewhere + except (OSError, RuntimeError) as exc: + logger.error("All fallback strategies failed: %s", exc) return Path.cwd() / "output" + +def _resolve_output_root_from_argv() -> Path | None: + try: + argv = list(sys.argv or []) + for i, token in enumerate(argv): + if token == "--output-directory" and i + 1 < len(argv): + return Path(str(argv[i + 1])).expanduser().resolve() + if isinstance(token, str) and token.startswith("--output-directory="): + return Path(token.split("=", 1)[1]).expanduser().resolve() + except Exception: + return None + return None + + +def _detect_comfy_root(current_file: Path) -> Path | None: + for parent in current_file.parents: + if (parent / "main.py").is_file() and (parent / "folder_paths.py").is_file(): + return parent + if len(current_file.parents) > 3: + return current_file.parents[3] + return None + OUTPUT_ROOT_PATH = _resolve_output_root() OUTPUT_ROOT = str(OUTPUT_ROOT_PATH) diff --git a/mjr_am_backend/custom_roots.py b/mjr_am_backend/custom_roots.py index 03c3fe77..fd270f66 100644 --- a/mjr_am_backend/custom_roots.py +++ b/mjr_am_backend/custom_roots.py @@ -13,7 +13,7 @@ import os import threading from datetime import datetime, timezone from pathlib import Path -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Tuple from uuid import uuid4 from .shared import Result, get_logger @@ -133,64 +133,122 @@ def list_custom_roots() -> Result[List[Dict[str, Any]]]: roots = store.get("roots") or [] cleaned: List[Dict[str, Any]] = [] for r in roots: - if not isinstance(r, dict): - continue - rid = str(r.get("id") or "").strip() - path = str(r.get("path") or "").strip() - if not rid or not path: - continue - normalized = _normalize_dir_path(path) - if not normalized: - cleaned.append( - { - "id": rid, - "path": path, - "label": str(r.get("label") or "").strip() or path, - "created_at": r.get("created_at"), - "offline": True, - "invalid": True, - } - ) - continue - exists = False - is_dir = False - try: - exists = normalized.exists() - is_dir = normalized.is_dir() - except Exception: - exists = False - is_dir = False - offline = not (exists and is_dir) - cleaned.append( - { - "id": rid, - "path": str(normalized), - "label": str(r.get("label") or "").strip() or normalized.name or str(normalized), - "created_at": r.get("created_at"), - "offline": bool(offline), - } - ) + normalized_row = _normalize_custom_root_row(r) + if normalized_row is not None: + cleaned.append(normalized_row) return Result.Ok(cleaned) +def _normalize_custom_root_row(row: Any) -> Optional[Dict[str, Any]]: + if not isinstance(row, dict): + return None + rid = str(row.get("id") or "").strip() + path = str(row.get("path") or "").strip() + if not rid or not path: + return None + normalized = _normalize_dir_path(path) + if not normalized: + return { + "id": rid, + "path": path, + "label": str(row.get("label") or "").strip() or path, + "created_at": row.get("created_at"), + "offline": True, + "invalid": True, + } + return _normalized_root_payload(row, rid=rid, normalized=normalized) + + +def _normalized_root_payload(row: Dict[str, Any], *, rid: str, normalized: Path) -> Dict[str, Any]: + exists, is_dir = _path_exists_and_is_dir(normalized) + offline = not (exists and is_dir) + return { + "id": rid, + "path": str(normalized), + "label": str(row.get("label") or "").strip() or normalized.name or str(normalized), + "created_at": row.get("created_at"), + "offline": bool(offline), + } + + +def _path_exists_and_is_dir(path: Path) -> Tuple[bool, bool]: + try: + return path.exists(), path.is_dir() + except Exception: + return False, False + + def add_custom_root(path: str, label: Optional[str] = None) -> Result[Dict[str, Any]]: """Add a custom root directory, or return the existing one if present.""" normalized = _normalize_dir_path(path) + validation = _validate_new_root_path(normalized) + if validation is not None: + return validation + + payload = _new_root_payload(normalized, label=label) + resolved = payload["resolved"] + normalized_key = payload["normalized_key"] + root_id = payload["root_id"] + created_at = payload["created_at"] + safe_label = payload["safe_label"] + + output_root, input_root = _resolve_builtin_roots() + overlap_err = _check_builtin_root_overlap(normalized, output_root, input_root) + if overlap_err is not None: + return overlap_err + + with _LOCK: + store = _read_store() + roots = store.get("roots") or [] + existing_result = _find_existing_or_overlap_root( + roots=roots, + normalized=normalized, + normalized_key=normalized_key, + fallback_id=root_id, + resolved=resolved, + safe_label=safe_label, + created_at=created_at, + ) + if existing_result is not None: + return existing_result + + roots.append(_root_row(root_id, resolved, safe_label, created_at)) + store["roots"] = roots + write_result = _write_store(store) + if not write_result.ok: + return write_result # type: ignore[return-value] + + return Result.Ok(_root_row(root_id, resolved, safe_label, created_at)) + + +def _validate_new_root_path(normalized: Optional[Path]) -> Optional[Result[Dict[str, Any]]]: if not normalized: return Result.Err("INVALID_INPUT", "Invalid path") if not normalized.exists(): return Result.Err("DIR_NOT_FOUND", f"Directory not found: {normalized}") if not normalized.is_dir(): return Result.Err("NOT_A_DIRECTORY", f"Not a directory: {normalized}") + return None + +def _new_root_payload(normalized: Path, *, label: Optional[str]) -> Dict[str, str]: resolved = str(normalized) - normalized_key = _canonical_path_key(resolved) - root_id = str(uuid4()) - created_at = _utc_now_iso() - safe_label = (label or "").strip() or normalized.name or resolved + return { + "resolved": resolved, + "normalized_key": _canonical_path_key(resolved), + "root_id": str(uuid4()), + "created_at": _utc_now_iso(), + "safe_label": (label or "").strip() or normalized.name or resolved, + } - output_root = None - input_root = None + +def _root_row(root_id: str, resolved: str, safe_label: str, created_at: str) -> Dict[str, Any]: + return {"id": root_id, "path": resolved, "label": safe_label, "created_at": created_at} + + +def _resolve_builtin_roots() -> tuple[Optional[Path], Optional[Path]]: + output_root: Optional[Path] = None + input_root: Optional[Path] = None try: if OUTPUT_ROOT: output_root = Path(OUTPUT_ROOT).resolve() @@ -198,50 +256,67 @@ def add_custom_root(path: str, label: Optional[str] = None) -> Result[Dict[str, output_root = None try: import folder_paths + input_root = Path(folder_paths.get_input_directory()).resolve() except Exception: input_root = None + return output_root, input_root - if output_root: - if _path_is_relative_to(normalized, output_root) or _path_is_relative_to(output_root, normalized): - return Result.Err("OVERLAP", f"Root overlaps with output: {output_root}") - if input_root: - if _path_is_relative_to(normalized, input_root) or _path_is_relative_to(input_root, normalized): - return Result.Err("OVERLAP", f"Root overlaps with input: {input_root}") - with _LOCK: - store = _read_store() - roots = store.get("roots") or [] - for r in roots: - if not isinstance(r, dict): - continue - existing_path = r.get("path") - if existing_path and _canonical_path_key(str(existing_path)) == normalized_key: - existing_id = str(r.get("id") or "") - return Result.Ok( - { - "id": existing_id or root_id, - "path": resolved, - "label": str(r.get("label") or safe_label), - "created_at": r.get("created_at") or created_at, - "already_exists": True, - } - ) - if existing_path: - try: - existing_resolved = Path(str(existing_path)).expanduser().resolve(strict=False) - if _path_is_relative_to(normalized, existing_resolved) or _path_is_relative_to(existing_resolved, normalized): - return Result.Err("OVERLAP", f"Root overlaps with existing: {existing_resolved}") - except Exception: - continue +def _check_builtin_root_overlap( + normalized: Path, + output_root: Optional[Path], + input_root: Optional[Path], +) -> Optional[Result[Dict[str, Any]]]: + if output_root and (_path_is_relative_to(normalized, output_root) or _path_is_relative_to(output_root, normalized)): + return Result.Err("OVERLAP", f"Root overlaps with output: {output_root}") + if input_root and (_path_is_relative_to(normalized, input_root) or _path_is_relative_to(input_root, normalized)): + return Result.Err("OVERLAP", f"Root overlaps with input: {input_root}") + return None - roots.append({"id": root_id, "path": resolved, "label": safe_label, "created_at": created_at}) - store["roots"] = roots - write_result = _write_store(store) - if not write_result.ok: - return write_result # type: ignore[return-value] - return Result.Ok({"id": root_id, "path": resolved, "label": safe_label, "created_at": created_at}) +def _find_existing_or_overlap_root( + *, + roots: List[Any], + normalized: Path, + normalized_key: str, + fallback_id: str, + resolved: str, + safe_label: str, + created_at: str, +) -> Optional[Result[Dict[str, Any]]]: + for root in roots: + if not isinstance(root, dict): + continue + existing_path = root.get("path") + same_path = existing_path and _canonical_path_key(str(existing_path)) == normalized_key + if same_path: + existing_id = str(root.get("id") or "") + return Result.Ok( + { + "id": existing_id or fallback_id, + "path": resolved, + "label": str(root.get("label") or safe_label), + "created_at": root.get("created_at") or created_at, + "already_exists": True, + } + ) + overlap = _existing_root_overlap(normalized, existing_path) + if overlap: + return Result.Err("OVERLAP", f"Root overlaps with existing: {overlap}") + return None + + +def _existing_root_overlap(normalized: Path, existing_path: Any) -> Optional[Path]: + if not existing_path: + return None + try: + existing_resolved = Path(str(existing_path)).expanduser().resolve(strict=False) + if _path_is_relative_to(normalized, existing_resolved) or _path_is_relative_to(existing_resolved, normalized): + return existing_resolved + except Exception: + return None + return None def remove_custom_root(root_id: str) -> Result[bool]: @@ -273,15 +348,26 @@ def resolve_custom_root(root_id: str) -> Result[Path]: roots_result = list_custom_roots() if not roots_result.ok: return Result.Err(roots_result.code, roots_result.error or "Failed to list custom roots") - for r in roots_result.data or []: - if str(r.get("id") or "") == rid: - normalized = _normalize_dir_path(str(r.get("path") or "")) - if not normalized: - return Result.Err("INVALID_INPUT", "Invalid stored path") - try: - if not normalized.exists() or not normalized.is_dir(): - return Result.Err("OFFLINE", f"Custom root is offline: {normalized}") - except Exception: - return Result.Err("OFFLINE", f"Custom root is offline: {normalized}") - return Result.Ok(normalized) + root_row = _find_custom_root_row_by_id(roots_result.data or [], rid) + if root_row is not None: + return _resolve_custom_root_path_from_row(root_row) return Result.Err("NOT_FOUND", f"Custom root not found: {rid}") + + +def _find_custom_root_row_by_id(rows: List[Dict[str, Any]], rid: str) -> Optional[Dict[str, Any]]: + for row in rows: + if str(row.get("id") or "") == rid: + return row + return None + + +def _resolve_custom_root_path_from_row(row: Dict[str, Any]) -> Result[Path]: + normalized = _normalize_dir_path(str(row.get("path") or "")) + if not normalized: + return Result.Err("INVALID_INPUT", "Invalid stored path") + try: + if not normalized.exists() or not normalized.is_dir(): + return Result.Err("OFFLINE", f"Custom root is offline: {normalized}") + except Exception: + return Result.Err("OFFLINE", f"Custom root is offline: {normalized}") + return Result.Ok(normalized) diff --git a/mjr_am_backend/deps.py b/mjr_am_backend/deps.py index ef9aeb86..5f8b0afa 100644 --- a/mjr_am_backend/deps.py +++ b/mjr_am_backend/deps.py @@ -32,6 +32,66 @@ from .config import ( logger = get_logger(__name__) +def _resolve_db_path(db_path: Optional[str]) -> str: + return db_path if db_path is not None else INDEX_DB + + +def _init_db_or_error(db_path: str) -> Result[Sqlite]: + logger.info(f"Initializing database: {db_path}") + try: + return Result.Ok(Sqlite(db_path, max_connections=DB_MAX_CONNECTIONS, timeout=DB_TIMEOUT)) + except Exception as exc: + logger.error("Failed to initialize database: %s", exc) + return Result.Err("DB_ERROR", f"Failed to initialize database: {exc}") + + +async def _migrate_db_or_error(db: Sqlite) -> Result[bool]: + migrate_result = await migrate_schema(db) + if not migrate_result.ok: + logger.error(f"Schema migration failed: {migrate_result.error}") + return Result.Err(migrate_result.code or "DB_ERROR", f"Failed to initialize database: {migrate_result.error}") + return Result.Ok(True) + + +def _init_tools_and_settings(db: Sqlite) -> tuple[ExifTool, FFProbe, AppSettings]: + exiftool = ExifTool(bin_name=EXIFTOOL_BIN or "exiftool", timeout=EXIFTOOL_TIMEOUT) + ffprobe = FFProbe(bin_name=FFPROBE_BIN or "ffprobe", timeout=FFPROBE_TIMEOUT) + settings_service = AppSettings(db) + return exiftool, ffprobe, settings_service + + +def _log_tool_availability(exiftool: ExifTool, ffprobe: FFProbe) -> None: + if exiftool.is_available(): + log_success(logger, "ExifTool is available") + else: + logger.warning("ExifTool not found - metadata extraction will be limited") + if ffprobe.is_available(): + log_success(logger, "ffprobe is available") + else: + logger.warning("ffprobe not found - video metadata will be limited") + + +def _build_services_dict( + db: Sqlite, + exiftool: ExifTool, + ffprobe: FFProbe, + metadata_service: MetadataService, + health_service: HealthService, + index_service: IndexService, + settings_service: AppSettings, +) -> dict: + return { + "db": db, + "exiftool": exiftool, + "ffprobe": ffprobe, + "metadata": metadata_service, + "health": health_service, + "index": index_service, + "settings": settings_service, + "duplicates": DuplicatesService(db), + } + + async def build_services(db_path: Optional[str] = None) -> Result[dict]: """ Build all services (DI container). @@ -44,43 +104,23 @@ async def build_services(db_path: Optional[str] = None) -> Result[dict]: """ logger.info("Building services...") - # Default database path from config - if db_path is None: - db_path = INDEX_DB + db_path = _resolve_db_path(db_path) + db_res = _init_db_or_error(db_path) + if not db_res.ok: + return Result.Err(db_res.code or "DB_ERROR", db_res.error or "Failed to initialize database") + db = db_res.data - # Initialize database - logger.info(f"Initializing database: {db_path}") - try: - db = Sqlite(db_path, max_connections=DB_MAX_CONNECTIONS, timeout=DB_TIMEOUT) - except Exception as exc: - logger.error("Failed to initialize database: %s", exc) - return Result.Err("DB_ERROR", f"Failed to initialize database: {exc}") - - # Run migrations - migrate_result = await migrate_schema(db) + migrate_result = await _migrate_db_or_error(db) if not migrate_result.ok: - logger.error(f"Schema migration failed: {migrate_result.error}") - return Result.Err(migrate_result.code or "DB_ERROR", f"Failed to initialize database: {migrate_result.error}") + return migrate_result # type: ignore[return-value] - # Initialize adapters - exiftool = ExifTool(bin_name=EXIFTOOL_BIN or "exiftool", timeout=EXIFTOOL_TIMEOUT) - ffprobe = FFProbe(bin_name=FFPROBE_BIN or "ffprobe", timeout=FFPROBE_TIMEOUT) - settings_service = AppSettings(db) + exiftool, ffprobe, settings_service = _init_tools_and_settings(db) try: await settings_service.ensure_security_bootstrap() except Exception as exc: logger.warning("Security bootstrap failed: %s", exc) - # Log tool availability - if exiftool.is_available(): - log_success(logger, "ExifTool is available") - else: - logger.warning("ExifTool not found - metadata extraction will be limited") - - if ffprobe.is_available(): - log_success(logger, "ffprobe is available") - else: - logger.warning("ffprobe not found - video metadata will be limited") + _log_tool_availability(exiftool, ffprobe) # Build services metadata_service = MetadataService( @@ -97,23 +137,22 @@ async def build_services(db_path: Optional[str] = None) -> Result[dict]: # Check for optional columns matches = await table_has_column(db, "asset_metadata", "tags_text") - + index_service = IndexService( db=db, metadata_service=metadata_service, has_tags_text_column=matches ) - services = { - "db": db, - "exiftool": exiftool, - "ffprobe": ffprobe, - "metadata": metadata_service, - "health": health_service, - "index": index_service, - "settings": settings_service, - "duplicates": DuplicatesService(db), - } + services = _build_services_dict( + db, + exiftool, + ffprobe, + metadata_service, + health_service, + index_service, + settings_service, + ) try: services["watcher_scope"] = await load_watcher_scope(db) except Exception: diff --git a/mjr_am_backend/features/audio/metadata.py b/mjr_am_backend/features/audio/metadata.py index 1bffed13..11658a12 100644 --- a/mjr_am_backend/features/audio/metadata.py +++ b/mjr_am_backend/features/audio/metadata.py @@ -28,19 +28,51 @@ def _extract_json_fields(container: Any) -> Tuple[Optional[Dict[str, Any]], Opti if not isinstance(parsed, dict): continue key_lower = str(key or "").lower() - if workflow is None and looks_like_comfyui_workflow(parsed): - workflow = parsed - if prompt is None and looks_like_comfyui_prompt_graph(parsed): - prompt = parsed - if workflow is None and "workflow" in key_lower and isinstance(parsed, dict): - workflow = parsed if looks_like_comfyui_workflow(parsed) else workflow - if prompt is None and "prompt" in key_lower and isinstance(parsed, dict): - prompt = parsed if looks_like_comfyui_prompt_graph(parsed) else prompt + workflow, prompt = _merge_workflow_prompt_candidate(workflow, prompt, parsed, key_lower) if workflow is not None and prompt is not None: break return (workflow, prompt) +def _merge_workflow_prompt_candidate( + workflow: Optional[Dict[str, Any]], + prompt: Optional[Dict[str, Any]], + parsed: Dict[str, Any], + key_lower: str, +) -> Tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]]]: + workflow = _workflow_candidate_or_existing(workflow, parsed, key_lower) + prompt = _prompt_candidate_or_existing(prompt, parsed, key_lower) + return workflow, prompt + + +def _workflow_candidate_or_existing( + workflow: Optional[Dict[str, Any]], + parsed: Dict[str, Any], + key_lower: str, +) -> Optional[Dict[str, Any]]: + if workflow is not None: + return workflow + if looks_like_comfyui_workflow(parsed): + return parsed + if "workflow" in key_lower and looks_like_comfyui_workflow(parsed): + return parsed + return workflow + + +def _prompt_candidate_or_existing( + prompt: Optional[Dict[str, Any]], + parsed: Dict[str, Any], + key_lower: str, +) -> Optional[Dict[str, Any]]: + if prompt is not None: + return prompt + if looks_like_comfyui_prompt_graph(parsed): + return parsed + if "prompt" in key_lower and looks_like_comfyui_prompt_graph(parsed): + return parsed + return prompt + + def _unwrap_workflow_prompt_container(container: Any) -> Tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]]]: """ Handle wrapper payloads embedded in a single text field/tag, e.g: @@ -51,21 +83,32 @@ def _unwrap_workflow_prompt_container(container: Any) -> Tuple[Optional[Dict[str wf = container.get("workflow") or container.get("Workflow") or None pr = container.get("prompt") or container.get("Prompt") or None + return _coerce_wrapped_workflow_prompt(wf, pr) - wf_out: Optional[Dict[str, Any]] = None - pr_out: Optional[Dict[str, Any]] = None - if isinstance(wf, dict) and looks_like_comfyui_workflow(wf): - wf_out = wf +def _coerce_wrapped_workflow_prompt( + workflow_value: Any, + prompt_value: Any, +) -> Tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]]]: + wf_out = _coerce_wrapped_workflow(workflow_value) + pr_out = _coerce_wrapped_prompt(prompt_value) + return wf_out, pr_out - if isinstance(pr, dict) and looks_like_comfyui_prompt_graph(pr): - pr_out = pr - elif isinstance(pr, str): - parsed = try_parse_json_text(pr) + +def _coerce_wrapped_workflow(workflow_value: Any) -> Optional[Dict[str, Any]]: + if isinstance(workflow_value, dict) and looks_like_comfyui_workflow(workflow_value): + return workflow_value + return None + + +def _coerce_wrapped_prompt(prompt_value: Any) -> Optional[Dict[str, Any]]: + if isinstance(prompt_value, dict) and looks_like_comfyui_prompt_graph(prompt_value): + return prompt_value + if isinstance(prompt_value, str): + parsed = try_parse_json_text(prompt_value) if isinstance(parsed, dict) and looks_like_comfyui_prompt_graph(parsed): - pr_out = parsed - - return (wf_out, pr_out) + return parsed + return None def _scan_for_workflow_prompt(*containers: Any) -> Tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]]]: @@ -79,33 +122,169 @@ def _scan_for_workflow_prompt(*containers: Any) -> Tuple[Optional[Dict[str, Any] if not isinstance(container, dict): continue - # 1) Fast key-based parse first - wf, pr = _extract_json_fields(container) - if workflow is None: - workflow = wf - if prompt is None: - prompt = pr - - # 2) Wrapper scan for JSON values under comments/descriptions - if workflow is None or prompt is None: - for _, value in container.items(): - parsed = parse_json_value(value) - if not isinstance(parsed, dict): - continue - wf2, pr2 = _unwrap_workflow_prompt_container(parsed) - if workflow is None and wf2 is not None: - workflow = wf2 - if prompt is None and pr2 is not None: - prompt = pr2 - if workflow is not None and prompt is not None: - break - + workflow, prompt = _scan_single_workflow_prompt_container(container, workflow=workflow, prompt=prompt) if workflow is not None and prompt is not None: break return (workflow, prompt) +def _scan_single_workflow_prompt_container( + container: Dict[str, Any], + *, + workflow: Optional[Dict[str, Any]], + prompt: Optional[Dict[str, Any]], +) -> Tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]]]: + wf, pr = _extract_json_fields(container) + workflow = workflow or wf + prompt = prompt or pr + if workflow is not None and prompt is not None: + return workflow, prompt + return _scan_wrapper_payload_candidates(container, workflow=workflow, prompt=prompt) + + +def _scan_wrapper_payload_candidates( + container: Dict[str, Any], + *, + workflow: Optional[Dict[str, Any]], + prompt: Optional[Dict[str, Any]], +) -> Tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]]]: + for value in container.values(): + parsed = parse_json_value(value) + if not isinstance(parsed, dict): + continue + wf2, pr2 = _unwrap_workflow_prompt_container(parsed) + workflow = workflow or wf2 + prompt = prompt or pr2 + if workflow is not None and prompt is not None: + break + return workflow, prompt + + +def _init_audio_metadata( + exif_data: Dict[str, Any], + ffprobe_data: Dict[str, Any], +) -> Dict[str, Any]: + return { + "raw": exif_data, + "raw_ffprobe": ffprobe_data, + "workflow": None, + "prompt": None, + "parameters": None, + "duration": None, + "audio_codec": None, + "sample_rate": None, + "channels": None, + "bitrate": None, + "quality": "none", + } + + +def _extract_ffprobe_parts(ffprobe_data: Dict[str, Any]) -> tuple[Dict[str, Any], Dict[str, Any], Dict[str, Any], List[Dict[str, Any]]]: + if not isinstance(ffprobe_data, dict): + return {}, {}, {}, [] + + audio_stream = _dict_field(ffprobe_data, "audio_stream") + fmt = _dict_field(ffprobe_data, "format") + format_tags = _dict_field(fmt, "tags") + stream_tag_dicts = _stream_tag_dicts(ffprobe_data.get("streams")) + return audio_stream, fmt, format_tags, stream_tag_dicts + + +def _dict_field(container: Any, key: str) -> Dict[str, Any]: + if not isinstance(container, dict): + return {} + value = container.get(key) or {} + return value if isinstance(value, dict) else {} + + +def _stream_tag_dicts(streams: Any) -> List[Dict[str, Any]]: + tags_list: List[Dict[str, Any]] = [] + if not isinstance(streams, list): + return tags_list + for stream in streams: + if not isinstance(stream, dict): + continue + tags = stream.get("tags") + if isinstance(tags, dict): + tags_list.append(tags) + return tags_list + + +def _apply_audio_technical_fields( + metadata: Dict[str, Any], + audio_stream: Dict[str, Any], + fmt: Dict[str, Any], +) -> None: + if isinstance(audio_stream, dict): + metadata["audio_codec"] = audio_stream.get("codec_name") + metadata["sample_rate"] = audio_stream.get("sample_rate") + metadata["channels"] = audio_stream.get("channels") + metadata["bitrate"] = audio_stream.get("bit_rate") + metadata["duration"] = audio_stream.get("duration") + if metadata.get("duration") is None and isinstance(fmt, dict): + metadata["duration"] = fmt.get("duration") + if metadata.get("bitrate") is None: + metadata["bitrate"] = fmt.get("bit_rate") + + +def _apply_workflow_prompt_quality( + metadata: Dict[str, Any], + workflow: Optional[Dict[str, Any]], + prompt: Optional[Dict[str, Any]], +) -> None: + if workflow is not None: + metadata["workflow"] = workflow + metadata["quality"] = "full" + if prompt is not None: + metadata["prompt"] = prompt + if metadata.get("quality") != "full": + metadata["quality"] = "partial" + + +def _collect_audio_text_candidates( + exif_data: Dict[str, Any], + fmt: Dict[str, Any], +) -> List[str]: + text_candidates: List[str] = [] + if isinstance(exif_data, dict): + text_candidates.extend(v for v in exif_data.values() if isinstance(v, str)) + if isinstance(fmt, dict) and isinstance(fmt.get("tags"), dict): + tags_dict = fmt.get("tags") + if isinstance(tags_dict, dict): + text_candidates.extend(v for v in tags_dict.values() if isinstance(v, str)) + return text_candidates + + +def _apply_auto1111_candidates( + metadata: Dict[str, Any], + text_candidates: List[str], +) -> None: + for text in text_candidates: + text_s = str(text or "").strip() + if text_s.startswith("{") or text_s.startswith("["): + continue + parsed = parse_auto1111_params(text) + if not parsed: + continue + metadata["parameters"] = text + for key, value in parsed.items(): + if key in ("workflow", "prompt") and isinstance(metadata.get(key), dict): + continue + metadata[key] = value + if metadata.get("quality") != "full": + metadata["quality"] = "partial" + break + + +def _apply_audio_rating_tags(metadata: Dict[str, Any], exif_data: Dict[str, Any]) -> None: + rating, tags = extract_rating_tags_from_exif(exif_data) + if rating is not None: + metadata["rating"] = rating + if tags: + metadata["tags"] = tags + + def extract_audio_metadata( file_path: str, exif_data: Optional[Dict[str, Any]] = None, @@ -119,90 +298,17 @@ def extract_audio_metadata( exif_data = exif_data or {} ffprobe_data = ffprobe_data or {} - metadata: Dict[str, Any] = { - "raw": exif_data, - "raw_ffprobe": ffprobe_data, - "workflow": None, - "prompt": None, - "parameters": None, - "duration": None, - "audio_codec": None, - "sample_rate": None, - "channels": None, - "bitrate": None, - "quality": "none", - } + metadata: Dict[str, Any] = _init_audio_metadata(exif_data, ffprobe_data) try: - audio_stream: Dict[str, Any] = {} - fmt: Dict[str, Any] = {} - if isinstance(ffprobe_data, dict): - audio_stream = ffprobe_data.get("audio_stream") or {} - fmt = ffprobe_data.get("format") or {} - if isinstance(audio_stream, dict): - metadata["audio_codec"] = audio_stream.get("codec_name") - metadata["sample_rate"] = audio_stream.get("sample_rate") - metadata["channels"] = audio_stream.get("channels") - metadata["bitrate"] = audio_stream.get("bit_rate") - metadata["duration"] = audio_stream.get("duration") - if metadata["duration"] is None and isinstance(fmt, dict): - metadata["duration"] = fmt.get("duration") - if metadata["bitrate"] is None: - metadata["bitrate"] = fmt.get("bit_rate") - - format_tags = fmt.get("tags") if isinstance(fmt, dict) and isinstance(fmt.get("tags"), dict) else {} - stream_tag_dicts = [] - if isinstance(ffprobe_data, dict) and isinstance(ffprobe_data.get("streams"), list): - for stream in ffprobe_data.get("streams") or []: - if not isinstance(stream, dict): - continue - tags = stream.get("tags") - if isinstance(tags, dict): - stream_tag_dicts.append(tags) - + audio_stream, fmt, format_tags, stream_tag_dicts = _extract_ffprobe_parts(ffprobe_data) + _apply_audio_technical_fields(metadata, audio_stream, fmt) workflow, prompt = _scan_for_workflow_prompt(exif_data, format_tags, *stream_tag_dicts) - - if workflow is not None: - metadata["workflow"] = workflow - metadata["quality"] = "full" - if prompt is not None: - metadata["prompt"] = prompt - if metadata["quality"] != "full": - metadata["quality"] = "partial" - - text_candidates: List[str] = [] - if isinstance(exif_data, dict): - text_candidates.extend(v for v in exif_data.values() if isinstance(v, str)) - if isinstance(fmt, dict) and isinstance(fmt.get("tags"), dict): - tags_dict = fmt.get("tags") - if isinstance(tags_dict, dict): - text_candidates.extend(v for v in tags_dict.values() if isinstance(v, str)) - for text in text_candidates: - t = str(text or "").strip() - if t.startswith("{") or t.startswith("["): - continue - parsed = parse_auto1111_params(text) - if not parsed: - continue - metadata["parameters"] = text - # Preserve already-parsed ComfyUI structures. Some free-text fields can - # produce a "prompt" string via A1111 parser and must not clobber a real - # prompt graph dict extracted from metadata tags. - for key, value in parsed.items(): - if key in ("workflow", "prompt") and isinstance(metadata.get(key), dict): - continue - metadata[key] = value - if metadata["quality"] != "full": - metadata["quality"] = "partial" - break - - rating, tags = extract_rating_tags_from_exif(exif_data) - if rating is not None: - metadata["rating"] = rating - if tags: - metadata["tags"] = tags - - if metadata["quality"] == "none" and (exif_data or ffprobe_data): + _apply_workflow_prompt_quality(metadata, workflow, prompt) + text_candidates = _collect_audio_text_candidates(exif_data, fmt) + _apply_auto1111_candidates(metadata, text_candidates) + _apply_audio_rating_tags(metadata, exif_data) + if metadata.get("quality") == "none" and (exif_data or ffprobe_data): metadata["quality"] = "partial" return Result.Ok(metadata, quality=metadata["quality"]) diff --git a/mjr_am_backend/features/browser/service.py b/mjr_am_backend/features/browser/service.py index a7e05527..b08f13fa 100644 --- a/mjr_am_backend/features/browser/service.py +++ b/mjr_am_backend/features/browser/service.py @@ -7,6 +7,7 @@ from __future__ import annotations import os import stat from pathlib import Path +from typing import Any from mjr_am_backend.shared import Result, classify_file from mjr_am_backend.routes.core import _is_within_root, _safe_rel_path @@ -33,59 +34,17 @@ def list_visible_subfolders(root_dir: Path, subfolder: str, root_id: str) -> Res """ List physical subfolders under `root_dir/subfolder` while enforcing root confinement. """ - try: - base = root_dir.resolve() - except OSError as exc: - return Result.Err("INVALID_INPUT", f"Invalid root directory: {exc}") - - rel = _safe_rel_path(subfolder or "") - if rel is None: - return Result.Err("INVALID_INPUT", "Invalid subfolder") - - target_dir = (base / rel) - try: - target_resolved = target_dir.resolve(strict=True) - except OSError: - return Result.Err("DIR_NOT_FOUND", f"Directory not found: {target_dir}") - - if not _is_within_root(target_resolved, base): - return Result.Err("INVALID_INPUT", "Subfolder outside root") + resolved = _resolve_browser_target_dir(root_dir, subfolder) + if not resolved.ok: + return resolved + base, target_resolved = resolved.data folders: list[dict] = [] try: for entry in target_resolved.iterdir(): - try: - if not entry.is_dir(): - continue - if _is_hidden_directory(entry): - continue - resolved_child = entry.resolve(strict=True) - if not _is_within_root(resolved_child, base): - continue - child_rel = str(resolved_child.relative_to(base)).replace("\\", "/") - folders.append( - { - "id": f"folder::{root_id}::{child_rel}", - "filename": entry.name, - "subfolder": child_rel, - "filepath": str(resolved_child), - "kind": "folder", - "type": "custom", - "root_id": str(root_id or ""), - "mtime": int(entry.stat().st_mtime), - "size": None, - "ext": "", - "rating": 0, - "tags": [], - "has_workflow": None, - "has_generation_data": None, - "width": None, - "height": None, - "duration": None, - } - ) - except Exception: - continue + row = _subfolder_row(entry, base, root_id) + if row is not None: + folders.append(row) except OSError as exc: return Result.Err("LIST_FAILED", f"Failed to list subfolders: {exc}") @@ -93,6 +52,55 @@ def list_visible_subfolders(root_dir: Path, subfolder: str, root_id: str) -> Res return Result.Ok(folders) +def _resolve_browser_target_dir(root_dir: Path, subfolder: str) -> Result[tuple[Path, Path]]: + try: + base = root_dir.resolve() + except OSError as exc: + return Result.Err("INVALID_INPUT", f"Invalid root directory: {exc}") + rel = _safe_rel_path(subfolder or "") + if rel is None: + return Result.Err("INVALID_INPUT", "Invalid subfolder") + target_dir = base / rel + try: + target_resolved = target_dir.resolve(strict=True) + except OSError: + return Result.Err("DIR_NOT_FOUND", f"Directory not found: {target_dir}") + if not _is_within_root(target_resolved, base): + return Result.Err("INVALID_INPUT", "Subfolder outside root") + return Result.Ok((base, target_resolved)) + + +def _subfolder_row(entry: Path, base: Path, root_id: str) -> dict | None: + try: + if not entry.is_dir() or _is_hidden_directory(entry): + return None + resolved_child = entry.resolve(strict=True) + if not _is_within_root(resolved_child, base): + return None + child_rel = str(resolved_child.relative_to(base)).replace("\\", "/") + return { + "id": f"folder::{root_id}::{child_rel}", + "filename": entry.name, + "subfolder": child_rel, + "filepath": str(resolved_child), + "kind": "folder", + "type": "custom", + "root_id": str(root_id or ""), + "mtime": int(entry.stat().st_mtime), + "size": None, + "ext": "", + "rating": 0, + "tags": [], + "has_workflow": None, + "has_generation_data": None, + "width": None, + "height": None, + "duration": None, + } + except Exception: + return None + + def _safe_name(path: Path) -> str: try: if path.name: @@ -102,6 +110,340 @@ def _safe_name(path: Path) -> str: return str(path) +def _normalize_browser_query_filters( + query: str, + kind_filter: str, + min_size_bytes: int, + max_size_bytes: int, + min_width: int, + min_height: int, + max_width: int, + max_height: int, +) -> dict[str, int | str | bool]: + q = str(query or "*").strip() + ql = q.lower() + browse_all = q in ("", "*") + filter_kind = str(kind_filter or "").strip().lower() + min_size, max_size = _normalized_size_bounds(min_size_bytes, max_size_bytes) + return { + "q": q, + "ql": ql, + "browse_all": browse_all, + "filter_kind": filter_kind, + "min_size": min_size, + "max_size": max_size, + "min_w": _nonnegative_int(min_width), + "min_h": _nonnegative_int(min_height), + "max_w": _nonnegative_int(max_width), + "max_h": _nonnegative_int(max_height), + } + + +def _nonnegative_int(value: Any) -> int: + try: + return max(0, int(value or 0)) + except Exception: + return 0 + + +def _normalized_size_bounds(min_size_bytes: int, max_size_bytes: int) -> tuple[int, int]: + min_size = _nonnegative_int(min_size_bytes) + max_size = _nonnegative_int(max_size_bytes) + if max_size > 0 and min_size > 0 and max_size < min_size: + max_size = min_size + return min_size, max_size + + +def _paginate_assets(assets: list[dict], limit: int, offset: int) -> list[dict]: + start = max(0, int(offset or 0)) + end = start + int(limit or 0) if int(limit or 0) > 0 else None + return assets[start:end] if end is not None else assets[start:] + + +def _build_browser_listing_payload( + assets: list[dict], + limit: int, + offset: int, + query: str, + *, + target: Path | None = None, +) -> dict: + payload: dict = { + "assets": assets, + "total": len(assets), + "limit": limit, + "offset": offset, + "query": query, + "scope": "custom", + } + if isinstance(target, Path): + payload["path"] = str(target) + payload["current_path"] = str(target) + payload["display_name"] = _safe_name(target) + return payload + + +def _make_browser_folder_entry(name: str, entry: Path, resolved_child: Path) -> dict: + return { + "id": f"folder::fs::{resolved_child}", + "filename": name, + "subfolder": str(resolved_child), + "filepath": str(resolved_child), + "kind": "folder", + "type": "custom", + "root_id": "", + "mtime": int(entry.stat().st_mtime), + "size": None, + "ext": "", + "rating": 0, + "tags": [], + "has_workflow": None, + "has_generation_data": None, + "width": None, + "height": None, + "duration": None, + } + + +def _make_browser_file_entry(target: Path, name: str, entry: Path, kind: str, size: int, mtime: int) -> dict: + return { + "id": None, + "filename": name, + "subfolder": str(target), + "filepath": str(entry), + "kind": kind, + "ext": entry.suffix.lower(), + "size": size, + "mtime": mtime, + "width": None, + "height": None, + "duration": None, + "rating": 0, + "tags": [], + "has_workflow": None, + "has_generation_data": None, + "type": "custom", + "root_id": "", + } + + +def _filter_browser_roots(roots: list[dict], *, filter_kind: str, browse_all: bool, ql: str) -> list[dict]: + if filter_kind: + return [] + if browse_all: + return roots + return [f for f in roots if ql in str(f.get("filename") or "").lower()] + + +def _collect_browser_folders_and_files( + target: Path, + *, + browse_all: bool, + ql: str, + filter_kind: str, + min_size: int, + max_size: int, +) -> Result[dict[str, list[dict]]]: + folders: list[dict] = [] + files: list[dict] = [] + try: + for entry in target.iterdir(): + try: + _collect_browser_single_entry( + entry=entry, + target=target, + folders=folders, + files=files, + browse_all=browse_all, + ql=ql, + filter_kind=filter_kind, + min_size=min_size, + max_size=max_size, + ) + except Exception: + continue + except OSError as exc: + return Result.Err("LIST_FAILED", f"Failed to list directory: {exc}") + return Result.Ok({"folders": folders, "files": files}) + + +def _collect_browser_single_entry( + *, + entry: Path, + target: Path, + folders: list[dict], + files: list[dict], + browse_all: bool, + ql: str, + filter_kind: str, + min_size: int, + max_size: int, +) -> None: + name = entry.name + if _skip_browser_entry_name(name, browse_all=browse_all, ql=ql): + return + if entry.is_dir(): + _append_browser_folder_if_visible(folders, name=name, entry=entry) + return + if not entry.is_file(): + return + kind = _browser_file_kind_or_none(name=name, filter_kind=filter_kind) + if kind is None: + return + st = entry.stat() + size = int(st.st_size or 0) + if not _size_within_bounds(size=size, min_size=min_size, max_size=max_size): + return + files.append(_make_browser_file_entry(target, name, entry, kind, size, int(st.st_mtime))) + + +def _skip_browser_entry_name(name: str, *, browse_all: bool, ql: str) -> bool: + if name.startswith("."): + return True + return (not browse_all) and (ql not in name.lower()) + + +def _append_browser_folder_if_visible(folders: list[dict], *, name: str, entry: Path) -> None: + if _is_hidden_directory(entry): + return + folders.append(_make_browser_folder_entry(name, entry, entry.resolve(strict=True))) + + +def _browser_file_kind_or_none(*, name: str, filter_kind: str) -> str | None: + kind = classify_file(name) + if kind == "unknown": + return None + if filter_kind and kind != filter_kind: + return None + return kind + + +def _size_within_bounds(*, size: int, min_size: int, max_size: int) -> bool: + if min_size > 0 and size < min_size: + return False + if max_size > 0 and size > max_size: + return False + return True + + +def _filter_browser_files_by_dimensions( + files: list[dict], + *, + min_w: int, + min_h: int, + max_w: int, + max_h: int, +) -> list[dict]: + if min_w <= 0 and min_h <= 0 and max_w <= 0 and max_h <= 0: + return files + + filtered_files: list[dict] = [] + for f in files: + if _browser_file_matches_dimensions(f, min_w=min_w, min_h=min_h, max_w=max_w, max_h=max_h): + filtered_files.append(f) + return filtered_files + + +def _browser_file_matches_dimensions( + file_row: dict, + *, + min_w: int, + min_h: int, + max_w: int, + max_h: int, +) -> bool: + w_raw = file_row.get("width") + h_raw = file_row.get("height") + if w_raw is None or h_raw is None: + return True + try: + w = int(w_raw or 0) + h = int(h_raw or 0) + except Exception: + return False + return _value_matches_bounds(value=w, min_value=min_w, max_value=max_w) and _value_matches_bounds( + value=h, min_value=min_h, max_value=max_h + ) + + +def _value_matches_bounds(*, value: int, min_value: int, max_value: int) -> bool: + if min_value > 0 and value < min_value: + return False + if max_value > 0 and value > max_value: + return False + return True + + +def _list_browser_root_mode_entries( + *, + q: str, + ql: str, + browse_all: bool, + filter_kind: str, + limit: int, + offset: int, +) -> Result[dict]: + roots_res = list_browser_roots() + if not roots_res.ok: + return Result.Err(roots_res.code, roots_res.error or "Failed to list roots") + roots = roots_res.data or [] + filtered = _filter_browser_roots( + roots, + filter_kind=filter_kind, + browse_all=browse_all, + ql=ql, + ) + payload = _build_browser_listing_payload(filtered, limit, offset, q) + payload["assets"] = _paginate_assets(filtered, limit, offset) + return Result.Ok(payload) + + +def _list_browser_path_mode_entries( + target: Path, + *, + q: str, + ql: str, + browse_all: bool, + filter_kind: str, + min_size: int, + max_size: int, + min_w: int, + min_h: int, + max_w: int, + max_h: int, + limit: int, + offset: int, +) -> Result[dict]: + collect_res = _collect_browser_folders_and_files( + target, + browse_all=browse_all, + ql=ql, + filter_kind=filter_kind, + min_size=min_size, + max_size=max_size, + ) + if not collect_res.ok: + return collect_res + data = collect_res.data if isinstance(collect_res.data, dict) else {} + folders = data.get("folders") if isinstance(data.get("folders"), list) else [] + files = data.get("files") if isinstance(data.get("files"), list) else [] + + folders.sort(key=_folder_sort_key) + files = _filter_browser_files_by_dimensions( + files, + min_w=min_w, + min_h=min_h, + max_w=max_w, + max_h=max_h, + ) + files.sort(key=lambda x: int(x.get("mtime") or 0), reverse=True) + hybrid = folders + files + + payload = _build_browser_listing_payload(hybrid, limit, offset, q, target=target) + payload["assets"] = _paginate_assets(hybrid, limit, offset) + return Result.Ok(payload) + + def list_browser_roots() -> Result[list[dict]]: """ List top-level filesystem roots for browser-mode custom scope. @@ -181,40 +523,35 @@ def list_filesystem_browser_entries( Browser-mode listing for custom scope (no pre-selected root). Returns folders + media files under absolute `path_value`, or roots when empty. """ - folders: list[dict] = [] - q = str(query or "*").strip() - ql = q.lower() - browse_all = q in ("", "*") - filter_kind = str(kind_filter or "").strip().lower() - min_size = max(0, int(min_size_bytes or 0)) - max_size = max(0, int(max_size_bytes or 0)) - if max_size > 0 and min_size > 0 and max_size < min_size: - max_size = min_size - min_w = max(0, int(min_width or 0)) - min_h = max(0, int(min_height or 0)) - max_w = max(0, int(max_width or 0)) - max_h = max(0, int(max_height or 0)) + opts = _normalize_browser_query_filters( + query, + kind_filter, + min_size_bytes, + max_size_bytes, + min_width, + min_height, + max_width, + max_height, + ) + q = str(opts["q"]) + ql = str(opts["ql"]) + browse_all = bool(opts["browse_all"]) + filter_kind = str(opts["filter_kind"]) + min_size = int(opts["min_size"]) + max_size = int(opts["max_size"]) + min_w = int(opts["min_w"]) + min_h = int(opts["min_h"]) + max_w = int(opts["max_w"]) + max_h = int(opts["max_h"]) + if not str(path_value or "").strip(): - roots_res = list_browser_roots() - if not roots_res.ok: - return Result.Err(roots_res.code, roots_res.error or "Failed to list roots") - folders = roots_res.data or [] - if filter_kind: - folders = [] - if not browse_all: - folders = [f for f in folders if ql in str(f.get("filename") or "").lower()] - total = len(folders) - start = max(0, int(offset or 0)) - end = start + int(limit or 0) if int(limit or 0) > 0 else None - return Result.Ok( - { - "assets": folders[start:end] if end is not None else folders[start:], - "total": total, - "limit": limit, - "offset": offset, - "query": q, - "scope": "custom", - } + return _list_browser_root_mode_entries( + q=q, + ql=ql, + browse_all=browse_all, + filter_kind=filter_kind, + limit=limit, + offset=offset, ) try: @@ -224,123 +561,18 @@ def list_filesystem_browser_entries( if not target.is_dir(): return Result.Err("INVALID_INPUT", "Path is not a directory") - files: list[dict] = [] - try: - for entry in target.iterdir(): - try: - name = entry.name - if name.startswith("."): - continue - if not browse_all and ql not in name.lower(): - continue - if entry.is_dir(): - if _is_hidden_directory(entry): - continue - resolved_child = entry.resolve(strict=True) - folders.append( - { - "id": f"folder::fs::{resolved_child}", - "filename": name, - "subfolder": str(resolved_child), - "filepath": str(resolved_child), - "kind": "folder", - "type": "custom", - "root_id": "", - "mtime": int(entry.stat().st_mtime), - "size": None, - "ext": "", - "rating": 0, - "tags": [], - "has_workflow": None, - "has_generation_data": None, - "width": None, - "height": None, - "duration": None, - } - ) - continue - if not entry.is_file(): - continue - kind = classify_file(name) - if kind == "unknown": - continue - if filter_kind and kind != filter_kind: - continue - st = entry.stat() - size = int(st.st_size or 0) - if min_size > 0 and size < min_size: - continue - if max_size > 0 and size > max_size: - continue - files.append( - { - "id": None, - "filename": name, - "subfolder": str(target), - "filepath": str(entry), - "kind": kind, - "ext": entry.suffix.lower(), - "size": size, - "mtime": int(st.st_mtime), - "width": None, - "height": None, - "duration": None, - "rating": 0, - "tags": [], - "has_workflow": None, - "has_generation_data": None, - "type": "custom", - "root_id": "", - } - ) - except Exception: - continue - except OSError as exc: - return Result.Err("LIST_FAILED", f"Failed to list directory: {exc}") - - folders.sort(key=_folder_sort_key) - if min_w > 0 or min_h > 0 or max_w > 0 or max_h > 0: - # Browser-mode filesystem rows usually don't have probed dimensions yet. - # Apply resolution filter only when dimensions are known. - filtered_files: list[dict] = [] - for f in files: - w_raw = f.get("width") - h_raw = f.get("height") - if w_raw is None or h_raw is None: - filtered_files.append(f) - continue - try: - w = int(w_raw or 0) - h = int(h_raw or 0) - if min_w > 0 and w < min_w: - continue - if min_h > 0 and h < min_h: - continue - if max_w > 0 and w > max_w: - continue - if max_h > 0 and h > max_h: - continue - except Exception: - continue - filtered_files.append(f) - files = filtered_files - - files.sort(key=lambda x: int(x.get("mtime") or 0), reverse=True) - hybrid = folders + files - total = len(hybrid) - start = max(0, int(offset or 0)) - end = start + int(limit or 0) if int(limit or 0) > 0 else None - page = hybrid[start:end] if end is not None else hybrid[start:] - return Result.Ok( - { - "assets": page, - "total": total, - "limit": limit, - "offset": offset, - "query": q, - "scope": "custom", - "path": str(target), - "current_path": str(target), - "display_name": _safe_name(target), - } + return _list_browser_path_mode_entries( + target, + q=q, + ql=ql, + browse_all=browse_all, + filter_kind=filter_kind, + min_size=min_size, + max_size=max_size, + min_w=min_w, + min_h=min_h, + max_w=max_w, + max_h=max_h, + limit=limit, + offset=offset, ) diff --git a/mjr_am_backend/features/collections/service.py b/mjr_am_backend/features/collections/service.py index 927a77b3..ee319898 100644 --- a/mjr_am_backend/features/collections/service.py +++ b/mjr_am_backend/features/collections/service.py @@ -109,6 +109,22 @@ class CollectionSummary: updated_at: str +def _read_collection_summary_row(path: Path) -> Optional[CollectionSummary]: + try: + data = json.loads(path.read_text(encoding="utf-8")) + except Exception: + return None + if not isinstance(data, dict): + return None + cid = str(data.get("id") or path.stem) + if not _safe_id(cid): + return None + name = str(data.get("name") or "").strip() or cid + count = int(len(data.get("items") or [])) + updated_at = str(data.get("updated_at") or "") + return CollectionSummary(cid, name, count, updated_at) + + class CollectionsService: """Filesystem-backed collections service (JSON files under `COLLECTIONS_DIR`).""" @@ -122,26 +138,18 @@ class CollectionsService: def list(self) -> Result[List[Dict[str, Any]]]: """List collections with basic metadata and item counts.""" - try: - base = self._base.resolve(strict=False) - except Exception as exc: - return Result.Err(ErrorCode.DB_ERROR, f"Collections dir unavailable: {exc}") + base_res = self._resolve_collections_base_dir() + if not base_res.ok: + return base_res + base = base_res.data items: List[CollectionSummary] = [] with self._lock: try: for p in base.glob("*.json"): - try: - data = json.loads(p.read_text(encoding="utf-8")) - except Exception: - continue - cid = str(data.get("id") or p.stem) - name = str(data.get("name") or "").strip() or cid - count = int(len(data.get("items") or [])) - updated_at = str(data.get("updated_at") or "") - if not _safe_id(cid): - continue - items.append(CollectionSummary(cid, name, count, updated_at)) + summary = _read_collection_summary_row(p) + if summary is not None: + items.append(summary) except Exception as exc: return Result.Err(ErrorCode.DB_ERROR, f"Failed to list collections: {exc}") @@ -149,6 +157,12 @@ class CollectionsService: items.sort(key=lambda x: (x.updated_at or "", x.id), reverse=True) return Result.Ok([i.__dict__ for i in items]) + def _resolve_collections_base_dir(self) -> Result[Path]: + try: + return Result.Ok(self._base.resolve(strict=False)) + except Exception as exc: + return Result.Err(ErrorCode.DB_ERROR, f"Collections dir unavailable: {exc}") + def create(self, name: str) -> Result[Dict[str, Any]]: """Create a new empty collection.""" cname = _safe_name(name) @@ -182,26 +196,32 @@ class CollectionsService: path = _collection_path(collection_id) if not path: return Result.Err(ErrorCode.INVALID_INPUT, "Invalid collection id") + loaded = self._load_collection_data(path) + if not loaded.ok: + return loaded + return self._build_collection_get_result(loaded.data or {}, collection_id) + + @staticmethod + def _load_collection_data(path: Path) -> Result[Dict[str, Any]]: try: - data = json.loads(path.read_text(encoding="utf-8")) + return Result.Ok(json.loads(path.read_text(encoding="utf-8"))) except FileNotFoundError: return Result.Err(ErrorCode.NOT_FOUND, "Collection not found") except Exception as exc: return Result.Err(ErrorCode.PARSE_ERROR, f"Failed to read collection: {exc}") + @staticmethod + def _build_collection_get_result(data: Dict[str, Any], collection_id: str) -> Result[Dict[str, Any]]: cid = str(data.get("id") or collection_id) if not _safe_id(cid): return Result.Err(ErrorCode.PARSE_ERROR, "Collection file is corrupted (bad id)") - name = str(data.get("name") or "").strip() or cid - items = data.get("items") if isinstance(data.get("items"), list) else [] - return Result.Ok( { "id": cid, - "name": name, + "name": str(data.get("name") or "").strip() or cid, "created_at": data.get("created_at") or "", "updated_at": data.get("updated_at") or "", - "items": items, + "items": data.get("items") if isinstance(data.get("items"), list) else [], } ) @@ -241,29 +261,113 @@ class CollectionsService: data["items"] = [] return path, data, None + @staticmethod + def _normalize_add_asset_item(asset: Dict[str, Any]) -> Optional[Dict[str, Any]]: + if not isinstance(asset, dict): + return None + fp = str(asset.get("filepath") or "").strip() + if not fp: + return None + root_id = CollectionsService._normalize_asset_root_id(asset) + kind = CollectionsService._normalize_asset_kind(asset, fp) + return { + "filepath": fp, + "filename": str(asset.get("filename") or "").strip(), + "subfolder": str(asset.get("subfolder") or "").strip(), + "type": str(asset.get("type") or "output").strip().lower(), + "root_id": root_id, + "kind": kind, + "added_at": _now_iso(), + } + + @staticmethod + def _normalize_asset_root_id(asset: Dict[str, Any]) -> Optional[str]: + root_id = str( + asset.get("root_id") + or asset.get("rootId") + or asset.get("custom_root_id") + or "" + ).strip() + return root_id or None + + @staticmethod + def _normalize_asset_kind(asset: Dict[str, Any], filepath: str) -> str: + return str(asset.get("kind") or classify_file(filepath) or "unknown").strip().lower() + + def _clean_add_assets_input(self, assets: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + cleaned: List[Dict[str, Any]] = [] + for asset in assets: + item = self._normalize_add_asset_item(asset) + if item is None: + continue + cleaned.append(item) + return cleaned + + @staticmethod + def _existing_items_with_indexes(raw_items: Any) -> Tuple[List[Dict[str, Any]], set[str], set[str]]: + existing: List[Dict[str, Any]] = ( + [it for it in raw_items if isinstance(it, dict)] if isinstance(raw_items, list) else [] + ) + seen: set[str] = set() + existing_seen: set[str] = set() + for item in existing: + existing_fp = item.get("filepath") + if not existing_fp: + continue + key = _normalize_fp(str(existing_fp)) + seen.add(key) + existing_seen.add(key) + return existing, seen, existing_seen + + def _append_collection_items( + self, + existing: List[Dict[str, Any]], + cleaned: List[Dict[str, Any]], + seen: set[str], + existing_seen: set[str], + ) -> Dict[str, int]: + added = 0 + skipped_existing = 0 + skipped_duplicate = 0 + skipped_limit = 0 + for item in cleaned: + if len(existing) >= _MAX_COLLECTION_ITEMS: + skipped_limit += 1 + continue + key = _normalize_fp(item["filepath"]) + if key in seen: + if key in existing_seen: + skipped_existing += 1 + else: + skipped_duplicate += 1 + continue + existing.append(item) + seen.add(key) + added += 1 + return { + "added": int(added), + "skipped_existing": int(skipped_existing), + "skipped_duplicate": int(skipped_duplicate), + "skipped_limit": int(skipped_limit), + } + + @staticmethod + def _write_collection_payload(path: Path, data: Dict[str, Any]) -> Optional[Result[Any]]: + try: + path.write_text( + json.dumps(data, ensure_ascii=False, separators=(",", ":")), + encoding="utf-8", + ) + except Exception as exc: + return Result.Err(ErrorCode.DB_ERROR, f"Failed to update collection: {exc}") + return None + def add_assets(self, collection_id: str, assets: List[Dict[str, Any]]) -> Result[Dict[str, Any]]: """Add assets (by filepath) to a collection (deduplicated, bounded).""" if not isinstance(assets, list) or not assets: return Result.Err(ErrorCode.INVALID_INPUT, "No assets provided") - cleaned: List[Dict[str, Any]] = [] - for a in assets: - if not isinstance(a, dict): - continue - fp = str(a.get("filepath") or "").strip() - if not fp: - continue - item = { - "filepath": fp, - "filename": str(a.get("filename") or "").strip(), - "subfolder": str(a.get("subfolder") or "").strip(), - "type": str(a.get("type") or "output").strip().lower(), - "root_id": str(a.get("root_id") or a.get("rootId") or a.get("custom_root_id") or "").strip() or None, - "kind": str(a.get("kind") or classify_file(fp) or "unknown").strip().lower(), - "added_at": _now_iso(), - } - cleaned.append(item) - + cleaned = self._clean_add_assets_input(assets) if not cleaned: return Result.Err(ErrorCode.INVALID_INPUT, "No valid assets provided") @@ -273,55 +377,23 @@ class CollectionsService: return err assert path is not None and data is not None - raw_items = data.get("items") - existing: List[Dict[str, Any]] = ( - [it for it in raw_items if isinstance(it, dict)] if isinstance(raw_items, list) else [] - ) - seen = set() - existing_seen = set() - for it in existing: - if isinstance(it, dict): - existing_fp = it.get("filepath") - if existing_fp: - k = _normalize_fp(str(existing_fp)) - seen.add(k) - existing_seen.add(k) - - added = 0 - skipped_existing = 0 - skipped_duplicate = 0 - skipped_limit = 0 - for it in cleaned: - if len(existing) >= _MAX_COLLECTION_ITEMS: - skipped_limit += 1 - continue - key = _normalize_fp(it["filepath"]) - if key in seen: - # Already in collection (or duplicated in the same request); report as skipped. - if key in existing_seen: - skipped_existing += 1 - else: - skipped_duplicate += 1 - continue - existing.append(it) - seen.add(key) - added += 1 + existing, seen, existing_seen = self._existing_items_with_indexes(data.get("items")) + counts = self._append_collection_items(existing, cleaned, seen, existing_seen) data["items"] = existing data["updated_at"] = _now_iso() - try: - path.write_text(json.dumps(data, ensure_ascii=False, separators=(",", ":")), encoding="utf-8") - except Exception as exc: - return Result.Err(ErrorCode.DB_ERROR, f"Failed to update collection: {exc}") + write_err = self._write_collection_payload(path, data) + if write_err: + return write_err return Result.Ok( { "id": str(data.get("id") or collection_id), - "added": int(added), - "skipped_existing": int(skipped_existing), - "skipped_duplicate": int(skipped_duplicate), - "skipped_limit": int(skipped_limit), + "added": counts["added"], + "skipped_existing": counts["skipped_existing"], + "skipped_duplicate": counts["skipped_duplicate"], + "skipped_limit": counts["skipped_limit"], "max_items": int(_MAX_COLLECTION_ITEMS), "count": len(existing), } @@ -329,11 +401,10 @@ class CollectionsService: def remove_filepaths(self, collection_id: str, filepaths: List[str]) -> Result[Dict[str, Any]]: """Remove items from a collection by filepath.""" - if not isinstance(filepaths, list) or not filepaths: - return Result.Err(ErrorCode.INVALID_INPUT, "No filepaths provided") - targets = set(_normalize_fp(str(p)) for p in filepaths if str(p or "").strip()) - if not targets: - return Result.Err(ErrorCode.INVALID_INPUT, "No valid filepaths provided") + targets_res = self._remove_targets(filepaths) + if not targets_res.ok: + return Result.Err(targets_res.code or ErrorCode.INVALID_INPUT, targets_res.error or "No valid filepaths provided") + targets = targets_res.data or set() with self._lock: path, data, err = self._load_for_update(collection_id) @@ -341,26 +412,37 @@ class CollectionsService: return err assert path is not None and data is not None - raw_items = data.get("items") - existing: List[Dict[str, Any]] = ( - [it for it in raw_items if isinstance(it, dict)] if isinstance(raw_items, list) else [] - ) - kept: List[Any] = [] - removed = 0 - for it in existing: - fp = it.get("filepath") - if fp and _normalize_fp(str(fp)) in targets: - removed += 1 - continue - kept.append(it) - + existing = self._collection_items_as_dicts(data.get("items")) + kept, removed = self._partition_removed_items(existing, targets) data["items"] = kept data["updated_at"] = _now_iso() - - try: - path.write_text(json.dumps(data, ensure_ascii=False, separators=(",", ":")), encoding="utf-8") - except Exception as exc: - return Result.Err(ErrorCode.DB_ERROR, f"Failed to update collection: {exc}") + write_err = self._write_collection_payload(path, data) + if write_err: + return write_err return Result.Ok({"id": str(data.get("id") or collection_id), "removed": int(removed), "count": len(kept)}) + @staticmethod + def _remove_targets(filepaths: List[str]) -> Result[set[str]]: + if not isinstance(filepaths, list) or not filepaths: + return Result.Err(ErrorCode.INVALID_INPUT, "No filepaths provided") + targets = set(_normalize_fp(str(p)) for p in filepaths if str(p or "").strip()) + if not targets: + return Result.Err(ErrorCode.INVALID_INPUT, "No valid filepaths provided") + return Result.Ok(targets) + + @staticmethod + def _collection_items_as_dicts(raw_items: Any) -> List[Dict[str, Any]]: + return [it for it in raw_items if isinstance(it, dict)] if isinstance(raw_items, list) else [] + + @staticmethod + def _partition_removed_items(existing: List[Dict[str, Any]], targets: set[str]) -> Tuple[List[Any], int]: + kept: List[Any] = [] + removed = 0 + for item in existing: + fp = item.get("filepath") + if fp and _normalize_fp(str(fp)) in targets: + removed += 1 + continue + kept.append(item) + return kept, removed diff --git a/mjr_am_backend/features/duplicates/service.py b/mjr_am_backend/features/duplicates/service.py index dbbdb578..6e8bae0e 100644 --- a/mjr_am_backend/features/duplicates/service.py +++ b/mjr_am_backend/features/duplicates/service.py @@ -59,6 +59,44 @@ def _hamming_hex(a: str, b: str) -> int: return 64 +def _build_known_hash_fields(row: Dict[str, Any]) -> tuple[str, str, str]: + return ( + str(row.get("hash_state") or ""), + str(row.get("content_hash") or ""), + str(row.get("phash") or ""), + ) + + +def _exact_group_item(row: Dict[str, Any]) -> Dict[str, Any]: + row_data = row or {} + return { + "id": _safe_int(row_data.get("id")), + "filepath": row_data.get("filepath"), + "filename": row_data.get("filename"), + "tags": row_data.get("tags") or "[]", + } + + +def _similar_pair_side(row: Dict[str, Any]) -> Dict[str, Any]: + return { + "id": _safe_int(row.get("id")), + "filepath": row.get("filepath"), + "filename": row.get("filename"), + } + + +def _similar_pair_row(distance: int, left: Dict[str, Any], right: Dict[str, Any]) -> Dict[str, Any]: + return { + "distance": distance, + "left": _similar_pair_side(left), + "right": _similar_pair_side(right), + } + + +def _phash_value(row: Dict[str, Any]) -> str: + return str((row or {}).get("phash") or "") + + class DuplicatesService: def __init__(self, db: Sqlite): self.db = db @@ -93,100 +131,122 @@ class DuplicatesService: async def _run_background(self, limit: int) -> None: try: - rows_res = await self.db.aquery( - """ - SELECT id, filepath, kind, size, mtime, content_hash, phash, hash_state - FROM assets - ORDER BY indexed_at DESC, id DESC - LIMIT ? - """, - (int(limit),), - ) + rows_res = await self._fetch_analysis_rows(limit) if not rows_res.ok: self._set_status(running=False, last_error=rows_res.error or "Query failed") return rows = rows_res.data or [] for row in rows: - aid = _safe_int((row or {}).get("id")) - fp = str((row or {}).get("filepath") or "") - kind = str((row or {}).get("kind") or "").lower() - size = _safe_int((row or {}).get("size")) - mtime = _safe_int((row or {}).get("mtime")) - current_state = f"{mtime}:{size}" - known_state = str((row or {}).get("hash_state") or "") - known_content = str((row or {}).get("content_hash") or "") - known_phash = str((row or {}).get("phash") or "") - path = Path(fp) - - self._status["processed"] = _safe_int(self._status.get("processed")) + 1 - if not aid or not fp or not path.exists() or not path.is_file(): - self._status["errors"] = _safe_int(self._status.get("errors")) + 1 + item = self._build_row_context(row) + self._inc_status("processed") + if not self._is_row_processable(item): + self._inc_status("errors") continue - if known_state == current_state and known_content and (kind != "image" or known_phash): + if self._row_is_up_to_date(item): continue - - try: - content_hash = await asyncio.to_thread(_compute_file_hash, path) - phash = None - if kind == "image": - phash = await asyncio.to_thread(_compute_phash_hex, path) - upd = await self.db.aexecute( - """ - UPDATE assets - SET content_hash = ?, phash = ?, hash_state = ? - WHERE id = ? - """, - (content_hash, phash, current_state, aid), - ) - if upd.ok: - self._status["updated"] = _safe_int(self._status.get("updated")) + 1 - else: - self._status["errors"] = _safe_int(self._status.get("errors")) + 1 - except Exception as exc: - self._status["errors"] = _safe_int(self._status.get("errors")) + 1 - self._status["last_error"] = str(exc) + await self._process_hash_update(item) except Exception as exc: self._set_status(last_error=str(exc)) finally: self._set_status(running=False) - async def get_alerts( - self, - roots: Optional[List[str]] = None, - max_groups: int = 6, - max_pairs: int = 10, - phash_distance: int = 6, - ) -> Result[Dict[str, Any]]: - roots = roots or [] - max_groups = max(1, min(50, int(max_groups or 6))) - max_pairs = max(1, min(100, int(max_pairs or 10))) - phash_distance = max(0, min(32, int(phash_distance or 6))) + async def _fetch_analysis_rows(self, limit: int): + return await self.db.aquery( + """ + SELECT id, filepath, kind, size, mtime, content_hash, phash, hash_state + FROM assets + ORDER BY indexed_at DESC, id DESC + LIMIT ? + """, + (int(limit),), + ) - where = "" - params: List[Any] = [] - if roots: - clauses = [] - for r in roots: - clauses.append("a.filepath LIKE ?") - params.append(str(Path(str(r)).resolve(strict=False)) + "%") - where = f"WHERE ({' OR '.join(clauses)})" + def _build_row_context(self, row: Dict[str, Any]) -> Dict[str, Any]: + row_data = row or {} + aid = _safe_int(row_data.get("id")) + fp = str(row_data.get("filepath") or "") + kind = str(row_data.get("kind") or "").lower() + size = _safe_int(row_data.get("size")) + mtime = _safe_int(row_data.get("mtime")) + known_state, known_content, known_phash = _build_known_hash_fields(row_data) + return { + "aid": aid, + "fp": fp, + "kind": kind, + "path": Path(fp), + "current_state": f"{mtime}:{size}", + "known_state": known_state, + "known_content": known_content, + "known_phash": known_phash, + } - dup_q = f""" - SELECT a.id, a.filepath, a.filename, a.content_hash, m.tags - FROM assets a - LEFT JOIN asset_metadata m ON m.asset_id = a.id - {where} - AND a.content_hash IS NOT NULL - AND a.content_hash IN ( - SELECT content_hash - FROM assets - WHERE content_hash IS NOT NULL - GROUP BY content_hash - HAVING COUNT(1) > 1 + def _is_row_processable(self, item: Dict[str, Any]) -> bool: + path = item["path"] + return bool(item["aid"] and item["fp"] and path.exists() and path.is_file()) + + def _row_is_up_to_date(self, item: Dict[str, Any]) -> bool: + if item["known_state"] != item["current_state"]: + return False + if not item["known_content"]: + return False + if item["kind"] == "image" and not item["known_phash"]: + return False + return True + + async def _process_hash_update(self, item: Dict[str, Any]) -> None: + try: + content_hash = await asyncio.to_thread(_compute_file_hash, item["path"]) + phash = await asyncio.to_thread(_compute_phash_hex, item["path"]) if item["kind"] == "image" else None + upd = await self.db.aexecute( + """ + UPDATE assets + SET content_hash = ?, phash = ?, hash_state = ? + WHERE id = ? + """, + (content_hash, phash, item["current_state"], item["aid"]), ) - ORDER BY a.content_hash, a.id - """ if where else """ + if upd.ok: + self._inc_status("updated") + else: + self._inc_status("errors") + except Exception as exc: + self._inc_status("errors") + self._status["last_error"] = str(exc) + + def _inc_status(self, key: str, amount: int = 1) -> None: + self._status[key] = _safe_int(self._status.get(key)) + int(amount) + + @staticmethod + def _alerts_where_clause(roots: List[str]) -> tuple[str, List[Any]]: + params: List[Any] = [] + if not roots: + return "", params + clauses: List[str] = [] + for root in roots: + clauses.append("a.filepath LIKE ?") + params.append(str(Path(str(root)).resolve(strict=False)) + "%") + return f"WHERE ({' OR '.join(clauses)})", params + + @staticmethod + def _exact_duplicates_query(where: str) -> str: + if where: + return f""" + SELECT a.id, a.filepath, a.filename, a.content_hash, m.tags + FROM assets a + LEFT JOIN asset_metadata m ON m.asset_id = a.id + {where} + AND a.content_hash IS NOT NULL + AND a.content_hash IN ( + SELECT content_hash + FROM assets + WHERE content_hash IS NOT NULL + GROUP BY content_hash + HAVING COUNT(1) > 1 + ) + ORDER BY a.content_hash, a.id + """ + return """ SELECT a.id, a.filepath, a.filename, a.content_hash, m.tags FROM assets a LEFT JOIN asset_metadata m ON m.asset_id = a.id @@ -200,39 +260,20 @@ class DuplicatesService: ) ORDER BY a.content_hash, a.id """ - dup_rows = await self.db.aquery(dup_q, tuple(params)) - if not dup_rows.ok: - return Result.Err("DB_ERROR", dup_rows.error or "Duplicate query failed") - groups: Dict[str, List[Dict[str, Any]]] = {} - for row in dup_rows.data or []: - key = str((row or {}).get("content_hash") or "") - if not key: - continue - groups.setdefault(key, []).append({ - "id": _safe_int((row or {}).get("id")), - "filepath": (row or {}).get("filepath"), - "filename": (row or {}).get("filename"), - "tags": (row or {}).get("tags") or "[]", - }) - - exact_groups = [] - for h, items in groups.items(): - if len(items) < 2: - continue - exact_groups.append({"hash": h, "count": len(items), "assets": items[:25]}) - exact_groups.sort(key=lambda g: _safe_int(g.get("count")), reverse=True) - exact_groups = exact_groups[:max_groups] - - sim_q = f""" - SELECT a.id, a.filepath, a.filename, a.phash - FROM assets a - {where} - AND a.kind = 'image' - AND a.phash IS NOT NULL - ORDER BY a.id DESC - LIMIT 800 - """ if where else """ + @staticmethod + def _similarity_query(where: str) -> str: + if where: + return f""" + SELECT a.id, a.filepath, a.filename, a.phash + FROM assets a + {where} + AND a.kind = 'image' + AND a.phash IS NOT NULL + ORDER BY a.id DESC + LIMIT 800 + """ + return """ SELECT a.id, a.filepath, a.filename, a.phash FROM assets a WHERE a.kind = 'image' @@ -240,45 +281,135 @@ class DuplicatesService: ORDER BY a.id DESC LIMIT 800 """ - sim_rows = await self.db.aquery(sim_q, tuple(params)) - if not sim_rows.ok: - return Result.Err("DB_ERROR", sim_rows.error or "Similarity query failed") - images = sim_rows.data or [] - similar_pairs = [] - for i in range(len(images)): - left = images[i] or {} - lp = str(left.get("phash") or "") + @staticmethod + def _build_exact_groups(rows: List[Dict[str, Any]], max_groups: int) -> List[Dict[str, Any]]: + groups: Dict[str, List[Dict[str, Any]]] = {} + for row in rows or []: + row_data = row or {} + key = str(row_data.get("content_hash") or "") + if not key: + continue + groups.setdefault(key, []).append(_exact_group_item(row_data)) + + exact_groups: List[Dict[str, Any]] = [] + for key, items in groups.items(): + if len(items) < 2: + continue + exact_groups.append({"hash": key, "count": len(items), "assets": items[:25]}) + exact_groups.sort(key=lambda g: _safe_int(g.get("count")), reverse=True) + return exact_groups[:max_groups] + + @staticmethod + def _build_similar_pairs( + images: List[Dict[str, Any]], + *, + phash_distance: int, + max_pairs: int, + ) -> List[Dict[str, Any]]: + similar_pairs: List[Dict[str, Any]] = [] + for i, left_raw in enumerate(images): + left = left_raw or {} + lp = _phash_value(left) if not lp: continue - for j in range(i + 1, len(images)): - right = images[j] or {} - rp = str(right.get("phash") or "") - if not rp: - continue - d = _hamming_hex(lp, rp) - if d <= phash_distance: - similar_pairs.append({ - "distance": d, - "left": {"id": _safe_int(left.get("id")), "filepath": left.get("filepath"), "filename": left.get("filename")}, - "right": {"id": _safe_int(right.get("id")), "filepath": right.get("filepath"), "filename": right.get("filename")}, - }) - if len(similar_pairs) >= max_pairs: - break - if len(similar_pairs) >= max_pairs: + reached_limit = DuplicatesService._append_similar_pairs_for_left( + similar_pairs, + left, + lp, + images[i + 1:], + phash_distance=phash_distance, + max_pairs=max_pairs, + ) + if reached_limit: break - similar_pairs.sort(key=lambda p: _safe_int(p.get("distance"), 64)) + return similar_pairs + @staticmethod + def _append_similar_pairs_for_left( + similar_pairs: List[Dict[str, Any]], + left: Dict[str, Any], + left_phash: str, + right_rows: List[Dict[str, Any]], + *, + phash_distance: int, + max_pairs: int, + ) -> bool: + for right_raw in right_rows: + right = right_raw or {} + right_phash = _phash_value(right) + if not right_phash: + continue + distance = _hamming_hex(left_phash, right_phash) + if distance > phash_distance: + continue + similar_pairs.append(_similar_pair_row(distance, left, right)) + if len(similar_pairs) >= max_pairs: + return True + return False + + async def get_alerts( + self, + roots: Optional[List[str]] = None, + max_groups: int = 6, + max_pairs: int = 10, + phash_distance: int = 6, + ) -> Result[Dict[str, Any]]: + roots, max_groups, max_pairs, phash_distance = self._normalize_alert_args( + roots=roots, + max_groups=max_groups, + max_pairs=max_pairs, + phash_distance=phash_distance, + ) + where, params = self._alerts_where_clause(roots) + exact_groups_res = await self._query_exact_groups(where, params, max_groups=max_groups) + if not exact_groups_res.ok: + return exact_groups_res + sim_rows_res = await self._query_similarity_rows(where, params) + if not sim_rows_res.ok: + return sim_rows_res + similar_pairs = self._build_similar_pairs( + sim_rows_res.data or [], + phash_distance=phash_distance, + max_pairs=max_pairs, + ) return Result.Ok({ - "exact_groups": exact_groups, + "exact_groups": exact_groups_res.data or [], "similar_pairs": similar_pairs, "status": dict(self._status), }) + @staticmethod + def _normalize_alert_args( + *, + roots: Optional[List[str]], + max_groups: int, + max_pairs: int, + phash_distance: int, + ) -> tuple[List[str], int, int, int]: + return ( + roots or [], + max(1, min(50, int(max_groups or 6))), + max(1, min(100, int(max_pairs or 10))), + max(0, min(32, int(phash_distance or 6))), + ) + + async def _query_exact_groups(self, where: str, params: List[Any], *, max_groups: int) -> Result[List[Dict[str, Any]]]: + dup_rows = await self.db.aquery(self._exact_duplicates_query(where), tuple(params)) + if not dup_rows.ok: + return Result.Err("DB_ERROR", dup_rows.error or "Duplicate query failed") + return Result.Ok(self._build_exact_groups(dup_rows.data or [], max_groups)) + + async def _query_similarity_rows(self, where: str, params: List[Any]) -> Result[List[Dict[str, Any]]]: + sim_rows = await self.db.aquery(self._similarity_query(where), tuple(params)) + if not sim_rows.ok: + return Result.Err("DB_ERROR", sim_rows.error or "Similarity query failed") + return Result.Ok(sim_rows.data or []) + async def merge_tags_for_group(self, keep_asset_id: int, merge_asset_ids: List[int]) -> Result[Dict[str, Any]]: keep_asset_id = int(keep_asset_id or 0) - merge_ids = [int(x) for x in (merge_asset_ids or []) if int(x or 0) > 0 and int(x or 0) != keep_asset_id] + merge_ids = self._normalize_merge_ids(keep_asset_id, merge_asset_ids) if not keep_asset_id or not merge_ids: return Result.Err("INVALID_INPUT", "Invalid keep/merge asset ids") @@ -291,31 +422,7 @@ class DuplicatesService: if not rows.ok: return Result.Err("DB_ERROR", rows.error or "Failed to load tags") - merged: List[str] = [] - seen = set() - for row in rows.data or []: - raw = (row or {}).get("tags") - tags = [] - if isinstance(raw, str): - try: - parsed = json.loads(raw) - if isinstance(parsed, list): - tags = parsed - except Exception: - tags = [] - elif isinstance(raw, list): - tags = raw - for t in tags: - if not isinstance(t, str): - continue - v = t.strip() - if not v: - continue - k = v.lower() - if k in seen: - continue - seen.add(k) - merged.append(v) + merged = self._collect_merged_tags(rows.data or []) upd = await self.db.aexecute( "UPDATE asset_metadata SET tags = ?, tags_text = ? WHERE asset_id = ?", @@ -324,3 +431,40 @@ class DuplicatesService: if not upd.ok: return Result.Err("DB_ERROR", upd.error or "Failed to update tags") return Result.Ok({"asset_id": keep_asset_id, "tags": merged, "merged_from": merge_ids}) + + def _normalize_merge_ids(self, keep_asset_id: int, merge_asset_ids: List[int]) -> List[int]: + out: List[int] = [] + for value in merge_asset_ids or []: + candidate = int(value or 0) + if candidate > 0 and candidate != keep_asset_id: + out.append(candidate) + return out + + def _collect_merged_tags(self, rows: List[Dict[str, Any]]) -> List[str]: + merged: List[str] = [] + seen: set[str] = set() + for row in rows: + for tag in self._extract_tags_from_row(row): + normalized = tag.strip() + if not normalized: + continue + key = normalized.lower() + if key in seen: + continue + seen.add(key) + merged.append(normalized) + return merged + + def _extract_tags_from_row(self, row: Dict[str, Any]) -> List[str]: + raw = (row or {}).get("tags") + if isinstance(raw, list): + return [value for value in raw if isinstance(value, str)] + if not isinstance(raw, str): + return [] + try: + parsed = json.loads(raw) + if isinstance(parsed, list): + return [value for value in parsed if isinstance(value, str)] + except Exception: + return [] + return [] diff --git a/mjr_am_backend/features/geninfo/parser.py b/mjr_am_backend/features/geninfo/parser.py index 540024fe..5cc23087 100644 --- a/mjr_am_backend/features/geninfo/parser.py +++ b/mjr_am_backend/features/geninfo/parser.py @@ -10,7 +10,7 @@ from __future__ import annotations import os from collections import deque from dataclasses import dataclass -from typing import Any, Dict, List, Optional, Set, Tuple, cast +from typing import Any, Callable, Dict, List, Optional, Set, Tuple, cast from ...shared import Result, get_logger @@ -36,6 +36,53 @@ SINK_CLASS_TYPES: Set[str] = { } +_VIDEO_SINK_TYPES: Set[str] = {"savevideo", "vhs_savevideo", "vhs_videocombine"} +_AUDIO_SINK_TYPES: Set[str] = {"saveaudio", "save_audio", "vhs_saveaudio"} +_IMAGE_SINK_TYPES: Set[str] = {"saveimage", "saveimagewebsocket", "saveanimatedwebp", "savegif"} + + +def _sink_group(ct: str) -> int: + if _sink_is_video(ct): + return 0 + if _sink_is_audio(ct): + return 1 + if _sink_is_image(ct): + return 2 + if _sink_is_preview(ct): + return 3 + return 4 + + +def _sink_is_video(ct: str) -> bool: + return ct in _VIDEO_SINK_TYPES or (("save" in ct) and ("video" in ct)) + + +def _sink_is_audio(ct: str) -> bool: + return ct in _AUDIO_SINK_TYPES or (("save" in ct) and ("audio" in ct)) + + +def _sink_is_image(ct: str) -> bool: + return ct in _IMAGE_SINK_TYPES or (("save" in ct) and ("image" in ct)) + + +def _sink_is_preview(ct: str) -> bool: + return ct == "previewimage" or "preview" in ct + + +def _sink_images_tiebreak(node: Dict[str, Any]) -> int: + try: + return 0 if _is_link(_inputs(node).get("images")) else 1 + except Exception: + return 1 + + +def _sink_node_id_tiebreak(node_id: Optional[str]) -> int: + try: + return -int(node_id) if node_id is not None else 0 + except Exception: + return 0 + + def _sink_priority(node: Dict[str, Any], node_id: Optional[str] = None) -> Tuple[int, int, int]: """ Rank sinks so we pick the "real" output when multiple sinks exist. @@ -44,36 +91,7 @@ def _sink_priority(node: Dict[str, Any], node_id: Optional[str] = None) -> Tuple attached to intermediate nodes and does not reliably reflect the final render. """ ct = _lower(_node_type(node)) - - # Group 0: Video Sinks (Explicit & Heuristic) - if ct in ("savevideo", "vhs_savevideo", "vhs_videocombine") or (("save" in ct) and ("video" in ct)): - group = 0 - # Group 1: Audio Sinks - elif ct in ("saveaudio", "save_audio", "vhs_saveaudio") or (("save" in ct) and ("audio" in ct)): - group = 1 - # Group 2: Image/Gif Sinks (Explicit & Heuristic) - elif ct in ("saveimage", "saveimagewebsocket", "saveanimatedwebp", "savegif") or (("save" in ct) and ("image" in ct)): - group = 2 - # Group 3: Preview - elif ct == "previewimage" or "preview" in ct: - group = 3 - # Group 4: Fallback - else: - group = 4 - - # Within group: prefer sinks that consume `images` (common for image outputs) - try: - has_images = 0 if _is_link(_inputs(node).get("images")) else 1 - except Exception: - has_images = 1 - - # Tie-breaker: prefer higher node IDs (likely added later / final output) - try: - nid_score = -int(node_id) if node_id is not None else 0 - except Exception: - nid_score = 0 - - prio = (group, has_images, nid_score) + prio = (_sink_group(ct), _sink_images_tiebreak(node), _sink_node_id_tiebreak(node_id)) # print(f"DEBUG: Sink {node_id} priority {prio}") return prio @@ -160,6 +178,15 @@ def _inputs(node: Any) -> Dict[str, Any]: def _lower(s: Any) -> str: return str(s or "").lower() + +def _is_numeric_like_text(value: str) -> bool: + return all(ch.isdigit() or ch.isspace() or ch in ".,+-" for ch in value) + + +def _has_control_chars(value: str) -> bool: + return any(ord(ch) < 9 for ch in value) + + def _looks_like_prompt_string(value: Any) -> bool: """ Conservative heuristic: accept human-ish prompt strings; reject numbers/gibberish. @@ -167,21 +194,13 @@ def _looks_like_prompt_string(value: Any) -> bool: if not isinstance(value, str): return False s = value.strip() - if not s: + if not s or len(s) < 6: return False - if len(s) < 6: + if _is_numeric_like_text(s): return False - # Reject numeric-only payloads (common for some nodes that also have `text`). - if all(ch.isdigit() or ch.isspace() or ch in ".,+-" for ch in s): + if _has_control_chars(s): return False - # Reject control-heavy strings (binary-ish). - bad = sum(1 for ch in s if ord(ch) < 9) - if bad > 0: - return False - # Require at least one letter (unicode-aware). - if not any(ch.isalpha() for ch in s): - return False - return True + return any(ch.isalpha() for ch in s) def _is_reroute(node: Dict[str, Any]) -> bool: @@ -271,13 +290,13 @@ def _find_candidate_sinks(nodes_by_id: Dict[str, Dict[str, Any]]) -> List[str]: if ct in SINK_CLASS_TYPES: sinks.append(node_id) continue - + # Heuristic for custom save nodes (e.g. WAS, Impact, CR, etc) # Must contain "save" or "preview" AND media hint if ("save" in ct or "preview" in ct) and ("image" in ct or "video" in ct or "audio" in ct): sinks.append(node_id) continue - + return sinks @@ -335,52 +354,48 @@ def _is_sampler(node: Dict[str, Any]) -> bool: ct = _lower(_node_type(node)) if not ct: return False - # KSamplerSelect is a sampler selector node, not a diffusion sampler. + if _is_named_sampler_type(ct): + return True + if _has_core_sampler_signature(node): + return True + return _is_custom_sampler(node, ct) + + +def _is_named_sampler_type(ct: str) -> bool: if "ksampler" in ct and "select" in ct: return False if "ksampler" in ct: return True - if "iterativelatentupscale" in ct: + if "iterativelatentupscale" in ct or "marigold" in ct: return True - # Marigold depth estimation (acts as sampler for depth maps) - if "marigold" in ct: - return True - # Kijai Flux inference sampler if "flux" in ct and ("sampler" in ct or "params" in ct): return True - - # Generic "Flux" nodes that act as samplers (e.g. "Flux2") - if ct == "flux2" or "flux_2" in ct: - return True + return ct == "flux2" or "flux_2" in ct - # Generic detection: any node with (steps + cfg + seed) is likely a sampler + +def _has_core_sampler_signature(node: Dict[str, Any]) -> bool: try: ins = _inputs(node) - has_steps = ins.get("steps") is not None - has_cfg = ins.get("cfg") is not None or ins.get("cfg_scale") is not None or ins.get("guidance") is not None - has_seed = ins.get("seed") is not None or ins.get("noise_seed") is not None - if has_steps and has_cfg and has_seed: - return True except Exception: - pass + return False + has_steps = ins.get("steps") is not None + has_cfg = any(ins.get(key) is not None for key in ("cfg", "cfg_scale", "guidance")) + has_seed = any(ins.get(key) is not None for key in ("seed", "noise_seed")) + return has_steps and has_cfg and has_seed - # WanVideoSampler / custom samplers - if ("sampler" in ct) and ("select" not in ct) and ("ksamplerselect" not in ct): - try: - ins = _inputs(node) - # Avoid matching unrelated nodes that happen to contain "sampler" in their name. - if _is_link(ins.get("model")): - return True - # Require at least one real sampling parameter (not just "sampler_name" selector nodes). - for k in ("steps", "cfg", "cfg_scale", "seed", "scheduler", "denoise"): - if ins.get(k) is not None: - return True - # Wan samplers often take `text_embeds` rather than conditioning. - if _is_link(ins.get("text_embeds")) or _is_link(ins.get("hyvid_embeds")): - return True - except (TypeError, ValueError, KeyError): - return False - return False + +def _is_custom_sampler(node: Dict[str, Any], ct: str) -> bool: + if "sampler" not in ct or "select" in ct or "ksamplerselect" in ct: + return False + try: + ins = _inputs(node) + except (TypeError, ValueError, KeyError): + return False + if _is_link(ins.get("model")): + return True + if any(ins.get(k) is not None for k in ("steps", "cfg", "cfg_scale", "seed", "scheduler", "denoise")): + return True + return _is_link(ins.get("text_embeds")) or _is_link(ins.get("hyvid_embeds")) def _is_advanced_sampler(node: Dict[str, Any]) -> bool: @@ -447,55 +462,55 @@ def _extract_posneg_from_text_embeds( def _select_primary_sampler( nodes_by_id: Dict[str, Dict[str, Any]], sink_node_id: str ) -> Tuple[Optional[str], str]: - sink = nodes_by_id.get(sink_node_id) - if not isinstance(sink, dict): - return None, "none" - start_link = _pick_sink_inputs(sink) - start_src = _walk_passthrough(nodes_by_id, start_link) if start_link else None - if not start_src: - return None, "none" - dist = _collect_upstream_nodes(nodes_by_id, start_src) - samplers: List[Tuple[str, int]] = [] - for nid, d in dist.items(): - node = nodes_by_id.get(nid) - if not isinstance(node, dict): - continue - if _is_sampler(node): - samplers.append((nid, d)) - if not samplers: - return None, "none" - samplers.sort(key=lambda x: x[1]) - best_depth = samplers[0][1] - best = [nid for nid, d in samplers if d == best_depth] - if len(best) == 1: - return best[0], "high" - # Ambiguous: multiple samplers equally close to sink - return best[0], "medium" + return _select_sampler_from_sink(nodes_by_id, sink_node_id, _is_sampler) def _select_advanced_sampler( nodes_by_id: Dict[str, Dict[str, Any]], sink_node_id: str ) -> Tuple[Optional[str], str]: - sink = nodes_by_id.get(sink_node_id) - if not isinstance(sink, dict): - return None, "none" - start_link = _pick_sink_inputs(sink) - start_src = _walk_passthrough(nodes_by_id, start_link) if start_link else None + return _select_sampler_from_sink(nodes_by_id, sink_node_id, _is_advanced_sampler) + + +def _select_sampler_from_sink( + nodes_by_id: Dict[str, Dict[str, Any]], + sink_node_id: str, + selector: Callable[[Dict[str, Any]], bool], +) -> Tuple[Optional[str], str]: + start_src = _sink_start_source(nodes_by_id, sink_node_id) if not start_src: return None, "none" + candidates = _upstream_sampler_candidates(nodes_by_id, start_src, selector) + return _best_candidate(candidates) + + +def _sink_start_source(nodes_by_id: Dict[str, Dict[str, Any]], sink_node_id: str) -> Optional[str]: + sink = nodes_by_id.get(sink_node_id) + if not isinstance(sink, dict): + return None + start_link = _pick_sink_inputs(sink) + return _walk_passthrough(nodes_by_id, start_link) if start_link else None + + +def _upstream_sampler_candidates( + nodes_by_id: Dict[str, Dict[str, Any]], + start_src: str, + selector: Callable[[Dict[str, Any]], bool], +) -> List[Tuple[str, int]]: dist = _collect_upstream_nodes(nodes_by_id, start_src) candidates: List[Tuple[str, int]] = [] - for nid, d in dist.items(): + for nid, depth in dist.items(): node = nodes_by_id.get(nid) - if not isinstance(node, dict): - continue - if _is_advanced_sampler(node): - candidates.append((nid, d)) + if isinstance(node, dict) and selector(node): + candidates.append((nid, depth)) + return candidates + + +def _best_candidate(candidates: List[Tuple[str, int]]) -> Tuple[Optional[str], str]: if not candidates: return None, "none" - candidates.sort(key=lambda x: x[1]) + candidates.sort(key=lambda item: item[1]) best_depth = candidates[0][1] - best = [nid for nid, d in candidates if d == best_depth] + best = [nid for nid, depth in candidates if depth == best_depth] if len(best) == 1: return best[0], "high" return best[0], "medium" @@ -506,33 +521,42 @@ def _select_any_sampler(nodes_by_id: Dict[str, Dict[str, Any]]) -> Tuple[Optiona Last resort when sinks exist but are not linked to the generation branch. Choose the "best" sampler-like node in the whole prompt graph. """ - candidates: List[Tuple[int, int, str]] = [] - for nid, node in nodes_by_id.items(): - if not isinstance(node, dict): - continue - if not _is_sampler(node): - continue - ins = _inputs(node) - score = 0 - if _is_link(ins.get("model")): - score += 3 - if _is_link(ins.get("positive")) or _is_link(ins.get("text_embeds")): - score += 3 - for k in ("steps", "cfg", "cfg_scale", "seed", "denoise", "scheduler"): - if ins.get(k) is not None: - score += 1 - # Prefer numeric node ids (stable ordering). - try: - n_int = int(nid) - except Exception: - n_int = 10**9 - candidates.append((-score, n_int, nid)) + candidates = _global_sampler_candidates(nodes_by_id) if not candidates: return None, "none" candidates.sort() return candidates[0][2], "low" +def _global_sampler_candidates(nodes_by_id: Dict[str, Dict[str, Any]]) -> List[Tuple[int, int, str]]: + candidates: List[Tuple[int, int, str]] = [] + for nid, node in nodes_by_id.items(): + if not isinstance(node, dict) or not _is_sampler(node): + continue + score = _sampler_candidate_score(_inputs(node)) + candidates.append((-score, _stable_numeric_node_id(nid), nid)) + return candidates + + +def _sampler_candidate_score(ins: Dict[str, Any]) -> int: + score = 0 + if _is_link(ins.get("model")): + score += 3 + if _is_link(ins.get("positive")) or _is_link(ins.get("text_embeds")): + score += 3 + for key in ("steps", "cfg", "cfg_scale", "seed", "denoise", "scheduler"): + if ins.get(key) is not None: + score += 1 + return score + + +def _stable_numeric_node_id(node_id: str) -> int: + try: + return int(node_id) + except Exception: + return 10**9 + + def _trace_sampler_name(nodes_by_id: Dict[str, Dict[str, Any]], link: Any) -> Optional[Tuple[str, str]]: src_id = _walk_passthrough(nodes_by_id, link) if not src_id: @@ -579,26 +603,7 @@ def _trace_scheduler_sigmas( steps = _scalar(ins.get("steps")) steps_confidence: Optional[str] = "high" if steps is not None else None if steps is None: - # Some video workflows use manual sigma schedules (e.g. ManualSigmas) instead of an explicit `steps` field. - # Best-effort: count the comma-separated values and treat it as (steps + 1). - try: - sigmas = ins.get("sigmas") - if isinstance(sigmas, str) and sigmas.strip(): - raw_parts = [p.strip() for p in sigmas.replace("\n", " ").split(",")] - parts = [p for p in raw_parts if p] - # Ensure these are numeric-ish to avoid accidentally treating random strings as step schedules. - numeric = 0 - for p in parts: - try: - float(p) - numeric += 1 - except Exception: - pass - if numeric >= 2: - steps = max(1, numeric - 1) - steps_confidence = "low" - except Exception: - pass + steps, steps_confidence = _steps_from_manual_sigmas(ins) scheduler = _scalar(ins.get("scheduler")) denoise = _scalar(ins.get("denoise")) model_link = ins.get("model") if _is_link(ins.get("model")) else None @@ -606,6 +611,32 @@ def _trace_scheduler_sigmas( return (steps, scheduler, denoise, model_link, (src_id, src), steps_confidence) +def _steps_from_manual_sigmas(ins: Dict[str, Any]) -> Tuple[Optional[Any], Optional[str]]: + # Some video workflows use manual sigma schedules instead of explicit `steps`. + numeric = _count_numeric_sigma_values(ins.get("sigmas")) + if numeric >= 2: + return max(1, numeric - 1), "low" + return None, None + + +def _count_numeric_sigma_values(sigmas: Any) -> int: + try: + if not (isinstance(sigmas, str) and sigmas.strip()): + return 0 + raw_parts = [p.strip() for p in sigmas.replace("\n", " ").split(",")] + parts = [p for p in raw_parts if p] + numeric = 0 + for p in parts: + try: + float(p) + numeric += 1 + except Exception: + continue + return numeric + except Exception: + return 0 + + def _trace_guidance_from_conditioning(nodes_by_id: Dict[str, Dict[str, Any]], conditioning_link: Any) -> Optional[Tuple[Any, str]]: start_id = _walk_passthrough(nodes_by_id, conditioning_link) if not start_id: @@ -646,24 +677,32 @@ def _extract_ksampler_widget_params(node: Dict[str, Any]) -> Dict[str, Any]: widgets = node.get("widgets_values") if not isinstance(widgets, list): return out - try: - if len(widgets) > 0: - out["seed"] = widgets[0] - if len(widgets) > 2: - out["steps"] = widgets[2] - if len(widgets) > 3: - out["cfg"] = widgets[3] - if len(widgets) > 4: - out["sampler_name"] = widgets[4] - if len(widgets) > 5: - out["scheduler"] = widgets[5] - if len(widgets) > 6: - out["denoise"] = widgets[6] - except Exception: - return {} + return _ksampler_values_from_widgets(widgets) + + +def _ksampler_values_from_widgets(widgets: List[Any]) -> Dict[str, Any]: + index_map = { + "seed": 0, + "steps": 2, + "cfg": 3, + "sampler_name": 4, + "scheduler": 5, + "denoise": 6, + } + out: Dict[str, Any] = {} + for field, idx in index_map.items(): + value = _widget_value_at(widgets, idx) + if value is not None: + out[field] = value return out +def _widget_value_at(widgets: List[Any], index: int) -> Optional[Any]: + if index < 0 or index >= len(widgets): + return None + return widgets[index] + + def _extract_lyrics_from_prompt_nodes(nodes_by_id: Dict[str, Dict[str, Any]]) -> Tuple[Optional[str], Optional[Any], Optional[str]]: """ Best-effort lyrics extraction for audio text-encode nodes (AceStep-like nodes). @@ -675,54 +714,83 @@ def _extract_lyrics_from_prompt_nodes(nodes_by_id: Dict[str, Dict[str, Any]]) -> ct = _lower(_node_type(node)) if "textencode" not in ct and "lyrics" not in ct: continue - ins = _inputs(node) - - is_acestep_task = ("acestep15tasktextencode" in ct) or ("acesteptasktextencode" in ct) - - lyrics = None - lyric_keys: Tuple[str, ...] = ("lyrics", "lyric", "lyric_text", "text_lyrics") - # ACEStep task encoders often store the full song text in a generic textbox. - if is_acestep_task: - lyric_keys = ("lyrics", "lyric", "lyric_text", "text_lyrics", "task_text", "task", "text") - - for key in lyric_keys: - v = ins.get(key) - if isinstance(v, str) and v.strip(): - lyrics = v.strip() - break - - strength = None - for key in ("lyrics_strength", "lyric_strength"): - v = _scalar(ins.get(key)) - if v is not None: - strength = v - break - - widgets = node.get("widgets_values") - if isinstance(widgets, list): - # AceStep: widgets_values[0]=tags, [1]=lyrics, [2]=lyrics_strength - if not lyrics and len(widgets) > 1 and isinstance(widgets[1], str) and widgets[1].strip(): - lyrics = widgets[1].strip() - if strength is None and len(widgets) > 2: - strength = _scalar(widgets[2]) - elif isinstance(widgets, dict): - for key in ("lyrics", "lyric_text", "text_lyrics", "task_text", "task", "text"): - v = widgets.get(key) - if isinstance(v, str) and v.strip(): - lyrics = lyrics or v.strip() - break - if strength is None: - for key in ("lyrics_strength", "lyric_strength"): - v = _scalar(widgets.get(key)) - if v is not None: - strength = v - break - + lyrics, strength = _extract_lyrics_from_node(node, ct) if isinstance(lyrics, str) and lyrics.strip(): return lyrics, strength, f"{_node_type(node)}:{nid}" return None, None, None +def _extract_lyrics_from_node(node: Dict[str, Any], ct: str) -> Tuple[Optional[str], Optional[Any]]: + ins = _inputs(node) + lyrics = _extract_lyrics_from_inputs(ins, ct) + strength = _extract_lyrics_strength(ins) + return _apply_lyrics_widget_fallback(node.get("widgets_values"), lyrics, strength) + + +def _extract_lyrics_from_inputs(ins: Dict[str, Any], ct: str) -> Optional[str]: + lyric_keys: Tuple[str, ...] = ("lyrics", "lyric", "lyric_text", "text_lyrics") + if "acestep15tasktextencode" in ct or "acesteptasktextencode" in ct: + lyric_keys = ("lyrics", "lyric", "lyric_text", "text_lyrics", "task_text", "task", "text") + for key in lyric_keys: + value = ins.get(key) + if isinstance(value, str) and value.strip(): + return value.strip() + return None + + +def _extract_lyrics_strength(ins: Dict[str, Any]) -> Optional[Any]: + for key in ("lyrics_strength", "lyric_strength"): + value = _scalar(ins.get(key)) + if value is not None: + return value + return None + + +def _apply_lyrics_widget_fallback(widgets: Any, lyrics: Optional[str], strength: Optional[Any]) -> Tuple[Optional[str], Optional[Any]]: + if isinstance(widgets, list): + return _lyrics_widget_fallback_from_list(widgets, lyrics, strength) + if isinstance(widgets, dict): + return _lyrics_widget_fallback_from_dict(widgets, lyrics, strength) + return lyrics, strength + + +def _lyrics_widget_fallback_from_list( + widgets: list[Any], lyrics: Optional[str], strength: Optional[Any] +) -> Tuple[Optional[str], Optional[Any]]: + if not lyrics and len(widgets) > 1 and isinstance(widgets[1], str) and widgets[1].strip(): + lyrics = widgets[1].strip() + if strength is None and len(widgets) > 2: + strength = _scalar(widgets[2]) + return lyrics, strength + + +def _lyrics_widget_fallback_from_dict( + widgets: Dict[str, Any], lyrics: Optional[str], strength: Optional[Any] +) -> Tuple[Optional[str], Optional[Any]]: + lyrics = lyrics or _first_non_empty_string( + widgets, ("lyrics", "lyric_text", "text_lyrics", "task_text", "task", "text") + ) + if strength is None: + strength = _first_non_none_scalar(widgets, ("lyrics_strength", "lyric_strength")) + return lyrics, strength + + +def _first_non_empty_string(source: Dict[str, Any], keys: Tuple[str, ...]) -> Optional[str]: + for key in keys: + value = source.get(key) + if isinstance(value, str) and value.strip(): + return value.strip() + return None + + +def _first_non_none_scalar(source: Dict[str, Any], keys: Tuple[str, ...]) -> Optional[Any]: + for key in keys: + value = _scalar(source.get(key)) + if value is not None: + return value + return None + + def _resolve_scalar_from_link(nodes_by_id: Dict[str, Dict[str, Any]], value: Any) -> Optional[Any]: src_id = _walk_passthrough(nodes_by_id, value) if not src_id: @@ -830,82 +898,21 @@ def _collect_text_encoder_nodes_from_conditioning( stack: List[Tuple[str, int]] = [(start_id, 0)] found: List[str] = [] - def _should_expand(node: Dict[str, Any]) -> bool: - ct = _lower(_node_type(node)) - # Negative "ConditioningZeroOut" chains don't represent a user-authored negative prompt. - # Avoid surfacing the positive prompt as the negative prompt when a graph uses - # ConditioningZeroOut to satisfy a required negative input. - if branch == "negative" and "conditioningzeroout" in ct: - return False - - # Optimization: Pass through Area Setters transparently - if "conditioningsetarea" in ct: - return True - - if _is_reroute(node): - return True - if "conditioning" in ct: - return True - ins = _inputs(node) - if any("conditioning" in str(k).lower() for k in ins.keys()): - return True - # Some pipelines (e.g. Wan/VHS wrappers) pass conditioning through nodes that - # expose `positive` / `negative` links but aren't named "Conditioning*". - try: - if branch in ("positive", "negative") and _is_link(ins.get(branch)): - return True - if _is_link(ins.get("positive")) or _is_link(ins.get("negative")): - return True - except (TypeError, ValueError, KeyError): - return False - return False - while stack and len(visited) < max_nodes: nid, depth = stack.pop() - - # Enforce depth limit - if depth > max_depth: - continue - - if nid in visited: + if depth > max_depth or nid in visited: continue visited.add(nid) node = nodes_by_id.get(nid) if not isinstance(node, dict): continue - if _looks_like_text_encoder(node): + if _looks_like_text_encoder(node) or _looks_like_conditioning_text(node): found.append(nid) continue - if _looks_like_conditioning_text(node): - found.append(nid) + if not _conditioning_should_expand(node, branch): continue - - if not _should_expand(node): - continue - - ins = _inputs(node) - - # If the caller requested a branch ("positive"/"negative") and this node has that - # explicit input, follow only that path to avoid mixing pos/neg prompts. - if branch in ("positive", "negative") and _is_link(ins.get(branch)): - v = ins.get(branch) - src_id = _walk_passthrough(nodes_by_id, v) - if src_id and src_id not in visited: - stack.append((src_id, depth + 1)) - continue - - for k, v in ins.items(): - k_s = str(k).lower() - if branch == "positive" and (k_s in ("negative", "neg", "negative_prompt") or k_s.startswith("negative_")): - continue - if branch == "negative" and (k_s in ("positive", "pos", "positive_prompt") or k_s.startswith("positive_")): - continue - - if _is_link(v): - src_id = _walk_passthrough(nodes_by_id, v) - if src_id and src_id not in visited: - stack.append((src_id, depth + 1)) + _expand_conditioning_frontier(nodes_by_id, stack, visited, _inputs(node), depth, branch) # Stable ordering: node id is usually numeric def _nid_key(x: str) -> Tuple[int, str]: @@ -918,6 +925,67 @@ def _collect_text_encoder_nodes_from_conditioning( return found +def _conditioning_should_expand(node: Dict[str, Any], branch: Optional[str]) -> bool: + ct = _lower(_node_type(node)) + if _conditioning_is_blocked_by_branch(ct, branch): + return False + if _conditioning_is_expand_node(ct, node): + return True + ins = _inputs(node) + if any("conditioning" in str(key).lower() for key in ins.keys()): + return True + try: + return _conditioning_has_branch_links(ins, branch) + except (TypeError, ValueError, KeyError): + return False + + +def _conditioning_is_blocked_by_branch(ct: str, branch: Optional[str]) -> bool: + return branch == "negative" and "conditioningzeroout" in ct + + +def _conditioning_is_expand_node(ct: str, node: Dict[str, Any]) -> bool: + return "conditioningsetarea" in ct or _is_reroute(node) or "conditioning" in ct + + +def _conditioning_has_branch_links(ins: Dict[str, Any], branch: Optional[str]) -> bool: + if branch in ("positive", "negative") and _is_link(ins.get(branch)): + return True + return _is_link(ins.get("positive")) or _is_link(ins.get("negative")) + + +def _expand_conditioning_frontier( + nodes_by_id: Dict[str, Dict[str, Any]], + stack: List[Tuple[str, int]], + visited: Set[str], + ins: Dict[str, Any], + depth: int, + branch: Optional[str], +) -> None: + if branch in ("positive", "negative") and _is_link(ins.get(branch)): + src_id = _walk_passthrough(nodes_by_id, ins.get(branch)) + if src_id and src_id not in visited: + stack.append((src_id, depth + 1)) + return + for key, value in ins.items(): + if not _conditioning_key_allowed(key, branch): + continue + if not _is_link(value): + continue + src_id = _walk_passthrough(nodes_by_id, value) + if src_id and src_id not in visited: + stack.append((src_id, depth + 1)) + + +def _conditioning_key_allowed(key: Any, branch: Optional[str]) -> bool: + key_s = str(key).lower() + if branch == "positive": + return key_s not in ("negative", "neg", "negative_prompt") and not key_s.startswith("negative_") + if branch == "negative": + return key_s not in ("positive", "pos", "positive_prompt") and not key_s.startswith("positive_") + return True + + def _collect_texts_from_conditioning( nodes_by_id: Dict[str, Dict[str, Any]], start_link: Any, max_nodes: int = DEFAULT_MAX_LINK_NODES, branch: Optional[str] = None ) -> List[Tuple[str, str]]: @@ -946,6 +1014,320 @@ def _collect_texts_from_conditioning( return out +def _first_model_string_from_inputs(ins: Dict[str, Any]) -> Optional[str]: + for key in ( + "ckpt_name", + "checkpoint", + "checkpoint_name", + "model_name", + "model", + "diffusion_name", + "diffusion", + "diffusion_model", + "unet_name", + "unet", + ): + value = ins.get(key) + if isinstance(value, str) and value.strip(): + return value.strip() + + for value in ins.values(): + if not isinstance(value, str): + continue + model_value = value.strip() + if not model_value: + continue + lower_value = model_value.lower().replace("\\", "/") + if any(lower_value.endswith(ext) for ext in _MODEL_EXTS): + return model_value + return None + + +def _is_lora_loader_node(ct: str, ins: Dict[str, Any]) -> bool: + return ("lora" in ct) or (ins.get("lora_name") is not None and _is_link(ins.get("model"))) + + +def _append_lora_entries( + node: Dict[str, Any], + node_id: str, + ins: Dict[str, Any], + confidence: str, + loras: List[Dict[str, Any]], +) -> None: + for key, value in ins.items(): + payload = _build_lora_payload_from_nested_value( + node=node, + node_id=node_id, + key=key, + value=value, + confidence=confidence, + ) + if payload is not None: + loras.append(payload) + + root_payload = _build_lora_payload_from_inputs(node=node, node_id=node_id, ins=ins, confidence=confidence) + if root_payload is not None: + loras.append(root_payload) + + +def _build_lora_payload_from_nested_value( + *, + node: Dict[str, Any], + node_id: str, + key: Any, + value: Any, + confidence: str, +) -> Optional[Dict[str, Any]]: + if not _is_nested_lora_key(key) or not isinstance(value, dict): + return None + if not _is_enabled_lora_value(value): + return None + name = _nested_lora_name(value) + if not name: + return None + return { + "name": name, + "strength_model": _nested_lora_strength(value), + "strength_clip": value.get("strength_clip") or value.get("clip_strength"), + "confidence": confidence, + "source": f"{_node_type(node)}:{node_id}:{key}", + } + + +def _is_nested_lora_key(key: Any) -> bool: + return str(key).lower().startswith("lora_") + + +def _is_enabled_lora_value(value: Dict[str, Any]) -> bool: + return value.get("on") is not False + + +def _nested_lora_name(value: Dict[str, Any]) -> Optional[str]: + return _clean_model_id(value.get("lora") or value.get("lora_name") or value.get("name")) + + +def _nested_lora_strength(value: Dict[str, Any]) -> Any: + return value.get("strength") or value.get("strength_model") or value.get("weight") or value.get("lora_strength") + + +def _build_lora_payload_from_inputs( + *, + node: Dict[str, Any], + node_id: str, + ins: Dict[str, Any], + confidence: str, +) -> Optional[Dict[str, Any]]: + name = _clean_model_id(ins.get("lora_name") or ins.get("lora") or ins.get("name")) + if not name: + return None + strength_model = ins.get("strength_model") or ins.get("strength") or ins.get("weight") or ins.get("lora_strength") + strength_clip = ins.get("strength_clip") or ins.get("clip_strength") + return { + "name": name, + "strength_model": strength_model, + "strength_clip": strength_clip, + "confidence": confidence, + "source": f"{_node_type(node)}:{node_id}", + } + + +def _is_diffusion_loader_node(ct: str) -> bool: + return ( + "loaddiffusionmodel" in ct + or "diffusionmodel" in ct + or "unetloader" in ct + or "loadunet" in ct + or ct == "unet" + or "videomodel" in ct + ) + + +def _is_generic_model_loader_node(ct: str) -> bool: + return ( + "modelloader" in ct + or "model_loader" in ct + or "model-loader" in ct + or "ltxvideomodel" in ct + or "wanvideomodel" in ct + or "hyvideomodel" in ct + or "cogvideomodel" in ct + ) + + +def _is_checkpoint_loader_node(ct: str, ins: Dict[str, Any]) -> bool: + if ins.get("ckpt_name") is not None: + return True + return any( + token in ct + for token in ( + "checkpointloader", + "checkpoint_loader", + "loadcheckpoint", + "load_checkpoint", + ) + ) + + +def _chain_result(next_link: Optional[Any], should_stop: bool) -> Tuple[Optional[Any], bool]: + return next_link, should_stop + + +def _handle_lora_chain_node( + node: Dict[str, Any], + node_id: str, + ct: str, + ins: Dict[str, Any], + loras: List[Dict[str, Any]], + confidence: str, +) -> Optional[Tuple[Optional[Any], bool]]: + if not _is_lora_loader_node(ct, ins): + return None + _append_lora_entries(node, node_id, ins, confidence, loras) + next_model = ins.get("model") + if _is_link(next_model): + return _chain_result(next_model, False) + return _chain_result(None, True) + + +def _handle_modelsampling_chain_node(ct: str, ins: Dict[str, Any]) -> Optional[Tuple[Optional[Any], bool]]: + if ("modelsampling" in ct or "model_sampling" in ct) and _is_link(ins.get("model")): + return _chain_result(ins.get("model"), False) + return None + + +def _handle_diffusion_loader_chain_node( + node: Dict[str, Any], + node_id: str, + ct: str, + ins: Dict[str, Any], + models: Dict[str, Dict[str, Any]], + confidence: str, +) -> Optional[Tuple[Optional[Any], bool]]: + if not _is_diffusion_loader_node(ct): + return None + source = f"{_node_type(node)}:{node_id}" + _set_model_entry_if_missing( + models, + "unet", + _clean_model_id(ins.get("unet_name") or ins.get("unet")), + confidence, + source, + ) + _set_model_entry_if_missing( + models, + "diffusion", + _clean_model_id( + ins.get("diffusion_name") or ins.get("diffusion") or ins.get("model_name") or ins.get("ckpt_name") or ins.get("model") + ), + confidence, + source, + ) + return _chain_result_from_model_input(ins) + + +def _set_model_entry_if_missing( + models: Dict[str, Dict[str, Any]], + key: str, + name: Optional[str], + confidence: str, + source: str, +) -> None: + if name and key not in models: + models[key] = {"name": name, "confidence": confidence, "source": source} + + +def _chain_result_from_model_input(ins: Dict[str, Any]) -> Tuple[Optional[Any], bool]: + next_model = ins.get("model") + if _is_link(next_model): + return _chain_result(next_model, False) + return _chain_result(None, True) + + +def _handle_generic_loader_chain_node( + node: Dict[str, Any], + node_id: str, + ct: str, + ins: Dict[str, Any], + models: Dict[str, Dict[str, Any]], + confidence: str, +) -> Optional[Tuple[Optional[Any], bool]]: + if not _is_generic_model_loader_node(ct): + return None + name = _clean_model_id(_first_model_string_from_inputs(ins)) + if name: + models.setdefault( + "checkpoint", + {"name": name, "confidence": confidence, "source": f"{_node_type(node)}:{node_id}"}, + ) + return _chain_result(None, True) + + +def _handle_checkpoint_loader_chain_node( + node: Dict[str, Any], + node_id: str, + ct: str, + ins: Dict[str, Any], + models: Dict[str, Dict[str, Any]], + confidence: str, +) -> Optional[Tuple[Optional[Any], bool]]: + if not _is_checkpoint_loader_node(ct, ins): + return None + ckpt = _clean_model_id(ins.get("ckpt_name") or ins.get("model_name")) + if ckpt: + models.setdefault( + "checkpoint", + {"name": ckpt, "confidence": confidence, "source": f"{_node_type(node)}:{node_id}"}, + ) + return _chain_result(None, True) + + +def _handle_switch_selector_chain_node(ct: str, ins: Dict[str, Any]) -> Optional[Tuple[Optional[Any], bool]]: + if ("switch" in ct or "selector" in ct) and not _is_link(ins.get("model")): + links = [value for value in ins.values() if _is_link(value)] + if len(links) == 1: + return _chain_result(links[0], False) + return None + + +def _handle_fallback_chain_node( + node: Dict[str, Any], + node_id: str, + ins: Dict[str, Any], + models: Dict[str, Dict[str, Any]], + confidence: str, +) -> Tuple[Optional[Any], bool]: + name = _clean_model_id(_first_model_string_from_inputs(ins)) + if name and "checkpoint" not in models: + models["checkpoint"] = {"name": name, "confidence": confidence, "source": f"{_node_type(node)}:{node_id}"} + next_model = ins.get("model") + if _is_link(next_model): + return _chain_result(next_model, False) + return _chain_result(None, True) + + +def _process_model_chain_node( + node: Dict[str, Any], + node_id: str, + ct: str, + ins: Dict[str, Any], + models: Dict[str, Dict[str, Any]], + loras: List[Dict[str, Any]], + confidence: str, +) -> Tuple[Optional[Any], bool]: + for handler in ( + lambda: _handle_lora_chain_node(node, node_id, ct, ins, loras, confidence), + lambda: _handle_modelsampling_chain_node(ct, ins), + lambda: _handle_diffusion_loader_chain_node(node, node_id, ct, ins, models, confidence), + lambda: _handle_generic_loader_chain_node(node, node_id, ct, ins, models, confidence), + lambda: _handle_checkpoint_loader_chain_node(node, node_id, ct, ins, models, confidence), + lambda: _handle_switch_selector_chain_node(ct, ins), + ): + handled = handler() + if handled is not None: + return handled + return _handle_fallback_chain_node(node, node_id, ins, models, confidence) + + def _trace_model_chain( nodes_by_id: Dict[str, Dict[str, Any]], model_link: Any, confidence: str ) -> Tuple[Dict[str, Dict[str, Any]], List[Dict[str, Any]]]: @@ -962,172 +1344,19 @@ def _trace_model_chain( node = nodes_by_id.get(node_id) if not isinstance(node, dict): break - ct = _lower(_node_type(node)) - ins = _inputs(node) - - def _first_model_string() -> Optional[str]: - for k in ( - "ckpt_name", - "checkpoint", - "checkpoint_name", - "model_name", - "model", - "diffusion_name", - "diffusion", - "diffusion_model", - "unet_name", - "unet", - ): - v = ins.get(k) - if isinstance(v, str) and v.strip(): - return v.strip() - # Last resort: scan for any model-like filename. - for v in ins.values(): - if not isinstance(v, str): - continue - s = v.strip() - if not s: - continue - lower_s = s.lower().replace("\\", "/") - if any(lower_s.endswith(ext) for ext in _MODEL_EXTS): - return s - return None - - # Some custom LoRA loader nodes do not contain "lora" in their class_type, - # but expose `lora_name` + a linked `model` input. Treat these as LoRA loaders. - if ("lora" in ct) or (ins.get("lora_name") is not None and _is_link(ins.get("model"))): - # rgthree "Power Lora Loader" stores multiple LoRAs under lora_1/lora_2/... - # objects instead of a flat `lora_name`. - for k, v in ins.items(): - if not str(k).lower().startswith("lora_"): - continue - if not isinstance(v, dict): - continue - enabled = v.get("on") - if enabled is False: - continue - name = _clean_model_id(v.get("lora") or v.get("lora_name") or v.get("name")) - if not name: - continue - strength = v.get("strength") or v.get("strength_model") or v.get("weight") or v.get("lora_strength") - loras.append( - { - "name": name, - "strength_model": strength, - "strength_clip": v.get("strength_clip") or v.get("clip_strength"), - "confidence": confidence, - "source": f"{_node_type(node)}:{node_id}:{k}", - } - ) - - name = _clean_model_id(ins.get("lora_name") or ins.get("lora") or ins.get("name")) - if name: - strength_model = ins.get("strength_model") or ins.get("strength") or ins.get("weight") or ins.get("lora_strength") - strength_clip = ins.get("strength_clip") or ins.get("clip_strength") - loras.append( - { - "name": name, - "strength_model": strength_model, - "strength_clip": strength_clip, - "confidence": confidence, - "source": f"{_node_type(node)}:{node_id}", - } - ) - - # If this node is LoRA-ish but didn't yield any LoRA entries, still follow the chain. - current_link = ins.get("model") if _is_link(ins.get("model")) else None - if current_link is not None: - continue - # No model link to follow; stop here. - break - - # Model sampling / patch nodes are common in video stacks (e.g. ModelSamplingSD3). - # They typically just transform a model object; follow their `model` input. - if ("modelsampling" in ct or "model_sampling" in ct) and _is_link(ins.get("model")): - current_link = ins.get("model") - continue - - # Common diffusion model loaders (video stacks, gguf, unet-only, etc.) - if ( - "loaddiffusionmodel" in ct - or "diffusionmodel" in ct - or "unetloader" in ct - or "loadunet" in ct - or "unet" == ct - or "videomodel" in ct # Helper for LTX/Wan video models - ): - unet = _clean_model_id(ins.get("unet_name") or ins.get("unet")) - diffusion = _clean_model_id(ins.get("diffusion_name") or ins.get("diffusion") or ins.get("model_name") or ins.get("ckpt_name") or ins.get("model")) - if unet and "unet" not in models: - models["unet"] = {"name": unet, "confidence": confidence, "source": f"{_node_type(node)}:{node_id}"} - if diffusion and "diffusion" not in models: - models["diffusion"] = {"name": diffusion, "confidence": confidence, "source": f"{_node_type(node)}:{node_id}"} - next_model = ins.get("model") - if _is_link(next_model): - current_link = next_model - continue - break - - # Generic "model loader" style custom nodes (e.g. WanVideoModelLoader) often expose - # only a model path/identifier without "ckpt_name" naming. - if ( - "modelloader" in ct - or "model_loader" in ct - or "model-loader" in ct - or "ltxvideomodel" in ct - or "wanvideomodel" in ct - or "hyvideomodel" in ct - or "cogvideomodel" in ct - ): - raw = _first_model_string() - name = _clean_model_id(raw) - if name: - # Many video wrappers use a single diffusion model file; map it to checkpoint. - models.setdefault( - "checkpoint", - {"name": name, "confidence": confidence, "source": f"{_node_type(node)}:{node_id}"}, - ) - break - - is_checkpoint_loader = any( - s in ct - for s in ( - "checkpointloader", - "checkpoint_loader", - "loadcheckpoint", - "load_checkpoint", - ) + next_link, should_stop = _process_model_chain_node( + node, + node_id, + _lower(_node_type(node)), + _inputs(node), + models, + loras, + confidence, ) - if is_checkpoint_loader or ins.get("ckpt_name") is not None: - ckpt = _clean_model_id(ins.get("ckpt_name") or ins.get("model_name")) - if ckpt: - models.setdefault( - "checkpoint", - {"name": ckpt, "confidence": confidence, "source": f"{_node_type(node)}:{node_id}"}, - ) + if should_stop: break - - # Some graphs insert generic "switch"/"selector" nodes in the model chain (e.g. rgthree). - # Follow the single upstream link when present to reach the actual loader node. - if ("switch" in ct or "selector" in ct) and not _is_link(ins.get("model")): - links = [v for v in ins.values() if _is_link(v)] - if len(links) == 1: - current_link = links[0] - continue - - # Heuristic: if the node exposes a model-like filename, record it as checkpoint. - raw = _first_model_string() - name = _clean_model_id(raw) - if name and "checkpoint" not in models: - models["checkpoint"] = {"name": name, "confidence": confidence, "source": f"{_node_type(node)}:{node_id}"} - - # Generic model passthrough - next_model = ins.get("model") - if _is_link(next_model): - current_link = next_model - continue - break + current_link = next_link return models, loras @@ -1144,33 +1373,44 @@ def _trace_named_loader(nodes_by_id: Dict[str, Dict[str, Any]], link: Any, keys: if not isinstance(node, dict): return None ins = _inputs(node) - # Special case: DualCLIPLoader-style nodes expose clip_name1/clip_name2. - if "clip_name1" in keys and "clip_name2" in keys: - c1 = _clean_model_id(ins.get("clip_name1")) - c2 = _clean_model_id(ins.get("clip_name2")) - if c1 and c2: - return { - "name": f"{c1} + {c2}", - "confidence": confidence, - "source": f"{_node_type(node)}:{node_id}", - } - for k in keys: - name = _clean_model_id(ins.get(k)) - if name: - return {"name": name, "confidence": confidence, "source": f"{_node_type(node)}:{node_id}"} - next_link = ins.get("clip") if _is_link(ins.get("clip")) else None - if _is_link(next_link): - current_link = next_link - continue - next_link = ins.get("vae") if _is_link(ins.get("vae")) else None - if _is_link(next_link): - current_link = next_link - continue - next_link = ins.get("model") if _is_link(ins.get("model")) else None - if _is_link(next_link): - current_link = next_link - continue + source = f"{_node_type(node)}:{node_id}" + named = _extract_named_loader_model(ins, keys, confidence=confidence, source=source) + if named is not None: + return named + current_link = _next_named_loader_link(ins) + if current_link is None: + return None + return None + + +def _extract_named_loader_model( + ins: Dict[str, Any], keys: Tuple[str, ...], *, confidence: str, source: str +) -> Optional[Dict[str, Any]]: + dual_clip = _extract_dual_clip_name(ins, keys) + if dual_clip is not None: + return {"name": dual_clip, "confidence": confidence, "source": source} + for key in keys: + name = _clean_model_id(ins.get(key)) + if name: + return {"name": name, "confidence": confidence, "source": source} + return None + + +def _extract_dual_clip_name(ins: Dict[str, Any], keys: Tuple[str, ...]) -> Optional[str]: + if "clip_name1" not in keys or "clip_name2" not in keys: return None + c1 = _clean_model_id(ins.get("clip_name1")) + c2 = _clean_model_id(ins.get("clip_name2")) + if c1 and c2: + return f"{c1} + {c2}" + return None + + +def _next_named_loader_link(ins: Dict[str, Any]) -> Optional[Any]: + for key in ("clip", "vae", "model"): + value = ins.get(key) + if _is_link(value): + return value return None @@ -1260,26 +1500,37 @@ def _trace_size(nodes_by_id: Dict[str, Dict[str, Any]], latent_link: Any, confid return None ct = _lower(_node_type(node)) ins = _inputs(node) - - if "emptylatentimage" in ct: - w = _scalar(ins.get("width")) - h = _scalar(ins.get("height")) - return _field_size(w, h, confidence, f"{_node_type(node)}:{node_id}") - - # Common resize/upscale latent nodes can carry explicit width/height - w = _scalar(ins.get("width")) - h = _scalar(ins.get("height")) - if w is not None and h is not None: - return _field_size(w, h, confidence, f"{_node_type(node)}:{node_id}") - - # Pass-through along latent/sample link - for key in ("samples", "latent", "latent_image", "image"): - v = ins.get(key) - if _is_link(v): - current_link = v - break - else: + size = _size_field_from_node(node, node_id, ct, ins, confidence) + if size: + return size + next_link = _next_latent_link(ins) + if not next_link: return None + current_link = next_link + return None + + +def _size_field_from_node( + node: Dict[str, Any], + node_id: str, + node_type: str, + ins: Dict[str, Any], + confidence: str, +) -> Optional[Dict[str, Any]]: + width = _scalar(ins.get("width")) + height = _scalar(ins.get("height")) + if "emptylatentimage" in node_type: + return _field_size(width, height, confidence, f"{_node_type(node)}:{node_id}") + if width is None or height is None: + return None + return _field_size(width, height, confidence, f"{_node_type(node)}:{node_id}") + + +def _next_latent_link(ins: Dict[str, Any]) -> Optional[Any]: + for key in ("samples", "latent", "latent_image", "image"): + value = ins.get(key) + if _is_link(value): + return value return None @@ -1300,38 +1551,47 @@ def _trace_guidance_value(nodes_by_id: Dict[str, Dict[str, Any]], start_link: An if nid in visited or depth > max_hops: continue visited.add(nid) - + node = nodes_by_id.get(nid) if not isinstance(node, dict): continue - + ins = _inputs(node) # Check if this node has guidance g_val = _scalar(ins.get("guidance")) if g_val is not None: return float(g_val), f"{_node_type(node)}:{nid}" - - # Traverse upstream - ct = _lower(_node_type(node)) - should_expand = False - - if _is_reroute(node): - should_expand = True - elif "conditioning" in ct or "guider" in ct or "flux" in ct: - should_expand = True - elif any("conditioning" in str(k).lower() for k in ins.keys()): - should_expand = True - - if should_expand: - # Add conditioning sources to stack - for k, v in ins.items(): - if "conditioning" in str(k).lower() and _is_link(v): - src = _walk_passthrough(nodes_by_id, v) - if src and src not in visited: - stack.append((src, depth + 1)) + + if not _guidance_should_expand(node, ins): + continue + for src in _iter_guidance_conditioning_sources(nodes_by_id, ins): + if src not in visited: + stack.append((src, depth + 1)) return None +def _guidance_should_expand(node: Dict[str, Any], ins: Dict[str, Any]) -> bool: + if _is_reroute(node): + return True + ct = _lower(_node_type(node)) + if "conditioning" in ct or "guider" in ct or "flux" in ct: + return True + return any("conditioning" in str(k).lower() for k in ins.keys()) + + +def _iter_guidance_conditioning_sources( + nodes_by_id: Dict[str, Dict[str, Any]], ins: Dict[str, Any] +) -> List[str]: + sources: List[str] = [] + for key, value in ins.items(): + if "conditioning" not in str(key).lower() or not _is_link(value): + continue + src = _walk_passthrough(nodes_by_id, value) + if src: + sources.append(src) + return sources + + def _extract_workflow_metadata(workflow: Any) -> Dict[str, Any]: meta = {} if isinstance(workflow, dict): @@ -1344,210 +1604,327 @@ def _extract_workflow_metadata(workflow: Any) -> Dict[str, Any]: +_INPUT_ROLE_PRIORITY: Tuple[str, ...] = ( + "first_frame", + "last_frame", + "mask/inpaint", + "depth", + "control_video", + "control_image", + "control", + "source", + "style/reference", + "frame_range", +) +_INPUT_FILENAME_KEYS: Tuple[str, ...] = ( + "image", + "video", + "filename", + "audio", + "file", + "media_source", + "path", + "image_path", + "video_path", + "audio_path", +) + + +def _subject_role_hints(subject: Optional[Dict[str, Any]]) -> Set[str]: + roles: Set[str] = set() + if not isinstance(subject, dict): + return roles + title = _lower(subject.get("_meta", {}).get("title", "") or subject.get("title", "")) + ins = _inputs(subject) + filename = _lower(str(ins.get("image", "") or ins.get("video", "") or "")) + _add_subject_role_if_match(roles, "first_frame", title=title, filename=filename, title_tokens=("first", "start"), filename_tokens=("first",)) + _add_subject_role_if_match(roles, "last_frame", title=title, filename=filename, title_tokens=("last", "end"), filename_tokens=("last",)) + _add_subject_role_if_match(roles, "control", title=title, filename=filename, title_tokens=("control",), filename_tokens=()) + _add_subject_role_if_match(roles, "mask/inpaint", title=title, filename=filename, title_tokens=("mask", "inpaint"), filename_tokens=()) + _add_subject_role_if_match(roles, "depth", title=title, filename=filename, title_tokens=("depth",), filename_tokens=()) + _add_subject_role_if_match( + roles, + "style/reference", + title=title, + filename=filename, + title_tokens=("reference", "ref", "style"), + filename_tokens=(), + ) + return roles + + +def _contains_any_token(value: str, tokens: Tuple[str, ...]) -> bool: + return any(token in value for token in tokens) + + +def _add_subject_role_if_match( + roles: Set[str], + role: str, + *, + title: str, + filename: str, + title_tokens: Tuple[str, ...], + filename_tokens: Tuple[str, ...], +) -> None: + if _contains_any_token(title, title_tokens) or _contains_any_token(filename, filename_tokens): + roles.add(role) + + +def _linked_input_from_frontier(node: Dict[str, Any], frontier: Set[str]) -> Optional[str]: + for key, value in _inputs(node).items(): + if not _is_link(value): + continue + resolved = _resolve_link(value) + if not resolved: + continue + src_id, _ = resolved + if str(src_id) in frontier: + return str(key).lower() + return None + + +def _classify_control_or_mask_role( + target_type: str, + hit_input_name: str, + subject_is_video: bool, +) -> Optional[str]: + if "ipadapter" in target_type: + return "style/reference" + if "controlnet" in target_type: + return "control_video" if subject_is_video else "control_image" + if "mask" in hit_input_name or "mask" in target_type or "inpaint" in target_type: + return "mask/inpaint" + if "depth" in target_type or "depth" in hit_input_name: + return "depth" + return None + + +def _classify_vace_or_range_role(target_type: str, hit_input_name: str) -> Optional[str]: + if _is_vace_target_type(target_type): + edge_role = _frame_edge_role(hit_input_name) + if "control" in hit_input_name or "reference" in hit_input_name: + return "control_video" + return edge_role or "source" + if _is_frame_range_target_type(target_type): + return _frame_edge_role(hit_input_name) or "frame_range" + return None + + +def _is_vace_target_type(target_type: str) -> bool: + return "vace" in target_type or "wanvace" in target_type + + +def _is_frame_range_target_type(target_type: str) -> bool: + return "starttoend" in target_type or "framerange" in target_type + + +def _frame_edge_role(hit_input_name: str) -> Optional[str]: + if "start" in hit_input_name or "first" in hit_input_name: + return "first_frame" + if "end" in hit_input_name or "last" in hit_input_name: + return "last_frame" + return None + + +def _classify_generic_source_role(target_type: str, hit_input_name: str) -> Optional[str]: + edge_role = _edge_role_from_input_name(hit_input_name) + if edge_role: + return edge_role + if _target_type_is_source_like(target_type): + return "source" + if _sampler_input_is_source(target_type, hit_input_name): + return "source" + return None + + +def _edge_role_from_input_name(hit_input_name: str) -> Optional[str]: + if "first" in hit_input_name or "start" in hit_input_name: + return "first_frame" + if "last" in hit_input_name or "end" in hit_input_name: + return "last_frame" + return None + + +def _target_type_is_source_like(target_type: str) -> bool: + return "img2vid" in target_type or "i2v" in target_type or "vaeencode" in target_type + + +def _sampler_input_is_source(target_type: str, hit_input_name: str) -> bool: + return "sampler" in target_type and ("image" in hit_input_name or "latent" in hit_input_name) + + +def _classify_downstream_input_role( + target_type: str, + hit_input_name: str, + subject_is_video: bool, +) -> Tuple[Optional[str], bool]: + role = _classify_control_or_mask_role(target_type, hit_input_name, subject_is_video) + if role: + return role, False + role = _classify_vace_or_range_role(target_type, hit_input_name) + if role: + return role, False + role = _classify_generic_source_role(target_type, hit_input_name) + if role: + return role, False + return None, True + + def _detect_input_role(nodes_by_id: Dict[str, Any], subject_node_id: str) -> str: """ - Determine how an input node is used (e.g. 'first_frame', 'last_frame', 'reference', + Determine how an input node is used (e.g. 'first_frame', 'last_frame', 'reference', 'control_video', 'control_image', 'mask/inpaint', 'source', 'style', 'depth'). Performs a limited-depth BFS downstream to handle intermediate nodes (Resize, VAE, etc). """ - roles = set() - frontier = {str(subject_node_id)} - visited = {str(subject_node_id)} - - # Also check the subject node's title for hints (e.g. "FIRST_FRAME", "LAST_FRAME") - subject = nodes_by_id.get(str(subject_node_id)) or nodes_by_id.get(subject_node_id) - if isinstance(subject, dict): - title = _lower(subject.get("_meta", {}).get("title", "") or subject.get("title", "")) - filename = "" - ins = _inputs(subject) - filename = _lower(str(ins.get("image", "") or ins.get("video", "") or "")) - - # Check title/filename for first/last frame hints - if "first" in title or "first" in filename or "start" in title: - roles.add("first_frame") - if "last" in title or "last" in filename or "end" in title: - roles.add("last_frame") - if "control" in title: - roles.add("control") - if "mask" in title or "inpaint" in title: - roles.add("mask/inpaint") - if "depth" in title: - roles.add("depth") - if "reference" in title or "ref" in title or "style" in title: - roles.add("style/reference") - - # Limit depth to avoid infinite loops or massive fan-out - for _ in range(8): - if not frontier: break - next_frontier = set() - - for nid, node in nodes_by_id.items(): - if str(nid) in visited: continue - if not isinstance(node, dict): continue - - # Check if 'node' inputs from 'frontier' - ins = _inputs(node) - linked = False - hit_input_name = "" - - for k, v in ins.items(): - if _is_link(v): - resolved = _resolve_link(v) - if not resolved: - continue - src_id, _ = resolved - if str(src_id) in frontier: - linked = True - hit_input_name = str(k).lower() - break - - if linked: - # Found a consumer - target_type = _lower(_node_type(node)) - title = _lower(node.get("_meta", {}).get("title", "") or node.get("title", "")) - - # 1. High specific roles - if "ipadapter" in target_type: - roles.add("style/reference") - elif "controlnet" in target_type: - if subject and _lower(_node_type(subject)).find("video") >= 0: - roles.add("control_video") - else: - roles.add("control_image") - elif "mask" in hit_input_name or "mask" in target_type or "inpaint" in target_type: - roles.add("mask/inpaint") - elif "depth" in target_type or "depth" in hit_input_name: - roles.add("depth") - - # 2. Video-specific nodes (WAN VACE, AnimateDiff, etc.) - elif "vace" in target_type or "wanvace" in target_type: - # VACE uses control video/image for video generation - if "control" in hit_input_name or "reference" in hit_input_name: - roles.add("control_video") - elif "start" in hit_input_name or "first" in hit_input_name: - roles.add("first_frame") - elif "end" in hit_input_name or "last" in hit_input_name: - roles.add("last_frame") - else: - roles.add("source") - - elif "starttoend" in target_type or "framerange" in target_type: - # Nodes like WanVideoVACEStartToEndFrame - if "start" in hit_input_name or "first" in hit_input_name: - roles.add("first_frame") - elif "end" in hit_input_name or "last" in hit_input_name: - roles.add("last_frame") - else: - # Check if the input connects to start_image or end_image - roles.add("frame_range") - - # 3. Frame position detection - elif "first" in hit_input_name or "start" in hit_input_name: - roles.add("first_frame") - elif "last" in hit_input_name or "end" in hit_input_name: - roles.add("last_frame") - elif "img2vid" in target_type or "i2v" in target_type: - roles.add("source") + subject_id = str(subject_node_id) + subject = nodes_by_id.get(subject_id) or nodes_by_id.get(subject_node_id) + roles = _subject_role_hints(subject if isinstance(subject, dict) else None) + subject_is_video = isinstance(subject, dict) and ("video" in _lower(_node_type(subject))) + roles.update(_detect_roles_from_downstream(nodes_by_id, subject_id, subject_is_video)) - # 4. VAE Encode -> likely source material (I2I/I2V) - elif "vaeencode" in target_type: - roles.add("source") - - # 5. If connected directly to sampler input (rare for images) - elif "sampler" in target_type: - if "image" in hit_input_name or "latent" in hit_input_name: - roles.add("source") - - else: - # Unknown/Intermediate node (Resize, Upscale, etc.) - # Add to frontier to trace what IT connects to - next_frontier.add(str(nid)) - + for role in _INPUT_ROLE_PRIORITY: + if role in roles: + return role + return "input" + + +def _detect_roles_from_downstream( + nodes_by_id: Dict[str, Any], + subject_id: str, + subject_is_video: bool, +) -> Set[str]: + roles: Set[str] = set() + frontier: Set[str] = {subject_id} + visited: Set[str] = {subject_id} + for _ in range(8): + if not frontier: + break + next_frontier = _scan_downstream_frontier( + nodes_by_id, + frontier, + visited, + subject_is_video, + roles, + ) visited.update(next_frontier) frontier = next_frontier + return roles - # Priority resolution (more specific roles first) - if "first_frame" in roles: return "first_frame" - if "last_frame" in roles: return "last_frame" - if "mask/inpaint" in roles: return "mask/inpaint" - if "depth" in roles: return "depth" - if "control_video" in roles: return "control_video" - if "control_image" in roles: return "control_image" - if "control" in roles: return "control" - if "source" in roles: return "source" - if "style/reference" in roles: return "style/reference" - if "frame_range" in roles: return "frame_range" - - return "input" # Default - generic input + +def _scan_downstream_frontier( + nodes_by_id: Dict[str, Any], + frontier: Set[str], + visited: Set[str], + subject_is_video: bool, + roles: Set[str], +) -> Set[str]: + next_frontier: Set[str] = set() + for nid, node in nodes_by_id.items(): + nid_s = str(nid) + if nid_s in visited or not isinstance(node, dict): + continue + hit_input_name = _linked_input_from_frontier(node, frontier) + if hit_input_name is None: + continue + role, should_expand = _classify_downstream_input_role( + _lower(_node_type(node)), + hit_input_name, + subject_is_video, + ) + if role: + roles.add(role) + if should_expand: + next_frontier.add(nid_s) + return next_frontier + + +def _loader_kind_from_node_type(node_type: str) -> Optional[str]: + ntype_clean = node_type.replace(" ", "").replace("_", "").replace("-", "") + is_image_load = _node_type_matches_any(ntype_clean, ("loadimage", "imageloader", "inputimage")) + is_video_load = _node_type_matches_any(ntype_clean, ("loadvideo", "videoloader", "inputvideo")) + is_audio_load = _node_type_matches_any(ntype_clean, ("loadaudio", "audioloader", "inputaudio")) + if "ipadapter" in ntype_clean and "image" in ntype_clean: + is_image_load = True + if is_audio_load: + return "audio" + if is_video_load: + return "video" + if is_image_load: + return "image" + return None + + +def _node_type_matches_any(node_type_clean: str, needles: Tuple[str, ...]) -> bool: + return any(needle in node_type_clean for needle in needles) + + +def _extract_loader_filename(node: Dict[str, Any], ins: Dict[str, Any]) -> Optional[str]: + direct = _loader_filename_from_inputs(ins) + if direct: + return direct + return _loader_filename_from_widgets(node.get("widgets_values")) + + +def _loader_filename_from_inputs(ins: Dict[str, Any]) -> Optional[str]: + for key in _INPUT_FILENAME_KEYS: + value = ins.get(key) + if isinstance(value, str) and value: + return value + return None + + +def _loader_filename_from_widgets(widgets: Any) -> Optional[str]: + if not isinstance(widgets, list): + return None + for widget_value in widgets: + if _looks_like_loader_filename(widget_value): + return str(widget_value) + return None + + +def _looks_like_loader_filename(value: Any) -> bool: + if not isinstance(value, str): + return False + return ("." in value or "/" in value or "\\" in value) and len(value) > 4 def _extract_input_files(nodes_by_id: Dict[str, Any]) -> List[Dict[str, Any]]: """ Extract input file references (LoadImage/LoadVideo/LoadAudio, etc) with usage context. """ - inputs = [] - seen = set() - + inputs: List[Dict[str, Any]] = [] + seen: Set[Tuple[str, str, str]] = set() + for nid, node in nodes_by_id.items(): if not isinstance(node, dict): continue - - ntype = _lower(_node_type(node)) - ntype_clean = ntype.replace(" ", "").replace("_", "").replace("-", "") - - # Standard input nodes (robust to spacing/naming variations) - is_image_load = "loadimage" in ntype_clean or "imageloader" in ntype_clean or "inputimage" in ntype_clean - is_video_load = "loadvideo" in ntype_clean or "videoloader" in ntype_clean or "inputvideo" in ntype_clean - is_audio_load = "loadaudio" in ntype_clean or "audioloader" in ntype_clean or "inputaudio" in ntype_clean - - # Also check for known specific nodes like LTX or IPAdapter that might not match above - if "ipadapter" in ntype_clean and "image" in ntype_clean: - is_image_load = True + node_type = _lower(_node_type(node)) + kind = _loader_kind_from_node_type(node_type) + if kind is None: + continue - if is_image_load or is_video_load or is_audio_load: - ins = _inputs(node) - - # 1. API Format inputs (broadened search for file/path keys) - filename = ( - ins.get("image") or ins.get("video") or ins.get("filename") or - ins.get("audio") or ins.get("file") or ins.get("media_source") or ins.get("path") or - ins.get("image_path") or ins.get("video_path") or ins.get("audio_path") - ) - - # 2. Workflow Format widgets_values (Heuristic) - if not filename: - widgets = node.get("widgets_values") - if isinstance(widgets, list) and len(widgets) > 0: - # Some nodes put the file at index 0, others index 1 (like "upload" switch) - # Iterate to find the first string that looks like a file path - for w in widgets: - if isinstance(w, str) and ( - "." in w or "/" in w or "\\" in w - ): - # Exclude simple extensions selection like "format: png" - if len(w) > 4: - filename = w - break + ins = _inputs(node) + filename = _extract_loader_filename(node, ins) + if not isinstance(filename, str) or not filename: + continue + + subfolder = ins.get("subfolder", "") or "" + dedupe_key = (filename, str(subfolder), node_type) + if dedupe_key in seen: + continue + seen.add(dedupe_key) + + inputs.append( + { + "filename": filename, + "subfolder": subfolder, + "type": kind, + "node_id": nid, + "folder_type": ins.get("type", "input"), + "role": _detect_input_role(nodes_by_id, nid), + } + ) - if isinstance(filename, str) and filename: - # Deduplicate by (filename, subfolder, type) - subfolder = ins.get("subfolder", "") or "" - key = (filename, str(subfolder), ntype) - - if key not in seen: - seen.add(key) - # Determine role - usage = _detect_input_role(nodes_by_id, nid) - - inputs.append({ - "filename": filename, - "subfolder": subfolder, - "type": "audio" if is_audio_load else ("video" if is_video_load else "image"), - "node_id": nid, - "folder_type": ins.get("type", "input"), - "role": usage - }) - return inputs @@ -1558,10 +1935,10 @@ def _collect_all_prompts_from_sinks( ) -> Tuple[List[str], List[str]]: """ Collect all distinct positive and negative prompts from multiple sinks. - + For multi-output workflows (like batch generation with different prompts), this gathers all unique prompts to show the user what variants were generated. - + Returns: (all_positive_prompts, all_negative_prompts) - deduplicated lists """ @@ -1569,169 +1946,207 @@ def _collect_all_prompts_from_sinks( all_negative: List[str] = [] seen_pos: Set[str] = set() seen_neg: Set[str] = set() - + for sink_id in sinks[:max_sinks]: try: - # Try primary sampler first - sampler_id, _ = _select_primary_sampler(nodes_by_id, sink_id) - if not sampler_id: - sampler_id, _ = _select_advanced_sampler(nodes_by_id, sink_id) - if not sampler_id: + sampler_node = _resolve_sink_sampler_node(nodes_by_id, sink_id) + if sampler_node is None: continue - - sampler_node = nodes_by_id.get(sampler_id) - if not isinstance(sampler_node, dict): - continue - ins = _inputs(sampler_node) - - # Extract positive prompt - if _is_link(ins.get("positive")): - items = _collect_texts_from_conditioning(nodes_by_id, ins.get("positive"), branch="positive") - for text, _ in items: - t = text.strip() - if t and t not in seen_pos: - seen_pos.add(t) - all_positive.append(t) - - # Extract negative prompt - if _is_link(ins.get("negative")): - items = _collect_texts_from_conditioning(nodes_by_id, ins.get("negative"), branch="negative") - for text, _ in items: - t = text.strip() - if t and t not in seen_neg: - seen_neg.add(t) - all_negative.append(t) - + _collect_prompt_branch_from_input(nodes_by_id, ins.get("positive"), "positive", seen_pos, all_positive) + _collect_prompt_branch_from_input(nodes_by_id, ins.get("negative"), "negative", seen_neg, all_negative) except Exception: continue - + return all_positive, all_negative +def _resolve_sink_sampler_node(nodes_by_id: Dict[str, Any], sink_id: str) -> Optional[Dict[str, Any]]: + sampler_id, _ = _select_primary_sampler(nodes_by_id, sink_id) + if not sampler_id: + sampler_id, _ = _select_advanced_sampler(nodes_by_id, sink_id) + if not sampler_id: + return None + sampler_node = nodes_by_id.get(sampler_id) + return sampler_node if isinstance(sampler_node, dict) else None + + +def _collect_prompt_branch_from_input( + nodes_by_id: Dict[str, Any], + input_value: Any, + branch: str, + seen: Set[str], + out: List[str], +) -> None: + if not _is_link(input_value): + return + items = _collect_texts_from_conditioning(nodes_by_id, input_value, branch=branch) + for text, _ in items: + stripped = text.strip() + if not stripped or stripped in seen: + continue + seen.add(stripped) + out.append(stripped) + + +def _workflow_sink_suffix(sink_type: str) -> str: + is_audio_out = ("audio" in sink_type) and ("image" not in sink_type) and ("video" not in sink_type) + is_video_out = (not is_audio_out) and ("video" in sink_type or "animate" in sink_type or "gif" in sink_type) + return "A" if is_audio_out else ("V" if is_video_out else "I") + + +def _scan_graph_input_kinds(nodes_by_id: Dict[str, Any]) -> Tuple[bool, bool, bool]: + has_image_input = False + has_video_input = False + has_audio_input = False + + for node in nodes_by_id.values(): + if not isinstance(node, dict): + continue + ct = _lower(_node_type(node)) + + if "vaeencode" in ct: + pixel_kind = _vae_pixel_input_kind(nodes_by_id, node) + if pixel_kind == "video": + has_video_input = True + elif pixel_kind == "image": + has_image_input = True + + if _is_image_loader_node_type(ct): + has_image_input = True + if _is_video_loader_node_type(ct): + has_video_input = True + if _is_audio_loader_node_type(ct): + has_audio_input = True + + return has_image_input, has_video_input, has_audio_input + + +def _is_image_loader_node_type(ct: str) -> bool: + return "loadimage" in ct or "imageloader" in ct + + +def _is_video_loader_node_type(ct: str) -> bool: + return "loadvideo" in ct or "videoloader" in ct + + +def _is_audio_loader_node_type(ct: str) -> bool: + return "loadaudio" in ct or "audioloader" in ct or ("load" in ct and "audio" in ct) + + +def _vae_pixel_input_kind(nodes_by_id: Dict[str, Any], vae_node: Dict[str, Any]) -> Optional[str]: + vae_ins = _inputs(vae_node) + pixel_link = vae_ins.get("pixels") or vae_ins.get("image") + if not _is_link(pixel_link): + return None + pix_src_id = _walk_passthrough(nodes_by_id, pixel_link) + if not pix_src_id: + return None + pct = _lower(_node_type(nodes_by_id.get(pix_src_id))) + if _is_video_loader_node_type(pct): + return "video" + if _is_image_loader_node_type(pct): + return "image" + return None + + +def _next_latent_upstream_id(nodes_by_id: Dict[str, Any], node: Dict[str, Any]) -> Optional[str]: + for key in ("samples", "latent", "latent_image"): + value = _inputs(node).get(key) + if _is_link(value): + return _walk_passthrough(nodes_by_id, value) + return None + + +def _trace_sampler_latent_input_kind( + nodes_by_id: Dict[str, Any], + sampler_id: Optional[str], + has_image_input: bool, + has_video_input: bool, +) -> Tuple[bool, bool]: + if not sampler_id: + return has_image_input, has_video_input + sampler = nodes_by_id.get(sampler_id) + if not isinstance(sampler, dict): + return has_image_input, has_video_input + + ins = _inputs(sampler) + latent_link = ins.get("latent_image") or ins.get("samples") or ins.get("latent") + if not _is_link(latent_link): + return has_image_input, has_video_input + + curr_id = _walk_passthrough(nodes_by_id, latent_link) + hops = 0 + while curr_id and hops < 15: + hops += 1 + node = nodes_by_id.get(curr_id) + if not isinstance(node, dict): + break + has_image_input, has_video_input, should_stop = _update_inputs_from_latent_node( + nodes_by_id, + node, + has_image_input=has_image_input, + has_video_input=has_video_input, + ) + if should_stop: + break + + curr_id = _next_latent_upstream_id(nodes_by_id, node) + + return has_image_input, has_video_input + + +def _update_inputs_from_latent_node( + nodes_by_id: Dict[str, Any], + node: Dict[str, Any], + *, + has_image_input: bool, + has_video_input: bool, +) -> Tuple[bool, bool, bool]: + ct = _lower(_node_type(node)) + if "emptylatent" in ct: + return has_image_input, has_video_input, True + if "vaeencode" in ct: + pixel_kind = _vae_pixel_input_kind(nodes_by_id, node) + if pixel_kind == "video": + has_video_input = True + else: + has_image_input = True + return has_image_input, has_video_input, True + if "loadlatent" in ct: + return True, has_video_input, True + return has_image_input, has_video_input, False + + +def _workflow_input_prefix(has_image_input: bool, has_video_input: bool, has_audio_input: bool) -> str: + if has_audio_input: + return "A" + if has_video_input: + return "V" + if has_image_input: + return "I" + return "T" + + def _determine_workflow_type(nodes_by_id: Dict[str, Any], sink_node_id: str, sampler_id: Optional[str]) -> str: """ Classify: T2I/I2I/T2V/I2V/V2V and audio variants (T2A/A2A) - + Improved detection: - Checks latent path (EmptyLatent vs VAEEncode) - Also scans for LoadImage/LoadVideo nodes feeding into VAEEncode anywhere in the graph (for reference-based workflows like Flux2 Redux, IP-Adapter, ControlNet, etc.) """ - # 1. Output Type sink = nodes_by_id.get(sink_node_id) - sink_type = _lower(_node_type(sink)) - # Heuristic: SaveVideo -> V, SaveAudio -> A, SaveImage -> I - is_audio_out = ("audio" in sink_type) and ("image" not in sink_type) and ("video" not in sink_type) - is_video_out = (not is_audio_out) and ("video" in sink_type or "animate" in sink_type or "gif" in sink_type) - suffix = "A" if is_audio_out else ("V" if is_video_out else "I") - - # 2. Input Type (Trace Latent Origin) - prefix = "T" # Default to Text - - # 2a. First, scan the entire graph for LoadImage/LoadVideo -> VAEEncode patterns - # This catches reference-based workflows (Flux2 Redux, IP-Adapter, ControlNet, etc.) - has_image_input = False - has_video_input = False - has_audio_input = False - - for nid, node in nodes_by_id.items(): - if not isinstance(node, dict): - continue - ct = _lower(_node_type(node)) - - # Check for VAEEncode nodes and trace their pixel input - if "vaeencode" in ct: - vae_ins = _inputs(node) - pixel_link = vae_ins.get("pixels") or vae_ins.get("image") - if _is_link(pixel_link): - # Trace pixel source - pix_src_id = _walk_passthrough(nodes_by_id, pixel_link) - if pix_src_id: - pix_node = nodes_by_id.get(pix_src_id) - pct = _lower(_node_type(pix_node)) - if "loadvideo" in pct or "videoloader" in pct: - has_video_input = True - elif "loadimage" in pct or "imageloader" in pct: - has_image_input = True - - # Also check direct LoadImage/LoadVideo that feed into conditioning nodes - # (IP-Adapter, ControlNet, etc. may not use VAEEncode) - if "loadimage" in ct or "imageloader" in ct: - # Check if this node's output goes to something other than preview - has_image_input = True - if "loadvideo" in ct or "videoloader" in ct: - has_video_input = True - if "loadaudio" in ct or "audioloader" in ct or ("load" in ct and "audio" in ct): - has_audio_input = True - - # 2b. Now check the main latent path for EmptyLatent (which would indicate pure T2X) - if sampler_id: - sampler = nodes_by_id.get(sampler_id) - if isinstance(sampler, dict): - ins = _inputs(sampler) - # Standard KSampler 'latent_image', or 'samples' input - latent_link = ins.get("latent_image") or ins.get("samples") or ins.get("latent") - - if _is_link(latent_link): - # Trace upstream to find if it comes from EmptyLatent (T2...) or VAEEncode (I2...) - curr_id = _walk_passthrough(nodes_by_id, latent_link) - hops = 0 - while curr_id and hops < 15: - hops += 1 - node = nodes_by_id.get(curr_id) - if not isinstance(node, dict): break - ct = _lower(_node_type(node)) - - if "emptylatent" in ct: - break - if "vaeencode" in ct: - # Found encoder -> I2 or V2 (main latent path) - # Check what feeds the VAE - vae_ins = _inputs(node) - pixel_link = vae_ins.get("pixels") or vae_ins.get("image") - if _is_link(pixel_link): - # Trace pixel source to see if it's "LoadVideo" or "LoadImage" - pix_src_id = _walk_passthrough(nodes_by_id, pixel_link) - if pix_src_id: - pix_node = nodes_by_id.get(pix_src_id) - pct = _lower(_node_type(pix_node)) - if "loadvideo" in pct or "videoloader" in pct: - has_video_input = True - else: - has_image_input = True - else: - has_image_input = True - else: - has_image_input = True - break - - # Input is latent directly (LoadLatent) - if "loadlatent" in ct: - has_image_input = True # Assume loaded latent came from image - break - - # Pass-through logic (LatentUpscale, Duplicate, etc.) - found_next = False - for k in ("samples", "latent", "latent_image"): - if _is_link(_inputs(node).get(k)): - curr_id = _walk_passthrough(nodes_by_id, _inputs(node).get(k)) - found_next = True - break - if not found_next: - break - - # 3. Determine prefix based on findings - # Priority: A (audio) > V (video) > I (image) > T (text-only) - if has_audio_input: - prefix = "A" - elif has_video_input: - prefix = "V" - elif has_image_input: - prefix = "I" - else: - prefix = "T" - + suffix = _workflow_sink_suffix(_lower(_node_type(sink))) + has_image_input, has_video_input, has_audio_input = _scan_graph_input_kinds(nodes_by_id) + has_image_input, has_video_input = _trace_sampler_latent_input_kind( + nodes_by_id, + sampler_id, + has_image_input, + has_video_input, + ) + prefix = _workflow_input_prefix(has_image_input, has_video_input, has_audio_input) return f"{prefix}2{suffix}" @@ -1747,14 +2162,10 @@ def parse_geninfo_from_prompt(prompt_graph: Any, workflow: Any = None) -> Result try: nodes_by_id = _normalize_graph_input(prompt_graph, workflow) except ValueError: - if workflow_meta: - return Result.Ok({"metadata": workflow_meta}) - return Result.Ok(None) + return _geninfo_metadata_only_result(workflow_meta) if not nodes_by_id: - if workflow_meta: - return Result.Ok({"metadata": workflow_meta}) - return Result.Ok(None) + return _geninfo_metadata_only_result(workflow_meta) # Find sink nodes try: @@ -1764,9 +2175,7 @@ def parse_geninfo_from_prompt(prompt_graph: Any, workflow: Any = None) -> Result sinks = [] if not sinks: - if workflow_meta: - return Result.Ok({"metadata": workflow_meta}) - return Result.Ok(None) + return _geninfo_metadata_only_result(workflow_meta) # Prefer "real" sinks (SaveVideo over PreviewImage, etc.) try: @@ -1774,29 +2183,139 @@ def parse_geninfo_from_prompt(prompt_graph: Any, workflow: Any = None) -> Result sink_id = sinks[0] sampler_id, sampler_conf = _select_primary_sampler(nodes_by_id, sink_id) - + # ... logic continues ... # Since I'm refactoring the monolithic block, I'll call a helper for the core logic # OR I can inline the rest if it's manageable. The main issue was the huge try/except. # Let's delegate the rest of extraction to a helper or just protect the extraction part. - + return _extract_geninfo(nodes_by_id, sinks, workflow_meta) except (ValueError, TypeError, KeyError, AttributeError) as e: logger.warning(f"GenInfo parsing failed: {e}") - if workflow_meta: - return Result.Ok({"metadata": workflow_meta}) - return Result.Ok(None) + return _geninfo_metadata_only_result(workflow_meta) except Exception as e: # Unexpected error - log with full traceback for debugging logger.exception(f"Unexpected error in GenInfo parsing: {e}") - if workflow_meta: - return Result.Ok({"metadata": workflow_meta}) - return Result.Ok(None) + return _geninfo_metadata_only_result(workflow_meta) + + +def _geninfo_metadata_only_result(workflow_meta: Optional[Dict[str, Any]]) -> Result[Optional[Dict[str, Any]]]: + if workflow_meta: + return Result.Ok({"metadata": workflow_meta}) + return Result.Ok(None) + +def _resolve_graph_target(prompt_graph: Any, workflow: Any) -> Optional[Dict[str, Any]]: + if isinstance(prompt_graph, dict) and prompt_graph: + return prompt_graph + if isinstance(workflow, dict) and "nodes" in workflow: + return workflow + return None + + +def _build_link_source_map(links: Any) -> Dict[int, Tuple[int, int]]: + link_to_source: Dict[int, Tuple[int, int]] = {} + if not isinstance(links, list): + return link_to_source + for link in links: + if isinstance(link, list) and len(link) >= 3: + link_id, src_node, src_slot = link[0], link[1], link[2] + link_to_source[link_id] = (src_node, src_slot) + return link_to_source + + +def _convert_litegraph_node( + node: Dict[str, Any], + link_to_source: Dict[int, Tuple[int, int]], +) -> Dict[str, Any]: + converted = _init_litegraph_converted_node(node) + raw_inputs = node.get("inputs", []) + widgets_values = node.get("widgets_values", []) + widgets_list = widgets_values if isinstance(widgets_values, list) else [] + converted_inputs = _populate_converted_inputs(converted, raw_inputs, widgets_list, link_to_source) + _merge_widget_dict_inputs(converted_inputs, widgets_values) + _set_text_fallback_from_widgets(converted_inputs, widgets_list, node) + return converted + + +def _init_litegraph_converted_node(node: Dict[str, Any]) -> Dict[str, Any]: + return { + "class_type": node.get("type"), + "type": node.get("type"), + "id": node.get("id"), + "inputs": {}, + "widgets_values": node.get("widgets_values"), + "outputs": node.get("outputs"), + "properties": node.get("properties"), + "title": node.get("title"), + "mode": node.get("mode"), + } + + +def _populate_converted_inputs( + converted: Dict[str, Any], + raw_inputs: Any, + widgets_list: List[Any], + link_to_source: Dict[int, Tuple[int, int]], +) -> Dict[str, Any]: + converted_inputs = converted.get("inputs") + if isinstance(raw_inputs, list) and isinstance(converted_inputs, dict): + _populate_converted_inputs_from_list(converted_inputs, raw_inputs, widgets_list, link_to_source) + return converted_inputs + if isinstance(raw_inputs, dict): + converted["inputs"] = cast(Dict[str, Any], raw_inputs) + converted_inputs = converted.get("inputs") + return cast(Dict[str, Any], converted_inputs or {}) + + +def _populate_converted_inputs_from_list( + converted_inputs: Dict[str, Any], + raw_inputs: List[Any], + widgets_list: List[Any], + link_to_source: Dict[int, Tuple[int, int]], +) -> None: + widget_idx = 0 + for inp in raw_inputs: + if not isinstance(inp, dict): + continue + name = inp.get("name") + if not name: + continue + link_id = inp.get("link") + if link_id is not None and link_id in link_to_source: + src_node_id, src_slot = link_to_source[link_id] + converted_inputs[name] = [str(src_node_id), src_slot] + continue + if "widget" in inp: + if widget_idx < len(widgets_list): + converted_inputs[name] = widgets_list[widget_idx] + widget_idx += 1 + + +def _merge_widget_dict_inputs(converted_inputs: Dict[str, Any], widgets_values: Any) -> None: + if not isinstance(widgets_values, dict): + return + for key, value in widgets_values.items(): + if key not in converted_inputs: + converted_inputs[key] = value + + +def _set_text_fallback_from_widgets(converted_inputs: Dict[str, Any], widgets_list: List[Any], node: Dict[str, Any]) -> None: + if not widgets_list or "text" in converted_inputs: + return + node_type_lower = _lower(node.get("type")) + if not any(token in node_type_lower for token in ("primitive", "string", "text", "encode")): + return + for widget_value in widgets_list: + if isinstance(widget_value, str) and len(widget_value.strip()) > 10: + converted_inputs["text"] = widget_value + converted_inputs["value"] = widget_value + return + def _normalize_graph_input(prompt_graph: Any, workflow: Any) -> Optional[Dict[str, Dict[str, Any]]]: """Normalize prompt graph or workflow into nodes_by_id dict. - + Handles both ComfyUI prompt-graph format and LiteGraph workflow format. LiteGraph format has: - nodes: list of {id, type, inputs: [{name, link}], widgets_values: [...]} @@ -1804,615 +2323,854 @@ def _normalize_graph_input(prompt_graph: Any, workflow: Any) -> Optional[Dict[st Prompt-graph format has: - dict with node_id as key: {class_type, inputs: {name: value}} """ - target_graph = prompt_graph - - if not isinstance(target_graph, dict) or not target_graph: - if isinstance(workflow, dict) and "nodes" in workflow: - target_graph = workflow - else: - return None + target_graph = _resolve_graph_target(prompt_graph, workflow) + if not isinstance(target_graph, dict): + return None nodes_by_id: Dict[str, Dict[str, Any]] = {} if "nodes" in target_graph and isinstance(target_graph["nodes"], list): - # LiteGraph format - need to convert inputs list to dict and apply widgets_values - links = target_graph.get("links", []) - - # Build link map: link_id -> (src_node_id, src_slot) - link_to_source: Dict[int, Tuple[int, int]] = {} - for link in links: - if isinstance(link, list) and len(link) >= 3: - link_id, src_node, src_slot = link[0], link[1], link[2] - link_to_source[link_id] = (src_node, src_slot) - + link_to_source = _build_link_source_map(target_graph.get("links", [])) for node in target_graph["nodes"]: - if isinstance(node, dict): - node_id = str(node.get("id")) - # Convert LiteGraph node to prompt-graph-like format - converted: Dict[str, Any] = { - "class_type": node.get("type"), - "type": node.get("type"), # Keep both for compatibility - "id": node.get("id"), - "inputs": {}, - # Preserve original fields for other parsers - "widgets_values": node.get("widgets_values"), - "outputs": node.get("outputs"), - "properties": node.get("properties"), - "title": node.get("title"), - "mode": node.get("mode"), - } - - raw_inputs = node.get("inputs", []) - widgets_values = node.get("widgets_values", []) - - # Build inputs dict from: - # 1. Links (for connected inputs) -> store as ["node_id", slot] - # 2. widgets_values (for widget inputs) - # Note: widgets_values can be a list or a dict (e.g., VHS nodes use dict) - widgets_list = widgets_values if isinstance(widgets_values, list) else [] - - if isinstance(raw_inputs, list): - widget_idx = 0 - for inp in raw_inputs: - if isinstance(inp, dict): - name = inp.get("name") - if not name: - continue - link_id = inp.get("link") - if link_id is not None and link_id in link_to_source: - # Connected input - store as [source_node_id, source_slot] - src_node_id, src_slot = link_to_source[link_id] - converted["inputs"][name] = [str(src_node_id), src_slot] - elif "widget" in inp: - # Widget input - try to get value from widgets_values - # LiteGraph stores widget values in order they appear - if widget_idx < len(widgets_list): - converted["inputs"][name] = widgets_list[widget_idx] - widget_idx += 1 - elif isinstance(raw_inputs, dict): - # Already in dict format (shouldn't happen for LiteGraph but handle it) - converted["inputs"] = cast(Dict[str, Any], raw_inputs) - - # Handle widgets_values as dict (VHS and some other nodes) - if isinstance(widgets_values, dict): - for k, v in widgets_values.items(): - converted_inputs = converted.get("inputs") - if isinstance(converted_inputs, dict) and k not in converted_inputs: - converted_inputs[k] = v - - # Also copy widgets_values to inputs if they have meaningful names - # This helps with nodes that have prompts in widgets but no widget definition - if widgets_list: - # For text encoder nodes, widgets_values[0] is usually the text - node_type_lower = _lower(node.get("type")) - converted_inputs = converted.get("inputs") - if isinstance(converted_inputs, dict) and "text" not in converted_inputs: - if any(x in node_type_lower for x in ["primitive", "string", "text", "encode"]): - for wv in widgets_list: - if isinstance(wv, str) and len(wv.strip()) > 10: - converted_inputs["text"] = wv - converted_inputs["value"] = wv - break - - nodes_by_id[node_id] = converted + if not isinstance(node, dict): + continue + node_id = str(node.get("id")) + nodes_by_id[node_id] = _convert_litegraph_node(node, link_to_source) else: - # Prompt-graph format - use as-is - for k, v in target_graph.items(): - if isinstance(v, dict): - nodes_by_id[str(k)] = v + for key, value in target_graph.items(): + if isinstance(value, dict): + nodes_by_id[str(key)] = value return nodes_by_id if nodes_by_id else None -def _extract_geninfo(nodes_by_id: Dict[str, Any], sinks: List[str], workflow_meta: Optional[Dict]) -> Result: - def _extract_tts_geninfo_fallback() -> Optional[Dict[str, Any]]: - text_node_id = None - text_node = None - engine_node_id = None - engine_node = None - for nid, node in nodes_by_id.items(): - if not isinstance(node, dict): - continue - ct = _lower(_node_type(node)) - if text_node is None and ("unifiedttstextnode" in ct or "tts_text_node" in ct or ("tts" in ct and "text" in ct)): - text_node_id = nid - text_node = node - if engine_node is None and ( - "ttsengine" in ct - or ("qwen" in ct and "tts" in ct and "engine" in ct) - or "engine_node" in ct and "tts" in ct - ): - engine_node_id = nid - engine_node = node - if text_node is not None and engine_node is not None: +def _set_named_field(out: Dict[str, Any], key: str, value: Any, source: str, confidence: str = "high") -> None: + if value is None: + return + name = str(value).strip() + if not name: + return + out[key] = {"name": name, "confidence": confidence, "source": source} + + +def _set_value_field(out: Dict[str, Any], key: str, value: Any, source: str, confidence: str = "high") -> None: + if value is None: + return + out[key] = {"value": value, "confidence": confidence, "source": source} + + +def _find_tts_nodes(nodes_by_id: Dict[str, Any]) -> Tuple[Optional[str], Optional[Dict[str, Any]], Optional[str], Optional[Dict[str, Any]]]: + text_node_id: Optional[str] = None + text_node: Optional[Dict[str, Any]] = None + engine_node_id: Optional[str] = None + engine_node: Optional[Dict[str, Any]] = None + + for nid, node in nodes_by_id.items(): + if not isinstance(node, dict): + continue + ct = _lower(_node_type(node)) + if text_node is None and _is_tts_text_node_type(ct): + text_node_id = str(nid) + text_node = node + if engine_node is None and _is_tts_engine_node_type(ct): + engine_node_id = str(nid) + engine_node = node + if text_node is not None and engine_node is not None: + break + + return text_node_id, text_node, engine_node_id, engine_node + + +def _is_tts_text_node_type(ct: str) -> bool: + return "unifiedttstextnode" in ct or "tts_text_node" in ct or ("tts" in ct and "text" in ct) + + +def _is_tts_engine_node_type(ct: str) -> bool: + return "ttsengine" in ct or ("qwen" in ct and "tts" in ct and "engine" in ct) or ("engine_node" in ct and "tts" in ct) + + +def _apply_tts_text_node_fields( + out: Dict[str, Any], + nodes_by_id: Dict[str, Any], + text_node_id: str, + text_node: Dict[str, Any], +) -> None: + tins = _inputs(text_node) + source = f"{_node_type(text_node)}:{text_node_id}" + out["sampler"] = {"name": str(_node_type(text_node) or "TTS"), "confidence": "high", "source": source} + + _apply_tts_text_direct_fields(out, tins, source) + _apply_tts_text_widget_fallback(out, text_node.get("widgets_values"), source) + _apply_tts_narrator_link_voice(out, nodes_by_id, tins) + + +def _apply_tts_text_direct_fields(out: Dict[str, Any], tins: Dict[str, Any], source: str) -> None: + text_value = tins.get("text") + if isinstance(text_value, str) and text_value.strip(): + out["positive"] = {"value": text_value.strip(), "confidence": "high", "source": source} + _set_value_field(out, "seed", _scalar(tins.get("seed")) or _scalar(tins.get("noise_seed")), source) + _set_named_field(out, "voice", _scalar(tins.get("narrator_voice")), source) + for key in ( + "enable_chunking", + "max_chars_per_chunk", + "chunk_combination_method", + "silence_between_chunks_ms", + "enable_audio_cache", + "batch_size", + "control_after_generate", + ): + _set_value_field(out, key, _scalar(tins.get(key)), source) + + +def _apply_tts_text_widget_fallback(out: Dict[str, Any], widgets: Any, source: str) -> None: + if not isinstance(widgets, list): + return + if "positive" not in out: + text = _first_long_widget_text(widgets) + if text: + out["positive"] = {"value": text, "confidence": "medium", "source": source} + if "seed" not in out: + seed = _first_nonnegative_int_scalar(widgets) + if seed is not None: + out["seed"] = {"value": seed, "confidence": "medium", "source": source} + if "voice" not in out: + voice_name = _widget_voice_name(widgets) + if voice_name: + out["voice"] = {"name": voice_name, "confidence": "medium", "source": source} + + +def _first_long_widget_text(widgets: List[Any]) -> Optional[str]: + for value in widgets: + if not isinstance(value, str): + continue + stripped = value.strip() + if len(stripped) > 20: + return stripped + return None + + +def _first_nonnegative_int_scalar(widgets: List[Any]) -> Optional[int]: + for value in widgets: + scalar_value = _scalar(value) + if isinstance(scalar_value, int) and scalar_value >= 0: + return scalar_value + return None + + +def _widget_voice_name(widgets: List[Any]) -> Optional[str]: + if len(widgets) < 2: + return None + voice = widgets[1] + if isinstance(voice, str) and voice.strip(): + return voice.strip() + return None + + +def _apply_tts_narrator_link_voice(out: Dict[str, Any], nodes_by_id: Dict[str, Any], tins: Dict[str, Any]) -> None: + narrator_link = tins.get("opt_narrator") + narrator_id = _walk_passthrough(nodes_by_id, narrator_link) if _is_link(narrator_link) else None + narrator_node = nodes_by_id.get(narrator_id) if narrator_id else None + if not isinstance(narrator_node, dict): + return + voice_name = _scalar(_inputs(narrator_node).get("voice_name")) + if voice_name is None or not str(voice_name).strip(): + return + out["voice"] = { + "name": str(voice_name).strip(), + "confidence": "high", + "source": f"{_node_type(narrator_node)}:{narrator_id}", + } + + +def _apply_tts_engine_node_fields(out: Dict[str, Any], engine_node_id: str, engine_node: Dict[str, Any]) -> None: + eins = _inputs(engine_node) + source = f"{_node_type(engine_node)}:{engine_node_id}" + _apply_tts_engine_direct_fields(out, eins, source) + _apply_tts_engine_widget_fields(out, engine_node, source) + + +def _apply_tts_engine_direct_fields(out: Dict[str, Any], eins: Dict[str, Any], source: str) -> None: + model_name = _clean_model_id(eins.get("model_size") or eins.get("model") or eins.get("checkpoint") or eins.get("model_name")) + if model_name: + out["checkpoint"] = {"name": model_name, "confidence": "high", "source": source} + out["models"] = {"checkpoint": out["checkpoint"]} + for key in ("device", "voice_preset", "instruct", "language"): + _set_value_field(out, key, _scalar(eins.get(key)), source) + for key in ( + "temperature", + "top_p", + "top_k", + "repetition_penalty", + "max_new_tokens", + "dtype", + "attn_implementation", + "x_vector_only_mode", + "use_torch_compile", + "use_cuda_graphs", + "compile_mode", + ): + _set_value_field(out, key, _scalar(eins.get(key)), source) + + +def _apply_tts_engine_widget_fields(out: Dict[str, Any], engine_node: Dict[str, Any], source: str) -> None: + ewidgets = engine_node.get("widgets_values") + if not isinstance(ewidgets, list): + return + ect = _lower(_node_type(engine_node)) + _apply_tts_engine_widget_checkpoint(out, ewidgets, source) + _apply_tts_engine_widget_language(out, ewidgets, ect, source) + if "qwen3ttsengine" in ect: + _apply_tts_engine_qwen_widgets(out, ewidgets, source) + + +def _apply_tts_engine_widget_checkpoint(out: Dict[str, Any], ewidgets: List[Any], source: str) -> None: + if "checkpoint" in out or not ewidgets: + return + guess_model = _clean_model_id(ewidgets[0]) + if not guess_model: + return + out["checkpoint"] = {"name": guess_model, "confidence": "medium", "source": source} + out["models"] = {"checkpoint": out["checkpoint"]} + + +def _apply_tts_engine_widget_language(out: Dict[str, Any], ewidgets: List[Any], ect: str, source: str) -> None: + if "language" in out: + return + guess_lang = _scalar(ewidgets[3]) if "qwen3ttsengine" in ect and len(ewidgets) > 3 else None + if guess_lang is None: + for value in ewidgets: + if isinstance(value, str) and value.strip() and value.strip().lower() not in ("auto", "default"): + guess_lang = value break + if guess_lang is not None: + out["language"] = {"value": str(guess_lang).strip(), "confidence": "medium", "source": source} - if not text_node and not engine_node: - return None - out: Dict[str, Any] = { - "engine": { - "parser_version": "geninfo-tts-v1", - "type": "tts", - "sink": "audio", - } +def _apply_tts_engine_qwen_widgets(out: Dict[str, Any], ewidgets: List[Any], source: str) -> None: + qwen_widget_indices = { + "device": 1, + "voice_preset": 2, + "top_k": 5, + "top_p": 6, + "temperature": 7, + "repetition_penalty": 8, + "max_new_tokens": 9, + } + for key, idx in qwen_widget_indices.items(): + if key in out or len(ewidgets) <= idx: + continue + value = _scalar(ewidgets[idx]) + if value is None: + continue + out[key] = { + "value": value if key not in ("device", "voice_preset") else str(value).strip(), + "confidence": "medium", + "source": source, } - if workflow_meta: - out["metadata"] = workflow_meta - if text_node: - tins = _inputs(text_node) - source = f"{_node_type(text_node)}:{text_node_id}" - # In TTS workflows, text node is the generation driver ("sampler"). - out["sampler"] = {"name": str(_node_type(text_node) or "TTS"), "confidence": "high", "source": source} - text_value = tins.get("text") - if isinstance(text_value, str) and text_value.strip(): - out["positive"] = {"value": text_value.strip(), "confidence": "high", "source": source} - seed_value = _scalar(tins.get("seed")) or _scalar(tins.get("noise_seed")) - if seed_value is not None: - out["seed"] = {"value": seed_value, "confidence": "high", "source": source} - narrator = _scalar(tins.get("narrator_voice")) - if narrator is not None and str(narrator).strip(): - out["voice"] = {"name": str(narrator).strip(), "confidence": "high", "source": source} - for key in ( - "enable_chunking", - "max_chars_per_chunk", - "chunk_combination_method", - "silence_between_chunks_ms", - "enable_audio_cache", - "batch_size", - "control_after_generate", - ): - value = _scalar(tins.get(key)) - if value is not None: - out[key] = {"value": value, "confidence": "high", "source": source} - # Workflow-export fallback: infer key fields from widgets_values. - widgets = text_node.get("widgets_values") - if isinstance(widgets, list): - if "positive" not in out: - for v in widgets: - if isinstance(v, str) and len(v.strip()) > 20: - out["positive"] = {"value": v.strip(), "confidence": "medium", "source": source} - break - if "seed" not in out: - for v in widgets: - sv = _scalar(v) - if isinstance(sv, int) and sv >= 0: - out["seed"] = {"value": sv, "confidence": "medium", "source": source} - break - if "voice" not in out and len(widgets) >= 2 and isinstance(widgets[1], str) and widgets[1].strip(): - out["voice"] = {"name": widgets[1].strip(), "confidence": "medium", "source": source} - # If narrator comes from CharacterVoices via opt_narrator link, resolve voice_name. - try: - narrator_link = tins.get("opt_narrator") - narrator_id = _walk_passthrough(nodes_by_id, narrator_link) if _is_link(narrator_link) else None - narrator_node = nodes_by_id.get(narrator_id) if narrator_id else None - if isinstance(narrator_node, dict): - nins = _inputs(narrator_node) - voice_name = _scalar(nins.get("voice_name")) - if voice_name is not None and str(voice_name).strip(): - out["voice"] = { - "name": str(voice_name).strip(), - "confidence": "high", - "source": f"{_node_type(narrator_node)}:{narrator_id}", - } - except Exception: - pass - if engine_node: - eins = _inputs(engine_node) - esource = f"{_node_type(engine_node)}:{engine_node_id}" - model_name = _clean_model_id(eins.get("model_size") or eins.get("model") or eins.get("checkpoint") or eins.get("model_name")) - if model_name: - out["checkpoint"] = {"name": model_name, "confidence": "high", "source": esource} - out["models"] = {"checkpoint": out["checkpoint"]} - device = _scalar(eins.get("device")) - if device is not None and str(device).strip(): - out["device"] = {"value": str(device).strip(), "confidence": "high", "source": esource} - voice_preset = _scalar(eins.get("voice_preset")) - if voice_preset is not None and str(voice_preset).strip(): - out["voice_preset"] = {"value": str(voice_preset).strip(), "confidence": "high", "source": esource} - instruct = _scalar(eins.get("instruct")) - if instruct is not None and str(instruct).strip(): - out["instruct"] = {"value": str(instruct).strip(), "confidence": "high", "source": esource} - language = _scalar(eins.get("language")) - if language is not None and str(language).strip(): - out["language"] = {"value": str(language).strip(), "confidence": "high", "source": esource} - for key in ( - "temperature", - "top_p", - "top_k", - "repetition_penalty", - "max_new_tokens", - "dtype", - "attn_implementation", - "x_vector_only_mode", - "use_torch_compile", - "use_cuda_graphs", - "compile_mode", - ): - value = _scalar(eins.get(key)) - if value is not None: - out[key] = {"value": value, "confidence": "high", "source": esource} - # Workflow-export fallback for engine widgets. - ewidgets = engine_node.get("widgets_values") - ect = _lower(_node_type(engine_node)) - if isinstance(ewidgets, list): - if "checkpoint" not in out and ewidgets: - guess_model = _clean_model_id(ewidgets[0]) - if guess_model: - out["checkpoint"] = {"name": guess_model, "confidence": "medium", "source": esource} - out["models"] = {"checkpoint": out["checkpoint"]} - if "language" not in out: - guess_lang = None - if "qwen3ttsengine" in ect and len(ewidgets) > 3: - guess_lang = _scalar(ewidgets[3]) - if guess_lang is None: - for v in ewidgets: - if isinstance(v, str) and v.strip() and v.strip().lower() not in ("auto", "default"): - guess_lang = v - break - if guess_lang is not None: - out["language"] = {"value": str(guess_lang).strip(), "confidence": "medium", "source": esource} - if "device" not in out and "qwen3ttsengine" in ect and len(ewidgets) > 1: - v = _scalar(ewidgets[1]) - if v is not None and str(v).strip(): - out["device"] = {"value": str(v).strip(), "confidence": "medium", "source": esource} - if "voice_preset" not in out and "qwen3ttsengine" in ect and len(ewidgets) > 2: - v = _scalar(ewidgets[2]) - if v is not None and str(v).strip(): - out["voice_preset"] = {"value": str(v).strip(), "confidence": "medium", "source": esource} - if "top_k" not in out and "qwen3ttsengine" in ect and len(ewidgets) > 5: - v = _scalar(ewidgets[5]) - if v is not None: - out["top_k"] = {"value": v, "confidence": "medium", "source": esource} - if "top_p" not in out and "qwen3ttsengine" in ect and len(ewidgets) > 6: - v = _scalar(ewidgets[6]) - if v is not None: - out["top_p"] = {"value": v, "confidence": "medium", "source": esource} - if "temperature" not in out and "qwen3ttsengine" in ect and len(ewidgets) > 7: - v = _scalar(ewidgets[7]) - if v is not None: - out["temperature"] = {"value": v, "confidence": "medium", "source": esource} - if "repetition_penalty" not in out and "qwen3ttsengine" in ect and len(ewidgets) > 8: - v = _scalar(ewidgets[8]) - if v is not None: - out["repetition_penalty"] = {"value": v, "confidence": "medium", "source": esource} - if "max_new_tokens" not in out and "qwen3ttsengine" in ect and len(ewidgets) > 9: - v = _scalar(ewidgets[9]) - if v is not None: - out["max_new_tokens"] = {"value": v, "confidence": "medium", "source": esource} +def _extract_tts_geninfo_fallback(nodes_by_id: Dict[str, Any], workflow_meta: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]: + text_node_id, text_node, engine_node_id, engine_node = _find_tts_nodes(nodes_by_id) + if not text_node and not engine_node: + return None - input_files = _extract_input_files(nodes_by_id) - if input_files: - out["inputs"] = input_files + out: Dict[str, Any] = { + "engine": { + "parser_version": "geninfo-tts-v1", + "type": "tts", + "sink": "audio", + } + } + if workflow_meta: + out["metadata"] = workflow_meta - if len(out.keys()) <= 1: - return None - return out + _apply_tts_text_node_fields_safe(out, nodes_by_id, text_node_id, text_node) + _apply_tts_engine_node_fields_safe(out, engine_node_id, engine_node) + _apply_tts_input_files(out, nodes_by_id) - sink_id = sinks[0] + if len(out.keys()) <= 1: + return None + return out + + +def _apply_tts_text_node_fields_safe( + out: Dict[str, Any], + nodes_by_id: Dict[str, Any], + text_node_id: Optional[str], + text_node: Any, +) -> None: + try: + if text_node_id and isinstance(text_node, dict): + _apply_tts_text_node_fields(out, nodes_by_id, text_node_id, text_node) + except Exception: + pass + + +def _apply_tts_engine_node_fields_safe( + out: Dict[str, Any], + engine_node_id: Optional[str], + engine_node: Any, +) -> None: + try: + if engine_node_id and isinstance(engine_node, dict): + _apply_tts_engine_node_fields(out, engine_node_id, engine_node) + except Exception: + pass + + +def _apply_tts_input_files(out: Dict[str, Any], nodes_by_id: Dict[str, Any]) -> None: + input_files = _extract_input_files(nodes_by_id) + if input_files: + out["inputs"] = input_files + + +def _find_special_sampler_id(nodes_by_id: Dict[str, Any]) -> Optional[str]: + for nid, node in nodes_by_id.items(): + ct = _lower(_node_type(node)) + if "marigold" in ct: + return str(nid) + if "instruction" in ct and "qwen" in ct: + return str(nid) + return None + + +def _select_sampler_context(nodes_by_id: Dict[str, Any], sink_id: str) -> Tuple[Optional[str], str, str]: sampler_id, sampler_conf = _select_primary_sampler(nodes_by_id, sink_id) sampler_mode = "primary" - + if not sampler_id: sampler_id, sampler_conf = _select_advanced_sampler(nodes_by_id, sink_id) - sampler_mode = "advanced" if sampler_id else sampler_mode - + if sampler_id: + sampler_mode = "advanced" + if not sampler_id: sampler_id, sampler_conf = _select_any_sampler(nodes_by_id) - sampler_mode = "global" if sampler_id else sampler_mode - - # Marigold / Qwen detection - if not sampler_id: - for nid, val in nodes_by_id.items(): - ct_lower = _lower(_node_type(val)) - if "marigold" in ct_lower: - sampler_id = nid - break - if "instruction" in ct_lower and "qwen" in ct_lower: - sampler_id = nid - break + if sampler_id: + sampler_mode = "global" if not sampler_id: - tts_fallback = _extract_tts_geninfo_fallback() - if tts_fallback: - return Result.Ok(tts_fallback) + special_sampler = _find_special_sampler_id(nodes_by_id) + if special_sampler: + sampler_id = special_sampler + sampler_conf = "low" + sampler_mode = "fallback" - out_fallback: Dict[str, Any] = {} - if workflow_meta: - out_fallback["metadata"] = workflow_meta - - input_files_fallback = _extract_input_files(nodes_by_id) - if input_files_fallback: - out_fallback["inputs"] = input_files_fallback - - if out_fallback: - return Result.Ok(out_fallback) - return Result.Ok(None) + return sampler_id, sampler_conf, sampler_mode - # For secondary traces (VAE, etc.), compute a stable upstream set from the sink input. - sink = nodes_by_id.get(sink_id) or {} - sink_link = _pick_sink_inputs(sink) - sink_start_id = _walk_passthrough(nodes_by_id, sink_link) if sink_link else None - sampler_node = nodes_by_id.get(sampler_id) or {} - sampler_source = f"{_node_type(sampler_node)}:{sampler_id}" - confidence = sampler_conf +def _build_no_sampler_result(nodes_by_id: Dict[str, Any], workflow_meta: Optional[Dict[str, Any]]) -> Result: + tts_fallback = _extract_tts_geninfo_fallback(nodes_by_id, workflow_meta) + if tts_fallback: + return Result.Ok(tts_fallback) - ins = _inputs(sampler_node) + out_fallback: Dict[str, Any] = {} + if workflow_meta: + out_fallback["metadata"] = workflow_meta + + input_files = _extract_input_files(nodes_by_id) + if input_files: + out_fallback["inputs"] = input_files + + if out_fallback: + return Result.Ok(out_fallback) + return Result.Ok(None) + + +def _source_from_items(items: List[Tuple[str, str]]) -> Optional[Tuple[str, str]]: + if not items: + return None + text = "\n".join([t for t, _ in items]).strip() + if not text: + return None + sources = [s for _, s in items] + source = sources[0] if len(set(sources)) <= 1 else f"{sources[0]} (+{len(sources)-1})" + return text, source + + +def _extract_prompt_from_conditioning(nodes_by_id: Dict[str, Any], link: Any, branch: Optional[str] = None) -> Optional[Tuple[str, str]]: + if not _is_link(link): + return None + items = _collect_texts_from_conditioning(nodes_by_id, link, branch=branch) + return _source_from_items(items) + + +def _apply_advanced_guider_prompt_trace( + nodes_by_id: Dict[str, Any], + guider_link: Any, + trace: Dict[str, Any], +) -> None: + guider_id = _walk_passthrough(nodes_by_id, guider_link) + guider_node = nodes_by_id.get(guider_id) if guider_id else None + if not isinstance(guider_node, dict): + return + + gins = _inputs(guider_node) + _apply_guider_conditioning_prompt_hints(nodes_by_id, gins, trace) + _apply_guider_positive_prompt_hints(nodes_by_id, gins, trace) + _apply_guider_negative_prompt_hints(nodes_by_id, gins, trace) + + cfg_val = _scalar(gins.get("cfg")) or _scalar(gins.get("cfg_scale")) or _scalar(gins.get("guidance")) + if cfg_val is not None: + trace["guider_cfg_value"] = cfg_val + trace["guider_cfg_source"] = f"{_node_type(guider_node)}:{guider_id}" + else: + found_guidance = _trace_guidance_value(nodes_by_id, gins.get("conditioning")) + if found_guidance: + trace["guider_cfg_value"], trace["guider_cfg_source"] = found_guidance + + if _is_link(gins.get("model")): + trace["guider_model_link"] = gins.get("model") + + +def _apply_guider_conditioning_prompt_hints(nodes_by_id: Dict[str, Any], gins: Dict[str, Any], trace: Dict[str, Any]) -> None: + conditioning = gins.get("conditioning") + if not _is_link(conditioning): + return + trace["conditioning_link"] = conditioning + extracted = _extract_prompt_from_conditioning(nodes_by_id, conditioning) + if extracted and not trace["pos_val"]: + trace["pos_val"] = extracted + + +def _apply_guider_positive_prompt_hints(nodes_by_id: Dict[str, Any], gins: Dict[str, Any], trace: Dict[str, Any]) -> None: + positive = gins.get("positive") + if not _is_link(positive): + return + trace["conditioning_link"] = trace["conditioning_link"] or positive + if trace["pos_val"]: + return + extracted = _extract_prompt_from_conditioning(nodes_by_id, positive, branch="positive") + if extracted: + trace["pos_val"] = extracted + + +def _apply_guider_negative_prompt_hints(nodes_by_id: Dict[str, Any], gins: Dict[str, Any], trace: Dict[str, Any]) -> None: + negative = gins.get("negative") + if trace.get("neg_val") or not _is_link(negative): + return + extracted = _extract_prompt_from_conditioning(nodes_by_id, negative, branch="negative") + if extracted: + trace["neg_val"] = extracted + + +def _extract_prompt_trace( + nodes_by_id: Dict[str, Any], + sampler_node: Dict[str, Any], + sampler_id: str, + ins: Dict[str, Any], + advanced: bool, +) -> Dict[str, Any]: + trace: Dict[str, Any] = { + "pos_val": None, + "neg_val": None, + "conditioning_link": None, + "guider_cfg_value": None, + "guider_cfg_source": None, + "guider_model_link": None, + } + + sampler_ct = _lower(_node_type(sampler_node)) + _apply_direct_sampler_prompt_hints(trace, sampler_node, sampler_id, ins, sampler_ct) + _apply_embed_prompt_hints(trace, nodes_by_id, ins) + _apply_conditioning_prompt_hints(trace, nodes_by_id, ins) + if advanced and _is_link(ins.get("guider")): + _apply_advanced_guider_prompt_trace(nodes_by_id, ins.get("guider"), trace) + _apply_prompt_text_fallback(trace, sampler_node, sampler_id, ins) + + return trace + + +def _apply_direct_sampler_prompt_hints( + trace: Dict[str, Any], + sampler_node: Dict[str, Any], + sampler_id: str, + ins: Dict[str, Any], + sampler_ct: str, +) -> None: + if "instruction" in sampler_ct and "qwen" in sampler_ct: + prompt_value = ins.get("instruction") or ins.get("text") + if isinstance(prompt_value, str) and prompt_value.strip(): + trace["pos_val"] = (prompt_value.strip(), f"{_node_type(sampler_node)}:{sampler_id}:instruction") + if "flux" in sampler_ct and "trainer" in sampler_ct: + prompt_value = ins.get("prompt") + if isinstance(prompt_value, str) and prompt_value.strip(): + trace["pos_val"] = (prompt_value.strip(), f"{_node_type(sampler_node)}:{sampler_id}:prompt") + + +def _apply_embed_prompt_hints(trace: Dict[str, Any], nodes_by_id: Dict[str, Any], ins: Dict[str, Any]) -> None: + if not (_is_link(ins.get("text_embeds")) or _is_link(ins.get("hyvid_embeds"))): + return + embeds_link = ins.get("text_embeds") or ins.get("hyvid_embeds") + pos_embed, neg_embed = _extract_posneg_from_text_embeds(nodes_by_id, embeds_link) + trace["pos_val"] = trace["pos_val"] or pos_embed + trace["neg_val"] = trace["neg_val"] or neg_embed + + +def _apply_conditioning_prompt_hints(trace: Dict[str, Any], nodes_by_id: Dict[str, Any], ins: Dict[str, Any]) -> None: + if _is_link(ins.get("positive")): + extracted = _extract_prompt_from_conditioning(nodes_by_id, ins.get("positive"), branch="positive") + if extracted: + trace["pos_val"] = extracted + trace["conditioning_link"] = ins.get("positive") + if _is_link(ins.get("negative")): + extracted = _extract_prompt_from_conditioning(nodes_by_id, ins.get("negative"), branch="negative") + if extracted: + trace["neg_val"] = extracted + + +def _apply_prompt_text_fallback(trace: Dict[str, Any], sampler_node: Dict[str, Any], sampler_id: str, ins: Dict[str, Any]) -> None: + if not trace.get("pos_val"): + trace["pos_val"] = _first_prompt_field( + ins, + ("positive_prompt", "prompt", "positive", "text", "text_g", "text_l"), + f"{_node_type(sampler_node)}:{sampler_id}", + ) + if not trace.get("neg_val"): + trace["neg_val"] = _first_prompt_field( + ins, + ("negative_prompt", "negative", "neg", "text_negative"), + f"{_node_type(sampler_node)}:{sampler_id}", + ) + + +def _first_prompt_field(ins: Dict[str, Any], keys: Tuple[str, ...], source_prefix: str) -> Optional[Tuple[str, str]]: + for key in keys: + value = ins.get(key) + if _looks_like_prompt_string(value): + return str(value).strip(), f"{source_prefix}:{key}" + return None + + +def _init_sampler_values(nodes_by_id: Dict[str, Any], sampler_node: Dict[str, Any], ins: Dict[str, Any]) -> Dict[str, Any]: + sampler_name = _sampler_name_from_inputs(sampler_node, ins) + seed_val = _seed_value_from_inputs(nodes_by_id, ins) + return { + "sampler_name": sampler_name, + "scheduler": _scalar(ins.get("scheduler")), + "steps": _scalar(ins.get("steps")) or _scalar(ins.get("denoise_steps")), + "cfg": _cfg_value_from_inputs(ins), + "denoise": _scalar(ins.get("denoise")), + "seed_val": seed_val, + } + + +def _sampler_name_from_inputs(sampler_node: Dict[str, Any], ins: Dict[str, Any]) -> Any: + sampler_name = _scalar(ins.get("sampler_name")) or _scalar(ins.get("sampler")) + if sampler_name: + return sampler_name + if "marigold" in _lower(_node_type(sampler_node)): + return _node_type(sampler_node) + return None + + +def _seed_value_from_inputs(nodes_by_id: Dict[str, Any], ins: Dict[str, Any]) -> Any: + seed_val = _scalar(ins.get("seed")) + if seed_val is not None: + return seed_val + if _is_link(ins.get("seed")): + return _resolve_scalar_from_link(nodes_by_id, ins.get("seed")) + return None + + +def _cfg_value_from_inputs(ins: Dict[str, Any]) -> Any: + for key in ("cfg", "cfg_scale", "guidance", "guidance_scale", "embedded_guidance_scale"): + value = _scalar(ins.get(key)) + if value is not None: + return value + return None + + +def _apply_widget_sampler_values(values: Dict[str, Any], sampler_node: Dict[str, Any]) -> None: + if not any(values.get(key) is None for key in ("sampler_name", "scheduler", "steps", "cfg", "denoise", "seed_val")): + return + ks_w = _extract_ksampler_widget_params(sampler_node) + if values.get("sampler_name") is None: + values["sampler_name"] = _scalar(ks_w.get("sampler_name")) + if values.get("scheduler") is None: + values["scheduler"] = _scalar(ks_w.get("scheduler")) + if values.get("steps") is None: + values["steps"] = _scalar(ks_w.get("steps")) + if values.get("cfg") is None: + values["cfg"] = _scalar(ks_w.get("cfg")) + if values.get("denoise") is None: + values["denoise"] = _scalar(ks_w.get("denoise")) + if values.get("seed_val") is None: + values["seed_val"] = _scalar(ks_w.get("seed")) + + +def _resolve_model_link_for_chain(ins: Dict[str, Any], trace: Dict[str, Any]) -> Any: + model_link_for_chain = ins.get("model") if _is_link(ins.get("model")) else None + if model_link_for_chain is None and _is_link(trace.get("guider_model_link")): + model_link_for_chain = trace.get("guider_model_link") + return model_link_for_chain + + +def _apply_advanced_sampler_values( + nodes_by_id: Dict[str, Any], + ins: Dict[str, Any], + values: Dict[str, Any], + trace: Dict[str, Any], + field_sources: Dict[str, str], + field_confidence: Dict[str, str], + model_link_for_chain: Any, +) -> Any: + _apply_advanced_sampler_name(nodes_by_id, ins, values) + model_link_for_chain = _apply_advanced_sigmas( + nodes_by_id, ins, values, field_sources, field_confidence, model_link_for_chain + ) + _apply_advanced_noise_seed(nodes_by_id, ins, values, field_sources) + _apply_advanced_cfg_from_conditioning(nodes_by_id, trace, values, field_sources) + return model_link_for_chain + + +def _apply_advanced_sampler_name(nodes_by_id: Dict[str, Any], ins: Dict[str, Any], values: Dict[str, Any]) -> None: + if not _is_link(ins.get("sampler")) or values.get("sampler_name"): + return + traced_sampler = _trace_sampler_name(nodes_by_id, ins.get("sampler")) + if traced_sampler: + values["sampler_name"] = traced_sampler[0] + + +def _apply_advanced_sigmas( + nodes_by_id: Dict[str, Any], + ins: Dict[str, Any], + values: Dict[str, Any], + field_sources: Dict[str, str], + field_confidence: Dict[str, str], + model_link_for_chain: Any, +) -> Any: + if not _is_link(ins.get("sigmas")): + return model_link_for_chain + st, sch, den, model_link, src, st_conf = _trace_scheduler_sigmas(nodes_by_id, ins.get("sigmas")) + src_name = src[1] if src and isinstance(src, tuple) and len(src) == 2 else None + _assign_advanced_sigma_values( + values, + field_sources, + field_confidence, + steps=st, + scheduler=sch, + denoise=den, + source_name=src_name, + steps_confidence=st_conf, + ) + if model_link and not model_link_for_chain: + return model_link + return model_link_for_chain + + +def _assign_advanced_sigma_values( + values: Dict[str, Any], + field_sources: Dict[str, str], + field_confidence: Dict[str, str], + *, + steps: Any, + scheduler: Any, + denoise: Any, + source_name: Optional[str], + steps_confidence: Optional[str], +) -> None: + steps_assigned = _assign_advanced_sigma_field(values, field_sources, "steps", steps, source_name) + _assign_advanced_sigma_field(values, field_sources, "scheduler", scheduler, source_name) + _assign_advanced_sigma_field(values, field_sources, "denoise", denoise, source_name) + if steps_assigned and steps_confidence: + field_confidence["steps"] = steps_confidence + + +def _assign_advanced_sigma_field( + values: Dict[str, Any], field_sources: Dict[str, str], key: str, value: Any, source_name: Optional[str] +) -> bool: + if value is None or values.get(key) is not None: + return False + values[key] = value + if source_name: + field_sources[key] = source_name + return True + + +def _apply_advanced_noise_seed( + nodes_by_id: Dict[str, Any], + ins: Dict[str, Any], + values: Dict[str, Any], + field_sources: Dict[str, str], +) -> None: + if not _is_link(ins.get("noise")) or values.get("seed_val") is not None: + return + traced_seed = _trace_noise_seed(nodes_by_id, ins.get("noise")) + if traced_seed: + values["seed_val"] = traced_seed[0] + field_sources["seed"] = traced_seed[1] + + +def _apply_advanced_cfg_from_conditioning( + nodes_by_id: Dict[str, Any], + trace: Dict[str, Any], + values: Dict[str, Any], + field_sources: Dict[str, str], +) -> None: + if not trace.get("conditioning_link") or values.get("cfg") is not None: + return + traced_cfg = _trace_guidance_from_conditioning(nodes_by_id, trace.get("conditioning_link")) + if traced_cfg: + values["cfg"] = traced_cfg[0] + field_sources["cfg"] = traced_cfg[1] + + +def _apply_guider_cfg_fallback(values: Dict[str, Any], trace: Dict[str, Any], field_sources: Dict[str, str]) -> None: + if values.get("cfg") is not None: + return + if trace.get("guider_cfg_value") is None: + return + values["cfg"] = trace.get("guider_cfg_value") + if trace.get("guider_cfg_source"): + field_sources["cfg"] = trace.get("guider_cfg_source") + + +def _extract_sampler_values( + nodes_by_id: Dict[str, Any], + sampler_node: Dict[str, Any], + sampler_id: str, + ins: Dict[str, Any], + advanced: bool, + confidence: str, + trace: Dict[str, Any], +) -> Dict[str, Any]: + _ = sampler_id, confidence field_sources: Dict[str, str] = {} field_confidence: Dict[str, str] = {} - # Advanced sampler pipelines (Flux/SD3): derive fields from orchestrator dependencies. - advanced = _is_advanced_sampler(sampler_node) - - # Qwen2-VL captioning workflows often don't have a sampler, but result in string output. - # If sink usage is text-based (rare for image gen tool, but possibly valid for caption), skip. - - # Detect Marigold (Depth) - often no sampler, just an estimator node or inserted in pipe. - # If the sink comes from a Marigold node, treat it as the "sampler". - if not sampler_id: - for nid, val in nodes_by_id.items(): - ct_lower = _lower(_node_type(val)) - if "marigold" in ct_lower: - # Marigold acts as the processor - sampler_id = nid - # Marigold usually takes 'image' input, not latent. - # Metadata to extract: seed, denoise_steps, ensemble_size. - break - # Qwen/Instruction nodes acted as sinks/processors - if "instruction" in ct_lower and "qwen" in ct_lower: - sampler_id = nid - # Needs special handling for prompt extraction later - break - - # Ensure sampler_node object is up to date if we just found an ID in the fallback - if sampler_id and (not sampler_node or str(sampler_node.get("id")) != str(sampler_id)): - sampler_node = nodes_by_id.get(sampler_id) or {} - ins = _inputs(sampler_node) - confidence = "low" # Fallback confidence - - # Prompts - pos_val: Optional[Tuple[str, str]] = None - neg_val: Optional[Tuple[str, str]] = None - conditioning_link = None - guider_cfg_value: Optional[Any] = None - guider_cfg_source: Optional[str] = None - guider_model_link: Optional[Any] = None - - # Handle Qwen Instruction Prompts (stored directly on the node) - if sampler_node and "instruction" in _lower(_node_type(sampler_node)) and "qwen" in _lower(_node_type(sampler_node)): - p = ins.get("instruction") or ins.get("text") - if isinstance(p, str) and p.strip(): - pos_val = (p.strip(), f"{_node_type(sampler_node)}:{sampler_id}:instruction") - - # FluxKohyaInferenceSampler (Kijai) - stores prompt directly - if "flux" in _lower(_node_type(sampler_node)) and "trainer" in _lower(_node_type(sampler_node)): - p = ins.get("prompt") - if isinstance(p, str) and p.strip(): - pos_val = (p.strip(), f"{_node_type(sampler_node)}:{sampler_id}:prompt") - - # Wan/video stacks: prompts are often encoded into text embeds rather than conditioning. - if _is_link(ins.get("text_embeds")) or _is_link(ins.get("hyvid_embeds")): - link = ins.get("text_embeds") or ins.get("hyvid_embeds") - p, n = _extract_posneg_from_text_embeds(nodes_by_id, link) - pos_val = pos_val or p - neg_val = neg_val or n - - if _is_link(ins.get("positive")): - items = _collect_texts_from_conditioning(nodes_by_id, ins.get("positive"), branch="positive") - if items: - text = "\n".join([t for t, _ in items]).strip() - srcs = [s for _, s in items] - source = srcs[0] if len(set(srcs)) <= 1 else f"{srcs[0]} (+{len(srcs)-1})" - if text: - pos_val = (text, source) - conditioning_link = ins.get("positive") - if _is_link(ins.get("negative")): - items = _collect_texts_from_conditioning(nodes_by_id, ins.get("negative"), branch="negative") - if items: - text = "\n".join([t for t, _ in items]).strip() - srcs = [s for _, s in items] - source = srcs[0] if len(set(srcs)) <= 1 else f"{srcs[0]} (+{len(srcs)-1})" - if text: - neg_val = (text, source) - - # Flux-style guidance pipelines: conditioning is passed through a guider node. - if advanced and _is_link(ins.get("guider")): - guider_id = _walk_passthrough(nodes_by_id, ins.get("guider")) - guider_node = nodes_by_id.get(guider_id) if guider_id else None - if isinstance(guider_node, dict): - gins = _inputs(guider_node) - # Some guiders pass a single `conditioning` link. - if _is_link(gins.get("conditioning")): - conditioning_link = gins.get("conditioning") - items = _collect_texts_from_conditioning(nodes_by_id, conditioning_link) - if items and not pos_val: - text = "\n".join([t for t, _ in items]).strip() - srcs = [s for _, s in items] - source = srcs[0] if len(set(srcs)) <= 1 else f"{srcs[0]} (+{len(srcs)-1})" - if text: - pos_val = (text, source) - - # Other guiders (e.g. CFGGuider in some video stacks) expose `positive` / `negative`. - if _is_link(gins.get("positive")): - conditioning_link = conditioning_link or gins.get("positive") - if not pos_val: - items = _collect_texts_from_conditioning(nodes_by_id, gins.get("positive"), branch="positive") - if items: - text = "\n".join([t for t, _ in items]).strip() - srcs = [s for _, s in items] - source = srcs[0] if len(set(srcs)) <= 1 else f"{srcs[0]} (+{len(srcs)-1})" - if text: - pos_val = (text, source) - - if _is_link(gins.get("negative")) and not neg_val: - items = _collect_texts_from_conditioning(nodes_by_id, gins.get("negative"), branch="negative") - if items: - text = "\n".join([t for t, _ in items]).strip() - srcs = [s for _, s in items] - source = srcs[0] if len(set(srcs)) <= 1 else f"{srcs[0]} (+{len(srcs)-1})" - if text: - neg_val = (text, source) - - # Guidance scale sometimes lives on the guider node. - cfg_val = _scalar(gins.get("cfg")) or _scalar(gins.get("cfg_scale")) or _scalar(gins.get("guidance")) - if cfg_val is not None: - guider_cfg_value = cfg_val - guider_cfg_source = f"{_node_type(guider_node)}:{guider_id}" - else: - # Flux-style: check if guidance is provided by an upstream node (e.g. FluxGuidance via conditioning) - cond_link = gins.get("conditioning") - found_guidance = _trace_guidance_value(nodes_by_id, cond_link) - if found_guidance: - guider_cfg_value, guider_cfg_source = found_guidance - - # Model chain for video stacks can also originate on the guider node. - if _is_link(gins.get("model")): - guider_model_link = gins.get("model") - - # Last-resort (still no guessing): some custom sampler nodes store prompt strings directly. - if not pos_val: - for k in ("positive_prompt", "prompt", "positive", "text", "text_g", "text_l"): - v = ins.get(k) - if _looks_like_prompt_string(v): - pos_val = (str(v).strip(), f"{_node_type(sampler_node)}:{sampler_id}:{k}") - break - if not neg_val: - for k in ("negative_prompt", "negative", "neg", "text_negative"): - v = ins.get(k) - if _looks_like_prompt_string(v): - neg_val = (str(v).strip(), f"{_node_type(sampler_node)}:{sampler_id}:{k}") - break - - # Sampler params - sampler_name = _scalar(ins.get("sampler_name")) or _scalar(ins.get("sampler")) - - # Check if this node class itself implies a specific sampler (Marigold/Depth) - if not sampler_name and "marigold" in _lower(_node_type(sampler_node)): - sampler_name = _node_type(sampler_node) - - scheduler = _scalar(ins.get("scheduler")) - steps = _scalar(ins.get("steps")) or _scalar(ins.get("denoise_steps")) # Marigold uses denoise_steps - # embedded_guidance_scale is used by HunyuanVideoSampler - cfg = _scalar(ins.get("cfg")) or _scalar(ins.get("cfg_scale")) or _scalar(ins.get("guidance")) or _scalar(ins.get("guidance_scale")) or _scalar(ins.get("embedded_guidance_scale")) - denoise = _scalar(ins.get("denoise")) - seed_val = _scalar(ins.get("seed")) - if seed_val is None and _is_link(ins.get("seed")): - seed_val = _resolve_scalar_from_link(nodes_by_id, ins.get("seed")) - - # LiteGraph workflows often keep KSampler scalar params in widgets_values. - if any(v is None for v in (sampler_name, scheduler, steps, cfg, denoise, seed_val)): - ks_w = _extract_ksampler_widget_params(sampler_node) - if sampler_name is None: - sampler_name = _scalar(ks_w.get("sampler_name")) - if scheduler is None: - scheduler = _scalar(ks_w.get("scheduler")) - if steps is None: - steps = _scalar(ks_w.get("steps")) - if cfg is None: - cfg = _scalar(ks_w.get("cfg")) - if denoise is None: - denoise = _scalar(ks_w.get("denoise")) - if seed_val is None: - seed_val = _scalar(ks_w.get("seed")) - - # Advanced sampler: sampler_name/steps/scheduler/denoise/seed may be stored on linked nodes. - model_link_for_chain = ins.get("model") if _is_link(ins.get("model")) else None - if model_link_for_chain is None and _is_link(guider_model_link): - model_link_for_chain = guider_model_link + values = _init_sampler_values(nodes_by_id, sampler_node, ins) + _apply_widget_sampler_values(values, sampler_node) + model_link_for_chain = _resolve_model_link_for_chain(ins, trace) if advanced: - if _is_link(ins.get("sampler")): - s = _trace_sampler_name(nodes_by_id, ins.get("sampler")) - if s and not sampler_name: - sampler_name = s[0] - if _is_link(ins.get("sigmas")): - st, sch, den, model_link, _src, st_conf = _trace_scheduler_sigmas(nodes_by_id, ins.get("sigmas")) - if st is not None and steps is None: - steps = st - if _src and isinstance(_src, tuple) and len(_src) == 2: - field_sources["steps"] = _src[1] - if st_conf: - field_confidence["steps"] = st_conf - if sch is not None and not scheduler: - scheduler = sch - if _src and isinstance(_src, tuple) and len(_src) == 2: - field_sources["scheduler"] = _src[1] - if den is not None and denoise is None: - denoise = den - if _src and isinstance(_src, tuple) and len(_src) == 2: - field_sources["denoise"] = _src[1] - if model_link and not model_link_for_chain: - model_link_for_chain = model_link - if _is_link(ins.get("noise")) and seed_val is None: - s = _trace_noise_seed(nodes_by_id, ins.get("noise")) - if s: - seed_val = s[0] - field_sources["seed"] = s[1] - if conditioning_link and cfg is None: - g = _trace_guidance_from_conditioning(nodes_by_id, conditioning_link) - if g: - cfg = g[0] - field_sources["cfg"] = g[1] + model_link_for_chain = _apply_advanced_sampler_values( + nodes_by_id, + ins, + values, + trace, + field_sources, + field_confidence, + model_link_for_chain, + ) + _apply_guider_cfg_fallback(values, trace, field_sources) - if cfg is None and guider_cfg_value is not None: - cfg = guider_cfg_value - if guider_cfg_source: - field_sources["cfg"] = guider_cfg_source + return { + **values, + "model_link_for_chain": model_link_for_chain, + "field_sources": field_sources, + "field_confidence": field_confidence, + } - # Model chain + LoRAs - models: Dict[str, Dict[str, Any]] = {} - loras: List[Dict[str, Any]] = [] + +def _collect_model_related_fields( + nodes_by_id: Dict[str, Any], + ins: Dict[str, Any], + confidence: str, + sink_start_id: Optional[str], + conditioning_link: Any, + model_link_for_chain: Any, +) -> Dict[str, Any]: + models, loras = _trace_models_and_loras(nodes_by_id, model_link_for_chain, confidence) + _ensure_upscaler_model(nodes_by_id, models) + size = _trace_size(nodes_by_id, ins.get("latent_image"), confidence) if _is_link(ins.get("latent_image")) else None + clip_skip = _trace_clip_skip_from_conditioning(nodes_by_id, conditioning_link, confidence) + clip = _trace_clip_from_text_encoder(nodes_by_id, conditioning_link, confidence) if conditioning_link and _is_link(conditioning_link) else None + vae = _trace_vae_from_sink(nodes_by_id, sink_start_id, confidence) if sink_start_id else None + + return { + "models": models, + "loras": loras, + "size": size, + "clip_skip": clip_skip, + "clip": clip, + "vae": vae, + } + + +def _trace_models_and_loras( + nodes_by_id: Dict[str, Any], + model_link_for_chain: Any, + confidence: str, +) -> Tuple[Dict[str, Dict[str, Any]], List[Dict[str, Any]]]: if model_link_for_chain and _is_link(model_link_for_chain): - models, loras = _trace_model_chain(nodes_by_id, model_link_for_chain, confidence) + return _trace_model_chain(nodes_by_id, model_link_for_chain, confidence) + return {}, [] - # Related model loaders (non-diffusion): e.g. latent upscalers in video workflows. - if "upscaler" not in models: - try: - for node_id, node in nodes_by_id.items(): - if not isinstance(node, dict): - continue - ct = _lower(_node_type(node)) - if "upscalemodelloader" not in ct and "upscale_model" not in ct and "latentupscale" not in ct: - continue - ins2 = _inputs(node) - name = _clean_model_id(ins2.get("model_name") or ins2.get("upscale_model") or ins2.get("upscale_model_name")) - if name: - models["upscaler"] = {"name": name, "confidence": "medium", "source": f"{_node_type(node)}:{node_id}"} - break - except Exception: - pass - # Size - size = None - if _is_link(ins.get("latent_image")): - size = _trace_size(nodes_by_id, ins.get("latent_image"), confidence) +def _ensure_upscaler_model(nodes_by_id: Dict[str, Any], models: Dict[str, Dict[str, Any]]) -> None: + if "upscaler" in models: + return + try: + for node_id, node in nodes_by_id.items(): + model_entry = _upscaler_model_entry(node, node_id) + if not model_entry: + continue + models["upscaler"] = model_entry + return + except Exception: + return - # Clip skip (from positive encoder's clip link) - clip_skip = None - encoder_id = None - if conditioning_link and _is_link(conditioning_link): - encoders = _collect_text_encoder_nodes_from_conditioning(nodes_by_id, conditioning_link, branch="positive") - encoder_id = encoders[0] if encoders else _walk_passthrough(nodes_by_id, conditioning_link) - pos_node = nodes_by_id.get(encoder_id) if encoder_id else None - if isinstance(pos_node, dict) and _is_link(_inputs(pos_node).get("clip")): - clip_skip = _trace_clip_skip(nodes_by_id, _inputs(pos_node).get("clip"), confidence) - # CLIP model (from text encoder clip input) - clip = None - if conditioning_link and _is_link(conditioning_link): - clip = _trace_clip_from_text_encoder(nodes_by_id, conditioning_link, confidence) +def _upscaler_model_entry(node: Any, node_id: Any) -> Optional[Dict[str, Any]]: + if not isinstance(node, dict): + return None + if not _is_upscaler_loader_type(_lower(_node_type(node))): + return None + name = _clean_model_id(_upscaler_model_name(_inputs(node))) + if not name: + return None + return {"name": name, "confidence": "medium", "source": f"{_node_type(node)}:{node_id}"} - # VAE model (from VAE decode on the sink path) - vae = None - if sink_start_id: - vae = _trace_vae_from_sink(nodes_by_id, sink_start_id, confidence) - # Determine Workflow Type (T2I, I2V, etc.) +def _is_upscaler_loader_type(node_type: str) -> bool: + return "upscalemodelloader" in node_type or "upscale_model" in node_type or "latentupscale" in node_type + + +def _upscaler_model_name(ins: Dict[str, Any]) -> Any: + return ins.get("model_name") or ins.get("upscale_model") or ins.get("upscale_model_name") + + +def _trace_clip_skip_from_conditioning(nodes_by_id: Dict[str, Any], conditioning_link: Any, confidence: str) -> Optional[Dict[str, Any]]: + if not conditioning_link or not _is_link(conditioning_link): + return None + encoders = _collect_text_encoder_nodes_from_conditioning(nodes_by_id, conditioning_link, branch="positive") + encoder_id = encoders[0] if encoders else _walk_passthrough(nodes_by_id, conditioning_link) + pos_node = nodes_by_id.get(encoder_id) if encoder_id else None + if not isinstance(pos_node, dict) or not _is_link(_inputs(pos_node).get("clip")): + return None + return _trace_clip_skip(nodes_by_id, _inputs(pos_node).get("clip"), confidence) + + +def _merge_models_payload(models: Dict[str, Dict[str, Any]], clip: Optional[Dict[str, Any]], vae: Optional[Dict[str, Any]]) -> Optional[Dict[str, Dict[str, Any]]]: + if not models and not clip and not vae: + return None + + merged: Dict[str, Dict[str, Any]] = {} + for key in ("checkpoint", "unet", "diffusion", "upscaler"): + if models.get(key): + merged[key] = models[key] + if clip: + merged["clip"] = clip + if vae: + merged["vae"] = vae + return merged or None + + +def _build_geninfo_payload( + nodes_by_id: Dict[str, Any], + sinks: List[str], + sink_id: str, + sampler_id: str, + sampler_mode: str, + sampler_source: str, + confidence: str, + workflow_meta: Optional[Dict[str, Any]], + trace: Dict[str, Any], + sampler_values: Dict[str, Any], + model_related: Dict[str, Any], +) -> Dict[str, Any]: wf_type = _determine_workflow_type(nodes_by_id, sink_id, sampler_id) - + out: Dict[str, Any] = { "engine": { "parser_version": "geninfo-v1", @@ -2424,95 +3182,147 @@ def _extract_geninfo(nodes_by_id: Dict[str, Any], sinks: List[str], workflow_met if workflow_meta: out["metadata"] = workflow_meta - if pos_val: - out["positive"] = {"value": pos_val[0], "confidence": confidence, "source": pos_val[1]} - if neg_val: - out["negative"] = {"value": neg_val[0], "confidence": confidence, "source": neg_val[1]} + _apply_trace_prompt_fields(out, trace, confidence) + _apply_lyrics_fields(out, nodes_by_id, sampler_source) + _apply_model_fields(out, model_related) + _apply_sampler_fields(out, sampler_values, confidence, sampler_source) + _apply_optional_model_metrics(out, model_related) + _apply_input_files_field(out, nodes_by_id) + _apply_multi_sink_prompt_fields(out, nodes_by_id, sinks) + return out + + +def _apply_trace_prompt_fields(out: Dict[str, Any], trace: Dict[str, Any], confidence: str) -> None: + if trace.get("pos_val"): + out["positive"] = {"value": trace["pos_val"][0], "confidence": confidence, "source": trace["pos_val"][1]} + if trace.get("neg_val"): + out["negative"] = {"value": trace["neg_val"][0], "confidence": confidence, "source": trace["neg_val"][1]} + + +def _apply_lyrics_fields(out: Dict[str, Any], nodes_by_id: Dict[str, Any], sampler_source: str) -> None: lyrics_text, lyrics_strength, lyrics_source = _extract_lyrics_from_prompt_nodes(nodes_by_id) if lyrics_text: out["lyrics"] = {"value": lyrics_text, "confidence": "high", "source": lyrics_source or sampler_source} if lyrics_strength is not None: - lyr_field = _field(lyrics_strength, "high", lyrics_source or sampler_source) - if lyr_field: - out["lyrics_strength"] = lyr_field + lyric_field = _field(lyrics_strength, "high", lyrics_source or sampler_source) + if lyric_field: + out["lyrics_strength"] = lyric_field - # Backward compatible: keep top-level `checkpoint` (best-effort) - if models: - preferred = ( - models.get("checkpoint") - or models.get("unet") - or models.get("diffusion") - or None + +def _apply_model_fields(out: Dict[str, Any], model_related: Dict[str, Any]) -> None: + models = model_related.get("models") or {} + loras = model_related.get("loras") or [] + clip = model_related.get("clip") + vae = model_related.get("vae") + _apply_preferred_checkpoint_field(out, models) + _set_if_present(out, "loras", loras) + _set_if_present(out, "clip", clip) + _set_if_present(out, "vae", vae) + merged_models = _merge_models_payload(models, clip, vae) + if merged_models: + out["models"] = merged_models + + +def _apply_preferred_checkpoint_field(out: Dict[str, Any], models: Dict[str, Any]) -> None: + if not models: + return + preferred = models.get("checkpoint") or models.get("unet") or models.get("diffusion") + if preferred: + out["checkpoint"] = preferred + + +def _set_if_present(out: Dict[str, Any], key: str, value: Any) -> None: + if value: + out[key] = value + + +def _apply_sampler_fields(out: Dict[str, Any], sampler_values: Dict[str, Any], confidence: str, sampler_source: str) -> None: + sampler_name_field = _field_name(sampler_values.get("sampler_name"), confidence, sampler_source) + if sampler_name_field: + out["sampler"] = sampler_name_field + scheduler_field = _field_name(sampler_values.get("scheduler"), confidence, sampler_source) + if scheduler_field: + out["scheduler"] = scheduler_field + for key in ("steps", "cfg", "seed", "denoise"): + value_key = "seed_val" if key == "seed" else key + field_value = _field( + sampler_values.get(value_key), + sampler_values.get("field_confidence", {}).get(key, confidence), + sampler_values.get("field_sources", {}).get(key, sampler_source), ) - if preferred: - out["checkpoint"] = preferred + if field_value: + out[key] = field_value - if loras: - out["loras"] = loras - if clip: - out["clip"] = clip - if vae: - out["vae"] = vae - if models or clip or vae: - merged: Dict[str, Dict[str, Any]] = {} - if models.get("checkpoint"): - merged["checkpoint"] = models["checkpoint"] - if models.get("unet"): - merged["unet"] = models["unet"] - if models.get("diffusion"): - merged["diffusion"] = models["diffusion"] - if models.get("upscaler"): - merged["upscaler"] = models["upscaler"] - if clip: - merged["clip"] = clip - if vae: - merged["vae"] = vae - if merged: - out["models"] = merged +def _apply_optional_model_metrics(out: Dict[str, Any], model_related: Dict[str, Any]) -> None: + if model_related.get("size"): + out["size"] = model_related["size"] + if model_related.get("clip_skip"): + out["clip_skip"] = model_related["clip_skip"] - sname = _field_name(sampler_name, confidence, sampler_source) - if sname: - out["sampler"] = sname - sch = _field_name(scheduler, confidence, sampler_source) - if sch: - out["scheduler"] = sch - for key, val in ( - ("steps", steps), - ("cfg", cfg), - ("seed", seed_val), - ("denoise", denoise), - ): - source = field_sources.get(key, sampler_source) - conf = field_confidence.get(key, confidence) - f = _field(val, conf, source) - if f: - out[key] = f - - if size: - out["size"] = size - - if clip_skip: - out["clip_skip"] = clip_skip - - # Extract inputs (images/videos used) +def _apply_input_files_field(out: Dict[str, Any], nodes_by_id: Dict[str, Any]) -> None: input_files = _extract_input_files(nodes_by_id) if input_files: out["inputs"] = input_files - # For multi-output workflows, collect all distinct prompts - # This helps users see all variations when a workflow generates multiple outputs - if len(sinks) > 1: - all_pos, all_neg = _collect_all_prompts_from_sinks(nodes_by_id, sinks) - # Only add if there are multiple distinct prompts - if len(all_pos) > 1: - out["all_positive_prompts"] = all_pos - if len(all_neg) > 1: - out["all_negative_prompts"] = all_neg - # Do-not-lie: if nothing useful extracted besides engine, return None. +def _apply_multi_sink_prompt_fields(out: Dict[str, Any], nodes_by_id: Dict[str, Any], sinks: List[str]) -> None: + if len(sinks) <= 1: + return + all_pos, all_neg = _collect_all_prompts_from_sinks(nodes_by_id, sinks) + if len(all_pos) > 1: + out["all_positive_prompts"] = all_pos + if len(all_neg) > 1: + out["all_negative_prompts"] = all_neg + + +def _extract_geninfo(nodes_by_id: Dict[str, Any], sinks: List[str], workflow_meta: Optional[Dict]) -> Result: + sink_id = sinks[0] + sampler_id, sampler_conf, sampler_mode = _select_sampler_context(nodes_by_id, sink_id) + + if not sampler_id: + return _build_no_sampler_result(nodes_by_id, workflow_meta) + + sink = nodes_by_id.get(sink_id) or {} + sink_link = _pick_sink_inputs(sink) + sink_start_id = _walk_passthrough(nodes_by_id, sink_link) if sink_link else None + + sampler_node = nodes_by_id.get(sampler_id) or {} + if not isinstance(sampler_node, dict): + return _build_no_sampler_result(nodes_by_id, workflow_meta) + + confidence = sampler_conf if sampler_conf != "none" else "low" + ins = _inputs(sampler_node) + sampler_source = f"{_node_type(sampler_node)}:{sampler_id}" + advanced = _is_advanced_sampler(sampler_node) + + trace = _extract_prompt_trace(nodes_by_id, sampler_node, sampler_id, ins, advanced) + sampler_values = _extract_sampler_values(nodes_by_id, sampler_node, sampler_id, ins, advanced, confidence, trace) + model_related = _collect_model_related_fields( + nodes_by_id, + ins, + confidence, + sink_start_id, + trace.get("conditioning_link"), + sampler_values.get("model_link_for_chain"), + ) + + out = _build_geninfo_payload( + nodes_by_id, + sinks, + sink_id, + sampler_id, + sampler_mode, + sampler_source, + confidence, + workflow_meta, + trace, + sampler_values, + model_related, + ) + if len(out.keys()) <= 1: if workflow_meta: return Result.Ok({"metadata": workflow_meta}) diff --git a/mjr_am_backend/features/health/service.py b/mjr_am_backend/features/health/service.py index 0c1f8727..b1853a8a 100644 --- a/mjr_am_backend/features/health/service.py +++ b/mjr_am_backend/features/health/service.py @@ -197,79 +197,23 @@ class HealthService: try: where_sql, where_params = self._roots_where(roots) if roots else ("1=1", []) params = tuple(where_params) + counts = await self._counter_query_results(where_sql, params) + last_scan_end = await self._metadata_value("last_scan_end") + last_index_end = await self._metadata_value("last_index_end") - total_result = await self.db.aexecute( - f"SELECT COUNT(*) as count FROM assets a WHERE {where_sql}", - params, - fetch=True, - ) - - kind_result = await self.db.aexecute( - f"SELECT a.kind, COUNT(*) as count FROM assets a WHERE {where_sql} GROUP BY a.kind", - params, - fetch=True, - ) - - rated_result = await self.db.aexecute( - f""" - SELECT COUNT(*) as count - FROM asset_metadata m - JOIN assets a ON a.id = m.asset_id - WHERE m.rating > 0 AND {where_sql} - """, - params, - fetch=True, - ) - - workflow_result = await self.db.aexecute( - f""" - SELECT COUNT(*) as count - FROM asset_metadata m - JOIN assets a ON a.id = m.asset_id - WHERE m.has_workflow = 1 AND {where_sql} - """, - params, - fetch=True, - ) - - generation_result = await self.db.aexecute( - f""" - SELECT COUNT(*) as count - FROM asset_metadata m - JOIN assets a ON a.id = m.asset_id - WHERE m.has_generation_data = 1 AND {where_sql} - """, - params, - fetch=True, - ) - - last_scan_result = await self.db.aexecute( - "SELECT value FROM metadata WHERE key = 'last_scan_end'", - fetch=True - ) - last_scan_end = last_scan_result.data[0]["value"] if last_scan_result.ok and last_scan_result.data else None - - last_index_result = await self.db.aexecute( - "SELECT value FROM metadata WHERE key = 'last_index_end'", - fetch=True - ) - last_index_end = last_index_result.data[0]["value"] if last_index_result.ok and last_index_result.data else None - - # Get counts by kind - by_kind = {} - if kind_result.ok and kind_result.data and len(kind_result.data) > 0: - by_kind = {row["kind"]: row["count"] for row in kind_result.data} - - workflow_count = workflow_result.data[0]["count"] if workflow_result.ok and workflow_result.data else 0 - generation_count = generation_result.data[0]["count"] if generation_result.ok and generation_result.data else 0 + by_kind = self._kind_counts(counts["kind"]) + total_count = self._result_count(counts["total"]) + rated_count = self._result_count(counts["rated"]) + workflow_count = self._result_count(counts["workflow"]) + generation_count = self._result_count(counts["generation"]) counters = { - "total_assets": total_result.data[0]["count"] if total_result.ok and total_result.data else 0, + "total_assets": total_count, "images": by_kind.get("image", 0), "videos": by_kind.get("video", 0), "audio": by_kind.get("audio", 0), "by_kind": by_kind, # Keep for backward compatibility - "rated": rated_result.data[0]["count"] if rated_result.ok and rated_result.data else 0, + "rated": rated_count, "with_workflow": workflow_count, "with_workflows": workflow_count, "with_generation_data": generation_count, @@ -289,6 +233,78 @@ class HealthService: logger.error(f"Failed to get counters: {e}") return Result.Err("DB_ERROR", str(e)) + async def _counter_query_results(self, where_sql: str, params: tuple[str, ...]) -> dict: + return { + "total": await self.db.aexecute( + f"SELECT COUNT(*) as count FROM assets a WHERE {where_sql}", + params, + fetch=True, + ), + "kind": await self.db.aexecute( + f"SELECT a.kind, COUNT(*) as count FROM assets a WHERE {where_sql} GROUP BY a.kind", + params, + fetch=True, + ), + "rated": await self.db.aexecute( + f""" + SELECT COUNT(*) as count + FROM asset_metadata m + JOIN assets a ON a.id = m.asset_id + WHERE m.rating > 0 AND {where_sql} + """, + params, + fetch=True, + ), + "workflow": await self.db.aexecute( + f""" + SELECT COUNT(*) as count + FROM asset_metadata m + JOIN assets a ON a.id = m.asset_id + WHERE m.has_workflow = 1 AND {where_sql} + """, + params, + fetch=True, + ), + "generation": await self.db.aexecute( + f""" + SELECT COUNT(*) as count + FROM asset_metadata m + JOIN assets a ON a.id = m.asset_id + WHERE m.has_generation_data = 1 AND {where_sql} + """, + params, + fetch=True, + ), + } + + async def _metadata_value(self, key: str) -> Optional[str]: + result = await self.db.aexecute( + "SELECT value FROM metadata WHERE key = ?", + (key,), + fetch=True, + ) + if result.ok and result.data: + return result.data[0].get("value") + return None + + @staticmethod + def _result_count(result: object) -> int: + try: + if getattr(result, "ok", False) and getattr(result, "data", None): + return int(result.data[0].get("count") or 0) + except Exception: + return 0 + return 0 + + @staticmethod + def _kind_counts(kind_result: object) -> dict: + try: + if getattr(kind_result, "ok", False) and getattr(kind_result, "data", None): + return {row["kind"]: row["count"] for row in kind_result.data} + except Exception: + return {} + return {} + def _get_tool_capabilities(self) -> dict: """Report availability of required tooling (ExifTool/FFprobe).""" return { diff --git a/mjr_am_backend/features/index/enricher.py b/mjr_am_backend/features/index/enricher.py index 50f32d39..c490947c 100644 --- a/mjr_am_backend/features/index/enricher.py +++ b/mjr_am_backend/features/index/enricher.py @@ -105,23 +105,12 @@ class MetadataEnricher: Args: filepaths: List of file paths to enqueue """ - cleaned = [str(p) for p in (filepaths or []) if p] + cleaned = self._clean_paths(filepaths) if not cleaned: return async with self._enrich_lock: - existing = {fp for _, fp in self._enrich_queue} - for fp in cleaned: - if fp in existing: - continue - self._enrich_queue.append((0, fp)) - existing.add(fp) - # Keep highest priority first (negative = higher priority). - self._enrich_queue.sort(key=lambda item: item[0]) - if self._enrich_running and self._enrich_task and not self._enrich_task.done(): - return - self._emit_status(True, queued=len(self._enrich_queue)) - self._enrich_running = True - self._enrich_task = asyncio.create_task(self._enrichment_worker()) + self._append_queue_items(cleaned, priority=0) + self._ensure_worker_started_locked(emit_on_running=False) async def enqueue(self, path: str, priority: int = 0) -> None: """ @@ -132,21 +121,40 @@ class MetadataEnricher: if not fp: return async with self._enrich_lock: - existing = {p for _, p in self._enrich_queue} + self._append_queue_items([fp], priority=priority) + self._ensure_worker_started_locked(emit_on_running=True) + + @staticmethod + def _clean_paths(filepaths: List[str]) -> List[str]: + return [str(p) for p in (filepaths or []) if p] + + @staticmethod + def _normalize_priority(priority: int) -> int: + try: + return int(priority) + except Exception: + return 0 + + def _append_queue_items(self, filepaths: List[str], priority: int) -> None: + existing = {fp for _, fp in self._enrich_queue} + prio = self._normalize_priority(priority) + for fp in filepaths: if fp in existing: - return - try: - prio = int(priority) - except Exception: - prio = 0 + continue self._enrich_queue.append((prio, fp)) - self._enrich_queue.sort(key=lambda item: item[0]) - if self._enrich_running and self._enrich_task and not self._enrich_task.done(): + existing.add(fp) + # Keep highest priority first (negative = higher priority). + self._enrich_queue.sort(key=lambda item: item[0]) + + def _ensure_worker_started_locked(self, emit_on_running: bool) -> None: + running = self._enrich_running and self._enrich_task and not self._enrich_task.done() + if running: + if emit_on_running: self._emit_status(True, queued=len(self._enrich_queue)) - return - self._emit_status(True, queued=len(self._enrich_queue)) - self._enrich_running = True - self._enrich_task = asyncio.create_task(self._enrichment_worker()) + return + self._emit_status(True, queued=len(self._enrich_queue)) + self._enrich_running = True + self._enrich_task = asyncio.create_task(self._enrichment_worker()) async def stop_enrichment(self, clear_queue: bool = True) -> None: """ @@ -198,24 +206,15 @@ class MetadataEnricher: queue_left = len(self._enrich_queue) self._emit_status(False, queue_left=queue_left) - async def _enrich_metadata_chunk(self, filepaths: List[str]) -> None: - """ - Process a chunk of files for metadata enrichment. - - Args: - filepaths: List of file paths to process - """ - cleaned = [str(p) for p in (filepaths or []) if p] - if not cleaned: - return - + async def _load_asset_ids_for_filepaths(self, filepaths: List[str]) -> Result[Dict[str, int]]: id_res = await self.db.aquery_in( "SELECT id, filepath FROM assets WHERE {IN_CLAUSE}", "filepath", - cleaned, + filepaths, ) if not id_res.ok: - return + return Result.Err(id_res.code, id_res.error or "Failed to query asset ids") + id_by_fp: Dict[str, int] = {} for row in id_res.data or []: if not isinstance(row, dict): @@ -227,182 +226,291 @@ class MetadataEnricher: aid = 0 if fp and aid: id_by_fp[str(fp)] = aid + return Result.Ok(id_by_fp) + def _build_state_hash_for_file(self, filepath: str) -> Optional[str]: + try: + stat = os.stat(filepath) + except Exception: + return None + mtime_ns = getattr(stat, "st_mtime_ns", int(stat.st_mtime * 1_000_000_000)) + return self._compute_state_hash(filepath, int(mtime_ns), int(stat.st_size)) + + @staticmethod + def _extract_dimensions(metadata_result: Result[Dict[str, Any]]) -> tuple[Any, Any, Any]: + width = None + height = None + duration = None + if metadata_result.ok and metadata_result.data: + meta = metadata_result.data + width = meta.get("width") + height = meta.get("height") + duration = meta.get("duration") + return width, height, duration + + def _build_update_item( + self, + asset_id: int, + filepath: str, + metadata_result: Result[Dict[str, Any]], + ) -> Dict[str, Any]: + width, height, duration = self._extract_dimensions(metadata_result) + return { + "asset_id": asset_id, + "width": width, + "height": height, + "duration": duration, + "metadata_result": metadata_result, + "filepath": filepath, + } + + async def _prepare_updates_from_cache( + self, + cleaned: List[str], + id_by_fp: Dict[str, int], + metadata_helpers_cls: Any, + ) -> tuple[List[Dict[str, Any]], List[tuple[str, int, str]]]: updates: List[Dict[str, Any]] = [] - retry_paths: List[str] = [] - - # Import helper locally to use shared logic (including FTS population) - from .metadata_helpers import MetadataHelpers - to_extract: List[tuple[str, int, str]] = [] for fp in cleaned: asset_id = id_by_fp.get(fp) if not asset_id: continue - try: - stat = os.stat(fp) - except Exception: - continue - mtime_ns = getattr(stat, "st_mtime_ns", int(stat.st_mtime * 1_000_000_000)) - state_hash = self._compute_state_hash(fp, int(mtime_ns), int(stat.st_size)) - # Try cache first to avoid tool work. - metadata_result = await MetadataHelpers.retrieve_cached_metadata(self.db, fp, state_hash) + state_hash = self._build_state_hash_for_file(fp) + if not state_hash: + continue + + metadata_result = await metadata_helpers_cls.retrieve_cached_metadata( + self.db, + fp, + state_hash, + ) if not (metadata_result and metadata_result.ok): to_extract.append((fp, int(asset_id), state_hash)) continue - if not metadata_result.ok: - metadata_result = self._metadata_error_payload(metadata_result, fp) + updates.append(self._build_update_item(int(asset_id), fp, metadata_result)) + return updates, to_extract - width = None - height = None - duration = None - if metadata_result.ok and metadata_result.data: - meta = metadata_result.data - width = meta.get("width") - height = meta.get("height") - duration = meta.get("duration") - - updates.append( - { - "asset_id": asset_id, - "width": width, - "height": height, - "duration": duration, - "metadata_result": metadata_result, - "filepath": fp, - } - ) - - if to_extract: - batch_results = await self.metadata.get_metadata_batch([fp for fp, _, _ in to_extract], scan_id=None) - for fp, asset_id, state_hash in to_extract: - metadata_result = batch_results.get(fp) - if not metadata_result: - metadata_result = self._metadata_error_payload(Result.Err("METADATA_MISSING", "No metadata returned"), fp) - elif not metadata_result.ok: - metadata_result = self._metadata_error_payload(metadata_result, fp) - else: - try: - await MetadataHelpers.store_metadata_cache(self.db, fp, state_hash, metadata_result) - except Exception: - pass - - width = None - height = None - duration = None - if metadata_result.ok and metadata_result.data: - meta = metadata_result.data - width = meta.get("width") - height = meta.get("height") - duration = meta.get("duration") - - updates.append( - { - "asset_id": asset_id, - "width": width, - "height": height, - "duration": duration, - "metadata_result": metadata_result, - "filepath": fp, - } + async def _prepare_updates_from_extraction( + self, + to_extract: List[tuple[str, int, str]], + metadata_helpers_cls: Any, + ) -> List[Dict[str, Any]]: + if not to_extract: + return [] + updates: List[Dict[str, Any]] = [] + batch_results = await self.metadata.get_metadata_batch( + [fp for fp, _, _ in to_extract], + scan_id=None, + ) + for fp, asset_id, state_hash in to_extract: + metadata_result = batch_results.get(fp) + if not metadata_result: + metadata_result = self._metadata_error_payload( + Result.Err("METADATA_MISSING", "No metadata returned"), + fp, ) + elif not metadata_result.ok: + metadata_result = self._metadata_error_payload(metadata_result, fp) + else: + try: + await metadata_helpers_cls.store_metadata_cache( + self.db, + fp, + state_hash, + metadata_result, + ) + except Exception: + pass + updates.append(self._build_update_item(asset_id, fp, metadata_result)) + return updates + + async def _notify_asset_updated(self, asset_id: int) -> None: + try: + res = await self.db.aquery( + """ + SELECT a.id, m.has_workflow AS has_workflow, + m.has_generation_data AS has_generation_data + FROM assets a + LEFT JOIN asset_metadata m ON a.id = m.asset_id + WHERE a.id = ? + LIMIT 1 + """, + (int(asset_id),), + ) + if res.ok and res.data: + payload = dict(res.data[0]) + from ...routes.registry import PromptServer + PromptServer.instance.send_sync("mjr-asset-updated", payload) + except Exception: + pass + + async def _apply_update_item(self, item: Dict[str, Any], metadata_helpers_cls: Any) -> Optional[str]: + asset_id = item.get("asset_id") + if not asset_id: + return None + meta_res = item.get("metadata_result") + if not isinstance(meta_res, Result): + return None + fp = str(item.get("filepath") or "") + + try: + update_error = await self._apply_metadata_update_transaction( + asset_id=asset_id, + item=item, + meta_res=meta_res, + metadata_helpers_cls=metadata_helpers_cls, + ) + if update_error is not None: + return update_error + if fp: + self._retry_counts.pop(fp, None) + await self._notify_asset_updated(int(asset_id)) + return None + except Exception as exc: + if _is_transient_db_error(exc): + logger.info( + "Metadata enrichment deferred for asset_id=%s due to DB contention/reset: %s", + asset_id, + exc, + ) + return fp or None + logger.warning("Metadata enrichment update failed for asset_id=%s: %s", asset_id, exc) + return None + + async def _apply_metadata_update_transaction( + self, + *, + asset_id: Any, + item: Dict[str, Any], + meta_res: Result[Dict[str, Any]], + metadata_helpers_cls: Any, + ) -> Optional[str]: + fp = str(item.get("filepath") or "") + async with self.db.lock_for_asset(asset_id): + async with self.db.atransaction(mode="deferred") as tx: + if not tx.ok: + logger.warning( + "Metadata enrichment skipped (transaction begin failed) for asset_id=%s: %s", + asset_id, + tx.error, + ) + return fp or None + await self._apply_metadata_update_queries(asset_id, item, meta_res, metadata_helpers_cls) + if not tx.ok: + logger.warning( + "Metadata enrichment commit failed for asset_id=%s: %s", + asset_id, + tx.error, + ) + return fp or None + return None + + async def _apply_metadata_update_queries( + self, + asset_id: Any, + item: Dict[str, Any], + meta_res: Result[Dict[str, Any]], + metadata_helpers_cls: Any, + ) -> None: + await self.db.aexecute( + """ + UPDATE assets + SET width = COALESCE(?, width), + height = COALESCE(?, height), + duration = COALESCE(?, duration), + updated_at = CURRENT_TIMESTAMP + WHERE id = ? + """, + ( + item.get("width"), + item.get("height"), + item.get("duration"), + asset_id, + ), + ) + await metadata_helpers_cls.write_asset_metadata_row( + self.db, + asset_id, + meta_res, + filepath=item.get("filepath"), + ) + + async def _apply_updates_and_collect_retries( + self, + updates: List[Dict[str, Any]], + metadata_helpers_cls: Any, + ) -> List[str]: + retry_paths: List[str] = [] + for item in updates: + retry_fp = await self._apply_update_item(item, metadata_helpers_cls) + if retry_fp: + retry_paths.append(retry_fp) + return retry_paths + + async def _requeue_retry_paths(self, retry_paths: List[str]) -> None: + if not retry_paths: + return + to_requeue: List[str] = [] + for fp in retry_paths: + try: + count = int(self._retry_counts.get(fp, 0)) + 1 + except Exception: + count = 1 + self._retry_counts[fp] = count + if count <= _MAX_ENRICH_RETRIES: + to_requeue.append(fp) + continue + logger.warning( + "Metadata enrichment dropped after max retries (%s): %s", + _MAX_ENRICH_RETRIES, + fp, + ) + self._retry_counts.pop(fp, None) + if to_requeue: + await asyncio.sleep(0.2) + await self.start_enrichment(to_requeue) + + async def _enrich_metadata_chunk(self, filepaths: List[str]) -> None: + """ + Process a chunk of files for metadata enrichment. + + Args: + filepaths: List of file paths to process + """ + cleaned = [str(p) for p in (filepaths or []) if p] + if not cleaned: + return + + id_map_res = await self._load_asset_ids_for_filepaths(cleaned) + if not id_map_res.ok: + return + id_by_fp = id_map_res.data or {} + + # Import helper locally to use shared logic (including FTS population) + from .metadata_helpers import MetadataHelpers + + updates, to_extract = await self._prepare_updates_from_cache( + cleaned, + id_by_fp, + MetadataHelpers, + ) + extracted_updates = await self._prepare_updates_from_extraction( + to_extract, + MetadataHelpers, + ) + if extracted_updates: + updates.extend(extracted_updates) if not updates: return - for item in updates: - asset_id = item.get("asset_id") - if not asset_id: - continue - meta_res = item.get("metadata_result") - if not isinstance(meta_res, Result): - continue - try: - async with self.db.lock_for_asset(asset_id): - # Use deferred tx to reduce lock contention with concurrent search/sort reads. - async with self.db.atransaction(mode="deferred") as tx: - if not tx.ok: - logger.warning( - "Metadata enrichment skipped (transaction begin failed) for asset_id=%s: %s", - asset_id, - tx.error, - ) - fp = str(item.get("filepath") or "") - if fp: - retry_paths.append(fp) - continue - await self.db.aexecute( - """ - UPDATE assets - SET width = COALESCE(?, width), - height = COALESCE(?, height), - duration = COALESCE(?, duration), - updated_at = CURRENT_TIMESTAMP - WHERE id = ? - """, - (item.get("width"), item.get("height"), item.get("duration"), asset_id), - ) - # Use shared helper to write metadata (ensures tags_text/FTS is populated with prompt/geninfo) - await MetadataHelpers.write_asset_metadata_row( - self.db, - asset_id, - meta_res, - filepath=item.get("filepath"), - ) - if not tx.ok: - logger.warning("Metadata enrichment commit failed for asset_id=%s: %s", asset_id, tx.error) - fp = str(item.get("filepath") or "") - if fp: - retry_paths.append(fp) - else: - fp = str(item.get("filepath") or "") - if fp: - self._retry_counts.pop(fp, None) - # Notify frontend so workflow dot updates promptly. - try: - res = await self.db.aquery( - """ - SELECT a.id, m.has_workflow AS has_workflow, - m.has_generation_data AS has_generation_data - FROM assets a - LEFT JOIN asset_metadata m ON a.id = m.asset_id - WHERE a.id = ? - LIMIT 1 - """, - (int(asset_id),), - ) - if res.ok and res.data: - payload = dict(res.data[0]) - from ...routes.registry import PromptServer - PromptServer.instance.send_sync("mjr-asset-updated", payload) - except Exception: - pass - except Exception as exc: - fp = str(item.get("filepath") or "") - if _is_transient_db_error(exc): - logger.info("Metadata enrichment deferred for asset_id=%s due to DB contention/reset: %s", asset_id, exc) - if fp: - retry_paths.append(fp) - else: - logger.warning("Metadata enrichment update failed for asset_id=%s: %s", asset_id, exc) - - if retry_paths: - to_requeue: List[str] = [] - for fp in retry_paths: - try: - c = int(self._retry_counts.get(fp, 0)) + 1 - except Exception: - c = 1 - self._retry_counts[fp] = c - if c <= _MAX_ENRICH_RETRIES: - to_requeue.append(fp) - else: - logger.warning("Metadata enrichment dropped after max retries (%s): %s", _MAX_ENRICH_RETRIES, fp) - self._retry_counts.pop(fp, None) - if to_requeue: - await asyncio.sleep(0.2) - await self.start_enrichment(to_requeue) + retry_paths = await self._apply_updates_and_collect_retries( + updates, + MetadataHelpers, + ) + await self._requeue_retry_paths(retry_paths) def get_queue_length(self) -> int: """Return pending enrichment queue length.""" diff --git a/mjr_am_backend/features/index/metadata_helpers.py b/mjr_am_backend/features/index/metadata_helpers.py index 3b4a1346..839c4ef8 100644 --- a/mjr_am_backend/features/index/metadata_helpers.py +++ b/mjr_am_backend/features/index/metadata_helpers.py @@ -28,6 +28,339 @@ _METADATA_CACHE_CLEANUP_LOCK = asyncio.Lock() _METADATA_CACHE_LAST_CLEANUP = 0.0 +def _metadata_quality_should_upgrade_sql() -> str: + quality_rank_excluded = ( + "CASE excluded.metadata_quality WHEN 'full' THEN 3 WHEN 'partial' THEN 2 " + "WHEN 'degraded' THEN 1 ELSE 0 END" + ) + quality_rank_current = ( + "CASE COALESCE(asset_metadata.metadata_quality, 'none') WHEN 'full' THEN 3 " + "WHEN 'partial' THEN 2 WHEN 'degraded' THEN 1 ELSE 0 END" + ) + return f"({quality_rank_excluded} >= {quality_rank_current})" + + +def _apply_metadata_json_size_guard( + asset_id: int, + metadata_result: Result[Dict[str, Any]], + metadata_raw_json: str, + filepath: Optional[str] = None, +) -> str: + max_bytes = _safe_max_metadata_bytes() + truncated, original_bytes, metadata_raw_json = _truncate_metadata_json_if_needed( + metadata_raw_json, max_bytes + ) + if not truncated: + return metadata_raw_json + _mark_metadata_truncated(metadata_result, original_bytes=original_bytes, max_bytes=max_bytes) + _log_metadata_truncation( + asset_id, + metadata_result, + filepath=filepath, + original_bytes=original_bytes, + max_bytes=max_bytes, + ) + return metadata_raw_json + + +def _safe_max_metadata_bytes() -> int: + try: + return int(MAX_METADATA_JSON_BYTES or 0) + except Exception: + return 0 + + +def _truncate_metadata_json_if_needed(metadata_raw_json: str, max_bytes: int) -> tuple[bool, int, str]: + if not (max_bytes and isinstance(metadata_raw_json, str)): + return False, 0, metadata_raw_json + try: + nbytes = len(metadata_raw_json.encode("utf-8", errors="ignore")) + except Exception: + nbytes = 0 + if nbytes <= max_bytes: + return False, 0, metadata_raw_json + try: + truncated = json.dumps( + {"_truncated": True, "original_bytes": int(nbytes)}, + ensure_ascii=False, + separators=(",", ":"), + ) + except Exception: + truncated = "{}" + return True, int(nbytes), truncated + + +def _mark_metadata_truncated(metadata_result: Result[Dict[str, Any]], *, original_bytes: int, max_bytes: int) -> None: + try: + metadata_result.meta["truncated"] = True + metadata_result.meta["original_bytes"] = int(original_bytes) + metadata_result.meta["max_bytes"] = int(max_bytes) + except Exception: + return + + +def _resolve_metadata_filepath( + metadata_result: Result[Dict[str, Any]], filepath: Optional[str] +) -> Optional[str]: + if filepath: + return filepath + if metadata_result and isinstance(metadata_result.data, dict): + fp = metadata_result.data.get("filepath") + return fp if isinstance(fp, str) and fp else None + return None + + +def _log_metadata_truncation( + asset_id: int, + metadata_result: Result[Dict[str, Any]], + *, + filepath: Optional[str], + original_bytes: int, + max_bytes: int, +) -> None: + try: + fp = _resolve_metadata_filepath(metadata_result, filepath) + suffix = f" filepath={fp}" if fp else "" + logger.warning( + "Metadata JSON truncated for asset_id=%s%s (bytes=%s > max=%s)", + asset_id, + suffix, + original_bytes, + max_bytes, + ) + except Exception: + return + + +def _sanitize_metadata_tags(raw_tags: Any) -> list[str]: + if not isinstance(raw_tags, list): + return [] + cleaned: list[str] = [] + seen: set[str] = set() + for item in raw_tags: + if not isinstance(item, str): + continue + tag = item.strip() + if not tag or len(tag) > MAX_TAG_LENGTH or tag in seen: + continue + seen.add(tag) + cleaned.append(tag) + return cleaned + + +def _extract_rating_and_tags(metadata_result: Result[Dict[str, Any]]) -> tuple[int, str, str]: + extracted_rating = 0 + extracted_tags_json = "[]" + extracted_tags_text = "" + if not (metadata_result and metadata_result.ok and metadata_result.data): + return extracted_rating, extracted_tags_json, extracted_tags_text + + meta = metadata_result.data + try: + rating_val = meta.get("rating") + if rating_val is not None: + extracted_rating = max(0, min(5, int(rating_val))) + except Exception: + extracted_rating = 0 + + try: + cleaned = _sanitize_metadata_tags(meta.get("tags")) + extracted_tags_json = json.dumps(cleaned, ensure_ascii=False) + extracted_tags_text = " ".join(cleaned) + except Exception: + extracted_tags_json = "[]" + extracted_tags_text = "" + return extracted_rating, extracted_tags_json, extracted_tags_text + + +def _collect_geninfo_extras(meta: Dict[str, Any]) -> list[str]: + extras: list[str] = [] + geninfo = meta.get("geninfo") + if not isinstance(geninfo, dict): + return extras + + extras.extend(_collect_named_values(geninfo.get("models"), "name")) + extras.extend(_collect_named_values(geninfo.get("loras"), "name")) + extras.extend(_collect_nested_text_values(geninfo, [("positive", "value"), ("engine", "type")])) + extras.extend(_collect_input_filenames(geninfo.get("inputs"))) + return extras + + +def _collect_named_values(container: Any, field: str) -> list[str]: + out: list[str] = [] + if isinstance(container, dict): + values = container.values() + elif isinstance(container, list): + values = container + else: + return out + for item in values: + if not isinstance(item, dict): + continue + value = str(item.get(field) or "").strip() + if value: + out.append(value) + return out + + +def _collect_nested_text_values(geninfo: Dict[str, Any], paths: list[tuple[str, str]]) -> list[str]: + out: list[str] = [] + for parent_key, child_key in paths: + parent = geninfo.get(parent_key) + if not isinstance(parent, dict): + continue + value = str(parent.get(child_key) or "").strip() + if value: + out.append(value) + return out + + +def _collect_input_filenames(inputs: Any) -> list[str]: + out: list[str] = [] + if not isinstance(inputs, list): + return out + for item in inputs: + if not isinstance(item, dict): + continue + fname = str(item.get("filename") or "").strip() + if fname: + out.append(fname) + return out + + +def _enrich_tags_text_with_metadata( + metadata_result: Result[Dict[str, Any]], + extracted_tags_text: str, +) -> str: + if not (metadata_result and metadata_result.ok and metadata_result.data): + return extracted_tags_text + meta = metadata_result.data + extras = _collect_geninfo_extras(meta) + extracted_tags_text = _append_extra_tags(extracted_tags_text, extras) + _build_metadata_payload_preview(meta, extracted_tags_text, extras) + return extracted_tags_text + + +def _append_extra_tags(extracted_tags_text: str, extras: list[str]) -> str: + if not extras: + return extracted_tags_text + extra_text = " ".join(extras) + if extracted_tags_text: + return f"{extracted_tags_text} {extra_text}" + return extra_text + + +def _build_metadata_payload_preview(meta: Dict[str, Any], extracted_tags_text: str, extras: list[str]) -> None: + # Keep this payload preparation for parity with previous behavior. + meta_fields = [extracted_tags_text] + extras + if meta.get("prompt"): + meta_fields.append(str(meta.get("prompt"))) + if meta.get("parameters"): + meta_fields.append(str(meta.get("parameters"))) + if meta.get("model"): + meta_fields.append(str(meta.get("model"))) + if meta.get("workflow_type"): + meta_fields.append(str(meta.get("workflow_type"))) + if meta.get("metadata_raw"): + meta_fields.append(str(meta.get("metadata_raw"))) + " ".join([str(f) for f in meta_fields if f]) + + +def _graph_has_sampler(graph: Any) -> bool: + try: + # Workflow export: dict with `nodes: []` and nodes have `type`. + if isinstance(graph, dict) and isinstance(graph.get("nodes"), list): + for node in graph.get("nodes") or []: + if _node_looks_like_sampler(node, require_sampling_inputs=False): + return True + return False + + # Prompt graph: dict of nodes with `class_type`. + if isinstance(graph, dict): + for node in graph.values(): + if _node_looks_like_sampler(node, require_sampling_inputs=True): + return True + return False + except Exception: + return False + return False + + +def _node_looks_like_sampler(node: Any, *, require_sampling_inputs: bool) -> bool: + if not isinstance(node, dict): + return False + ct = _sampler_class_type(node) + if not ct or not _sampler_class_matches(ct): + return False + if not require_sampling_inputs: + return True + return _sampler_inputs_look_sampling(node.get("inputs")) + + +def _sampler_class_type(node: Dict[str, Any]) -> str: + return str(node.get("class_type") or node.get("type") or "").lower() + + +def _sampler_class_matches(ct: str) -> bool: + if "select" in ct: + return False + if "ksampler" in ct or "samplercustom" in ct: + return True + return "sampler" in ct + + +def _sampler_inputs_look_sampling(inputs: Any) -> bool: + if not isinstance(inputs, dict): + return False + return any(key in inputs for key in ("steps", "cfg", "cfg_scale", "seed", "denoise")) + + +def _coerce_json_dict(value: Any) -> Optional[Dict[str, Any]]: + if isinstance(value, dict): + return value + if isinstance(value, str): + parsed = try_parse_json_text(value) + return parsed if isinstance(parsed, dict) else None + return None + + +def _metadata_presence_flags(meta: Dict[str, Any]) -> Tuple[bool, bool]: + raw_workflow = meta.get("workflow") + raw_prompt = meta.get("prompt") + workflow_obj = _coerce_json_dict(raw_workflow) + prompt_obj = _coerce_json_dict(raw_prompt) + + has_workflow = _has_workflow_payload(meta, raw_workflow, raw_prompt, workflow_obj, prompt_obj) + has_generation_data = _has_generation_payload(meta, raw_workflow, raw_prompt, workflow_obj, prompt_obj) + return has_workflow, has_generation_data + + +def _has_workflow_payload( + meta: Dict[str, Any], + raw_workflow: Any, + raw_prompt: Any, + workflow_obj: Optional[Dict[str, Any]], + prompt_obj: Optional[Dict[str, Any]], +) -> bool: + workflow_ok = bool(workflow_obj and looks_like_comfyui_workflow(workflow_obj)) + prompt_ok = bool(prompt_obj and looks_like_comfyui_prompt_graph(prompt_obj)) + return bool(workflow_ok or prompt_ok or raw_workflow or raw_prompt or meta.get("parameters")) + + +def _has_generation_payload( + meta: Dict[str, Any], + raw_workflow: Any, + raw_prompt: Any, + workflow_obj: Optional[Dict[str, Any]], + prompt_obj: Optional[Dict[str, Any]], +) -> bool: + if meta.get("parameters") or meta.get("geninfo") or meta.get("model") or meta.get("seed"): + return True + if _graph_has_sampler(prompt_obj or raw_prompt): + return True + return _graph_has_sampler(workflow_obj or raw_workflow) + + class MetadataHelpers: """ Shared utilities for metadata operations. @@ -52,87 +385,9 @@ class MetadataHelpers: metadata_quality = "none" metadata_raw_json = "{}" - def _graph_has_sampler(graph: Any) -> bool: - try: - # Workflow export: dict with `nodes: []` and nodes have `type`. - if isinstance(graph, dict) and isinstance(graph.get("nodes"), list): - for node in graph.get("nodes") or []: - if not isinstance(node, dict): - continue - ct = str(node.get("type") or node.get("class_type") or "").lower() - if not ct: - continue - if "ksampler" in ct and "select" not in ct: - return True - if "samplercustom" in ct: - return True - if "sampler" in ct and "select" not in ct: - return True - return False - - # Prompt graph: dict of nodes with `class_type`. - if isinstance(graph, dict): - for node in graph.values(): - if not isinstance(node, dict): - continue - ct = str(node.get("class_type") or node.get("type") or "").lower() - if not ct: - continue - if "ksampler" in ct and "select" not in ct: - return True - if "samplercustom" in ct: - return True - if "sampler" in ct and "select" not in ct: - # Require a couple of sampling-ish keys to avoid false positives. - ins = node.get("inputs") - if isinstance(ins, dict) and any(k in ins for k in ("steps", "cfg", "cfg_scale", "seed", "denoise")): - return True - return False - except Exception: - return False - return False - - def _coerce_json_dict(value: Any) -> Optional[Dict[str, Any]]: - if isinstance(value, dict): - return value - if isinstance(value, str): - parsed = try_parse_json_text(value) - return parsed if isinstance(parsed, dict) else None - return None - if metadata_result and metadata_result.ok and isinstance(metadata_result.data, dict) and metadata_result.data: meta = metadata_result.data - - raw_workflow = meta.get("workflow") - raw_prompt = meta.get("prompt") - - workflow_obj = _coerce_json_dict(raw_workflow) - prompt_obj = _coerce_json_dict(raw_prompt) - - workflow_ok = bool(workflow_obj and looks_like_comfyui_workflow(workflow_obj)) - prompt_ok = bool(prompt_obj and looks_like_comfyui_prompt_graph(prompt_obj)) - - # Backward-compatible: if the extractor stored non-empty strings but we can't parse, - # still count it as "has_workflow" so Workflow-only filters don't hide the asset. - workflow_present = bool(raw_workflow) - prompt_present = bool(raw_prompt) - - has_workflow = bool( - workflow_ok or - prompt_ok or - workflow_present or - prompt_present or - meta.get("parameters") - ) - - has_generation_data = bool( - meta.get("parameters") or - meta.get("geninfo") or - meta.get("model") or - meta.get("seed") or - _graph_has_sampler(prompt_obj or raw_prompt) or - _graph_has_sampler(workflow_obj or raw_workflow) - ) + has_workflow, has_generation_data = _metadata_presence_flags(meta) metadata_quality = meta.get("quality", "none") metadata_raw_json = json.dumps(meta) @@ -198,153 +453,14 @@ class MetadataHelpers: db_has_workflow = MetadataHelpers._bool_to_db(has_workflow) db_has_generation = MetadataHelpers._bool_to_db(has_generation_data) - # Guard against pathological/huge JSON payloads; keep DB stable. - try: - max_bytes = int(MAX_METADATA_JSON_BYTES or 0) - except Exception: - max_bytes = 0 - truncated = False - original_bytes = 0 - if max_bytes and isinstance(metadata_raw_json, str): - try: - nbytes = len(metadata_raw_json.encode("utf-8", errors="ignore")) - except Exception: - nbytes = 0 - if nbytes > max_bytes: - truncated = True - original_bytes = int(nbytes) - try: - metadata_raw_json = json.dumps( - {"_truncated": True, "original_bytes": int(nbytes)}, - ensure_ascii=False, - separators=(",", ":"), - ) - except Exception: - metadata_raw_json = "{}" - if truncated: - try: - metadata_result.meta["truncated"] = True - metadata_result.meta["original_bytes"] = int(original_bytes) - metadata_result.meta["max_bytes"] = int(max_bytes) - except Exception: - pass - try: - fp = filepath - if not fp and metadata_result and isinstance(metadata_result.data, dict): - fp = metadata_result.data.get("filepath") - suffix = f" filepath={fp}" if isinstance(fp, str) and fp else "" - logger.warning( - "Metadata JSON truncated for asset_id=%s%s (bytes=%s > max=%s)", - asset_id, - suffix, - original_bytes, - max_bytes, - ) - except Exception: - pass - - extracted_rating = 0 - extracted_tags_json = "[]" - extracted_tags_text = "" - if metadata_result and metadata_result.ok and metadata_result.data: - meta = metadata_result.data - try: - r = meta.get("rating") - if r is not None: - extracted_rating = max(0, min(5, int(r))) - except Exception: - extracted_rating = 0 - try: - t = meta.get("tags") - if isinstance(t, list): - cleaned = [] - seen = set() - for item in t: - if not isinstance(item, str): - continue - tag = item.strip() - if not tag or len(tag) > MAX_TAG_LENGTH or tag in seen: - continue - seen.add(tag) - cleaned.append(tag) - extracted_tags_json = json.dumps(cleaned, ensure_ascii=False) - extracted_tags_text = " ".join(cleaned) - except Exception: - extracted_tags_json = "[]" - extracted_tags_text = "" - - # Enrich tags_text with GenInfo (Models, LoRAs, Prompts) for FTS - if metadata_result and metadata_result.ok and metadata_result.data: - meta = metadata_result.data - extras = [] - geninfo = meta.get("geninfo") - if isinstance(geninfo, dict): - # Models - models = geninfo.get("models", {}) - if isinstance(models, dict): - for m_key, m_val in models.items(): - if isinstance(m_val, dict) and "name" in m_val: - n = str(m_val["name"]).strip() - if n: extras.append(n) - - # LoRAs - loras = geninfo.get("loras", []) - if isinstance(loras, list): - for lora in loras: - if isinstance(lora, dict) and "name" in lora: - n = str(lora["name"]).strip() - if n: extras.append(n) - - # Positive Prompt - pos_obj = geninfo.get("positive") - if isinstance(pos_obj, dict): - p_val = pos_obj.get("value") - if isinstance(p_val, str) and p_val.strip(): - extras.append(p_val.strip()) - - # Workflow Type (T2I, I2V, etc.) - engine = geninfo.get("engine") - if isinstance(engine, dict): - wf_type = engine.get("type") - if isinstance(wf_type, str) and wf_type.strip(): - extras.append(wf_type.strip()) # e.g. "I2V", "T2I" - - # Input files (to find assets using a specific source) - inputs = geninfo.get("inputs") - if isinstance(inputs, list): - for inp in inputs: - if isinstance(inp, dict): - fname = inp.get("filename") - if isinstance(fname, str) and fname.strip(): - extras.append(fname.strip()) - - role = inp.get("role") - if isinstance(role, str) and role.strip() and role not in ("secondary",): - # Index roles too? Maybe overkill to add "first_frame" to index, - # but "mask" or "source" could be useful. - # Let's keep it simple for now and just index the filenames. - pass - - if extras: - if extracted_tags_text: - extracted_tags_text += " " + " ".join(extras) - else: - extracted_tags_text = " ".join(extras) - - # Concat all metadata fields for metadata_text - meta_fields = [extracted_tags_text] + extras - # Add prompts, parameters, model, workflow, etc. from metadata_raw - if meta.get("prompt"): - meta_fields.append(str(meta.get("prompt"))) - if meta.get("parameters"): - meta_fields.append(str(meta.get("parameters"))) - if meta.get("model"): - meta_fields.append(str(meta.get("model"))) - if meta.get("workflow_type"): - meta_fields.append(str(meta.get("workflow_type"))) - if meta.get("metadata_raw"): - meta_fields.append(str(meta.get("metadata_raw"))) - " ".join([str(f) for f in meta_fields if f]) + metadata_raw_json = _apply_metadata_json_size_guard( + asset_id, + metadata_result, + metadata_raw_json, + filepath=filepath, + ) + extracted_rating, extracted_tags_json, extracted_tags_text = _extract_rating_and_tags(metadata_result) + extracted_tags_text = _enrich_tags_text_with_metadata(metadata_result, extracted_tags_text) # Import existing OS/file metadata when DB has defaults, without overriding user edits. # - rating: only set if current rating is 0 @@ -353,9 +469,7 @@ class MetadataHelpers: # We keep the best known metadata_quality for a file unless a newer extraction # yields an equal-or-better quality. This prevents counters/progress from # "going backwards" during background enrichment retries. - quality_rank_excluded = "CASE excluded.metadata_quality WHEN 'full' THEN 3 WHEN 'partial' THEN 2 WHEN 'degraded' THEN 1 ELSE 0 END" - quality_rank_current = "CASE COALESCE(asset_metadata.metadata_quality, 'none') WHEN 'full' THEN 3 WHEN 'partial' THEN 2 WHEN 'degraded' THEN 1 ELSE 0 END" - should_upgrade = f"({quality_rank_excluded} >= {quality_rank_current})" + should_upgrade = _metadata_quality_should_upgrade_sql() return await db.aexecute( """ @@ -501,18 +615,7 @@ class MetadataHelpers: @staticmethod async def _maybe_cleanup_metadata_cache(db: Sqlite) -> None: global _METADATA_CACHE_LAST_CLEANUP - try: - max_entries = int(METADATA_CACHE_MAX or 0) - except Exception: - max_entries = 0 - try: - ttl_seconds = float(METADATA_CACHE_TTL_SECONDS or 0) - except Exception: - ttl_seconds = 0.0 - try: - interval = float(METADATA_CACHE_CLEANUP_INTERVAL_SECONDS or 0) - except Exception: - interval = 0.0 + max_entries, ttl_seconds, interval = _cache_cleanup_config() if max_entries <= 0 and ttl_seconds <= 0: return @@ -528,89 +631,23 @@ class MetadataHelpers: _METADATA_CACHE_LAST_CLEANUP = now if ttl_seconds > 0: - try: - cutoff = datetime.now(timezone.utc) - timedelta(seconds=ttl_seconds) - cutoff_str = cutoff.strftime("%Y-%m-%d %H:%M:%S") - await db.aexecute( - "DELETE FROM metadata_cache WHERE last_updated < ?", - (cutoff_str,), - ) - except Exception: - pass + await _cleanup_cache_by_ttl(db, ttl_seconds) if max_entries > 0: - try: - count_res = await db.aquery("SELECT COUNT(1) AS count FROM metadata_cache") - if count_res.ok and count_res.data: - total = int(count_res.data[0].get("count") or 0) - else: - total = 0 - if total > max_entries: - to_remove = max(0, total - max_entries) - if to_remove: - await db.aexecute( - """ - DELETE FROM metadata_cache - WHERE filepath IN ( - SELECT filepath - FROM metadata_cache - ORDER BY last_updated ASC - LIMIT ? - ) - """, - (to_remove,), - ) - except Exception: - pass + await _cleanup_cache_by_max_entries(db, max_entries) @staticmethod async def store_metadata_cache(db: Sqlite, filepath: str, state_hash: str, metadata_result: Result[Dict[str, Any]]) -> Result[Any]: """ Store metadata in cache for future use. - - Args: - db: Database adapter instance - filepath: File path - state_hash: State hash for validation - metadata_result: Metadata result to cache - - Returns: - Result from database operation """ if not metadata_result.ok or not metadata_result.data: return Result.Err("CACHE_SKIPPED", "No metadata to cache") - metadata_raw = json.dumps( - metadata_result.data, - ensure_ascii=False, - separators=(",", ":"), - sort_keys=True - ) - try: - max_bytes = int(MAX_METADATA_JSON_BYTES or 0) - except Exception: - max_bytes = 0 - if max_bytes and isinstance(metadata_raw, str): - try: - nbytes = len(metadata_raw.encode("utf-8", errors="ignore")) - except Exception: - nbytes = 0 - if nbytes > max_bytes: - try: - logger.warning( - "Metadata cache skipped (payload too large) for %s (bytes=%s > max=%s)", - filepath, - int(nbytes), - int(max_bytes), - ) - except Exception: - pass - return Result.Err( - "CACHE_SKIPPED", - "Metadata JSON too large to cache", - original_bytes=int(nbytes), - max_bytes=int(max_bytes), - ) + metadata_raw = MetadataHelpers._metadata_json_payload(metadata_result.data) + size_check = MetadataHelpers._metadata_payload_size_guard(metadata_raw, filepath) + if size_check is not None: + return size_check metadata_hash = MetadataHelpers.compute_metadata_hash(metadata_raw) write_res = await db.aexecute( @@ -634,33 +671,114 @@ class MetadataHelpers: pass return write_res + @staticmethod + def _metadata_json_payload(metadata: Dict[str, Any]) -> str: + return json.dumps( + metadata, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ) + + @staticmethod + def _metadata_payload_size_guard(metadata_raw: str, filepath: str) -> Optional[Result[Any]]: + max_bytes = MetadataHelpers._max_metadata_json_bytes() + if not max_bytes: + return None + nbytes = MetadataHelpers._metadata_payload_size_bytes(metadata_raw) + if nbytes <= max_bytes: + return None + MetadataHelpers._log_metadata_payload_too_large(filepath, nbytes, max_bytes) + return Result.Err( + "CACHE_SKIPPED", + "Metadata JSON too large to cache", + original_bytes=int(nbytes), + max_bytes=int(max_bytes), + ) + + @staticmethod + def _max_metadata_json_bytes() -> int: + try: + return int(MAX_METADATA_JSON_BYTES or 0) + except Exception: + return 0 + + @staticmethod + def _metadata_payload_size_bytes(metadata_raw: str) -> int: + try: + return len(metadata_raw.encode("utf-8", errors="ignore")) + except Exception: + return 0 + + @staticmethod + def _log_metadata_payload_too_large(filepath: str, nbytes: int, max_bytes: int) -> None: + try: + logger.warning( + "Metadata cache skipped (payload too large) for %s (bytes=%s > max=%s)", + filepath, + int(nbytes), + int(max_bytes), + ) + except Exception: + pass + @staticmethod def compute_metadata_hash(raw_json: str) -> str: - """ - Compute hash of metadata JSON. - - Args: - raw_json: JSON string to hash - - Returns: - MD5 hash of the JSON - """ return hashlib.md5(raw_json.encode("utf-8")).hexdigest() @staticmethod async def set_metadata_value(db: Sqlite, key: str, value: str) -> Result[Any]: - """ - Persist a simple key/value in the metadata table. - - Args: - db: Database adapter instance - key: Metadata key - value: Metadata value - - Returns: - Result from database operation - """ return await db.aexecute( "INSERT OR REPLACE INTO metadata (key, value) VALUES (?, ?)", (key, value) ) + + +def _cache_cleanup_config() -> tuple[int, float, float]: + try: + max_entries = int(METADATA_CACHE_MAX or 0) + except Exception: + max_entries = 0 + try: + ttl_seconds = float(METADATA_CACHE_TTL_SECONDS or 0) + except Exception: + ttl_seconds = 0.0 + try: + interval = float(METADATA_CACHE_CLEANUP_INTERVAL_SECONDS or 0) + except Exception: + interval = 0.0 + return max_entries, ttl_seconds, interval + + +async def _cleanup_cache_by_ttl(db: Sqlite, ttl_seconds: float) -> None: + try: + cutoff = datetime.now(timezone.utc) - timedelta(seconds=ttl_seconds) + cutoff_str = cutoff.strftime("%Y-%m-%d %H:%M:%S") + await db.aexecute("DELETE FROM metadata_cache WHERE last_updated < ?", (cutoff_str,)) + except Exception: + return + + +async def _cleanup_cache_by_max_entries(db: Sqlite, max_entries: int) -> None: + try: + count_res = await db.aquery("SELECT COUNT(1) AS count FROM metadata_cache") + total = int(count_res.data[0].get("count") or 0) if count_res.ok and count_res.data else 0 + if total <= max_entries: + return + to_remove = max(0, total - max_entries) + if not to_remove: + return + await db.aexecute( + """ + DELETE FROM metadata_cache + WHERE filepath IN ( + SELECT filepath + FROM metadata_cache + ORDER BY last_updated ASC + LIMIT ? + ) + """, + (to_remove,), + ) + except Exception: + return diff --git a/mjr_am_backend/features/index/scanner.py b/mjr_am_backend/features/index/scanner.py index 6e60dce7..5d9788a9 100644 --- a/mjr_am_backend/features/index/scanner.py +++ b/mjr_am_backend/features/index/scanner.py @@ -144,37 +144,59 @@ class IndexScanner: Best-effort diagnosis for batch transaction failures. Returns (filepath, reason) where filepath may be None if unknown. """ + message, message_lower = self._batch_error_messages(batch_error) + if self._is_unique_filepath_error(message_lower): + diagnosed = self._diagnose_unique_filepath_error(prepared) + if diagnosed is not None: + return diagnosed + fp = self._first_prepared_filepath(prepared) + return fp, (message or type(batch_error).__name__) + + @staticmethod + def _batch_error_messages(batch_error: Exception) -> tuple[str, str]: try: message = str(batch_error or "") - message_lower = message.lower() except Exception: - message = "" - message_lower = "" + return "", "" + return message, message.lower() - # Common hard failure: duplicate filepath in the same batch payload. - if "unique constraint failed" in message_lower and "assets.filepath" in message_lower: - seen: set[str] = set() - for entry in prepared: - fp = str(entry.get("filepath") or entry.get("file_path") or "").strip() - if not fp: - continue - key = fp.lower() if IS_WINDOWS else fp - if key in seen: - return fp, "duplicate filepath in batch payload (UNIQUE assets.filepath)" - seen.add(key) - # Could also be a duplicate against existing DB rows. - for entry in prepared: - fp = str(entry.get("filepath") or entry.get("file_path") or "").strip() - if fp: - return fp, "filepath conflicts with existing database row (UNIQUE assets.filepath)" - return None, "UNIQUE constraint on assets.filepath" + @staticmethod + def _is_unique_filepath_error(message_lower: str) -> bool: + return "unique constraint failed" in message_lower and "assets.filepath" in message_lower - # Generic fallback: provide first actionable entry. + def _diagnose_unique_filepath_error(self, prepared: List[Dict[str, Any]]) -> Optional[tuple[Optional[str], str]]: + duplicate = self._first_duplicate_filepath_in_batch(prepared) + if duplicate: + return duplicate, "duplicate filepath in batch payload (UNIQUE assets.filepath)" + fp = self._first_prepared_filepath(prepared) + if fp: + return fp, "filepath conflicts with existing database row (UNIQUE assets.filepath)" + return None, "UNIQUE constraint on assets.filepath" + + @staticmethod + def _first_duplicate_filepath_in_batch(prepared: List[Dict[str, Any]]) -> Optional[str]: + seen: set[str] = set() for entry in prepared: - fp = str(entry.get("filepath") or entry.get("file_path") or "").strip() + fp = IndexScanner._prepared_filepath(entry) + if not fp: + continue + key = fp.lower() if IS_WINDOWS else fp + if key in seen: + return fp + seen.add(key) + return None + + @staticmethod + def _first_prepared_filepath(prepared: List[Dict[str, Any]]) -> Optional[str]: + for entry in prepared: + fp = IndexScanner._prepared_filepath(entry) if fp: - return fp, (message or type(batch_error).__name__) - return None, (message or type(batch_error).__name__) + return fp + return None + + @staticmethod + def _prepared_filepath(entry: Dict[str, Any]) -> str: + return str(entry.get("filepath") or entry.get("file_path") or "").strip() @staticmethod def _drain_walk_queue(q: "Queue[Optional[Path]]", max_items: int) -> list[Optional[Path]]: @@ -247,12 +269,9 @@ class IndexScanner: to_enrich: List[str] = [] async with self._scan_lock: dir_path = Path(directory) - - if not dir_path.exists(): - return Result.Err("DIR_NOT_FOUND", f"Directory not found: {directory}") - - if not dir_path.is_dir(): - return Result.Err("NOT_A_DIRECTORY", f"Not a directory: {directory}") + validation_error = self._validate_scan_directory(dir_path, directory) + if validation_error is not None: + return validation_error scan_id = str(uuid4()) self._current_scan_id = scan_id @@ -267,103 +286,20 @@ class IndexScanner: files_root=str(dir_path) ) - stats: Dict[str, Any] = { - "scanned": 0, - "added": 0, - "updated": 0, - "skipped": 0, - "errors": 0, - "batch_fallbacks": 0, - "skipped_state_changed": 0, - "start_time": datetime.now().isoformat() - } - - def _stream_batch_target(scanned_count: int) -> int: - try: - n = int(scanned_count or 0) - except (ValueError, TypeError): - n = 0 - if n <= 0: - return max(1, int(SCAN_BATCH_INITIAL)) - if n <= SCAN_BATCH_SMALL_THRESHOLD: - return max(int(SCAN_BATCH_MIN), int(SCAN_BATCH_SMALL)) - if n <= SCAN_BATCH_MED_THRESHOLD: - return max(int(SCAN_BATCH_MIN), int(SCAN_BATCH_MED)) - if n <= SCAN_BATCH_LARGE_THRESHOLD: - return max(int(SCAN_BATCH_MIN), int(SCAN_BATCH_LARGE)) - return max(int(SCAN_BATCH_MIN), int(SCAN_BATCH_XL)) + stats: Dict[str, Any] = self._new_scan_stats() try: - # IMPORTANT: os.walk() / Path.iterdir() are blocking and can freeze aiohttp. - # Walk the filesystem in a dedicated thread and consume results in batches. - loop = asyncio.get_running_loop() - stop_event = threading.Event() - q: "Queue[Optional[Path]]" = Queue(maxsize=max(1000, int(SCAN_BATCH_XL) * 4)) - - walk_future = loop.run_in_executor( - _FS_WALK_EXECUTOR, - self._walk_and_enqueue, - dir_path, - recursive, - stop_event, - q, + await self._run_scan_streaming_loop( + dir_path=dir_path, + directory=directory, + recursive=recursive, + incremental=incremental, + source=source, + root_id=root_id, + fast=fast, + stats=stats, + to_enrich=to_enrich, ) - - batch: List[Path] = [] - done = False - try: - while not done: - target = _stream_batch_target(stats["scanned"]) - pulled = await asyncio.to_thread(self._drain_walk_queue, q, target) - if not pulled: - # Yield and try again. - await asyncio.sleep(0) - continue - - for file_path in pulled: - if file_path is None: - done = True - break - - batch.append(file_path) - stats["scanned"] += 1 - - if len(batch) >= _stream_batch_target(stats["scanned"]): - await self._scan_stream_batch( - batch=batch, - base_dir=directory, - incremental=incremental, - source=source, - root_id=root_id, - fast=fast, - stats=stats, - to_enrich=to_enrich, - ) - batch = [] - await asyncio.sleep(0) - - if batch: - await self._scan_stream_batch( - batch=batch, - base_dir=directory, - incremental=incremental, - source=source, - root_id=root_id, - fast=fast, - stats=stats, - to_enrich=to_enrich, - ) - batch = [] - await asyncio.sleep(0) - except asyncio.CancelledError: - stop_event.set() - raise - finally: - stop_event.set() - try: - await asyncio.wait_for(walk_future, timeout=2.0) - except Exception: - pass finally: stats["end_time"] = datetime.now().isoformat() duration = time.perf_counter() - scan_start @@ -404,6 +340,165 @@ class IndexScanner: return Result.Ok(stats) + @staticmethod + def _new_scan_stats() -> Dict[str, Any]: + return { + "scanned": 0, + "added": 0, + "updated": 0, + "skipped": 0, + "errors": 0, + "batch_fallbacks": 0, + "skipped_state_changed": 0, + "start_time": datetime.now().isoformat(), + } + + @staticmethod + def _validate_scan_directory(dir_path: Path, directory: str) -> Optional[Result[Dict[str, Any]]]: + if not dir_path.exists(): + return Result.Err("DIR_NOT_FOUND", f"Directory not found: {directory}") + if not dir_path.is_dir(): + return Result.Err("NOT_A_DIRECTORY", f"Not a directory: {directory}") + return None + + @staticmethod + def _stream_batch_target(scanned_count: int) -> int: + try: + n = int(scanned_count or 0) + except (ValueError, TypeError): + n = 0 + if n <= 0: + return max(1, int(SCAN_BATCH_INITIAL)) + if n <= SCAN_BATCH_SMALL_THRESHOLD: + return max(int(SCAN_BATCH_MIN), int(SCAN_BATCH_SMALL)) + if n <= SCAN_BATCH_MED_THRESHOLD: + return max(int(SCAN_BATCH_MIN), int(SCAN_BATCH_MED)) + if n <= SCAN_BATCH_LARGE_THRESHOLD: + return max(int(SCAN_BATCH_MIN), int(SCAN_BATCH_LARGE)) + return max(int(SCAN_BATCH_MIN), int(SCAN_BATCH_XL)) + + async def _run_scan_streaming_loop( + self, + *, + dir_path: Path, + directory: str, + recursive: bool, + incremental: bool, + source: str, + root_id: Optional[str], + fast: bool, + stats: Dict[str, Any], + to_enrich: List[str], + ) -> None: + loop = asyncio.get_running_loop() + stop_event = threading.Event() + q: "Queue[Optional[Path]]" = Queue(maxsize=max(1000, int(SCAN_BATCH_XL) * 4)) + walk_future = loop.run_in_executor( + _FS_WALK_EXECUTOR, + self._walk_and_enqueue, + dir_path, + recursive, + stop_event, + q, + ) + try: + await self._consume_scan_queue( + q=q, + directory=directory, + incremental=incremental, + source=source, + root_id=root_id, + fast=fast, + stats=stats, + to_enrich=to_enrich, + stop_event=stop_event, + ) + finally: + stop_event.set() + try: + await asyncio.wait_for(walk_future, timeout=2.0) + except Exception: + pass + + async def _consume_scan_queue( + self, + *, + q: "Queue[Optional[Path]]", + directory: str, + incremental: bool, + source: str, + root_id: Optional[str], + fast: bool, + stats: Dict[str, Any], + to_enrich: List[str], + stop_event: threading.Event, + ) -> None: + batch: List[Path] = [] + done = False + try: + while not done: + target = self._stream_batch_target(stats["scanned"]) + pulled = await asyncio.to_thread(self._drain_walk_queue, q, target) + if not pulled: + await asyncio.sleep(0) + continue + for file_path in pulled: + if file_path is None: + done = True + break + batch.append(file_path) + stats["scanned"] += 1 + if len(batch) >= self._stream_batch_target(stats["scanned"]): + await self._process_scan_batch( + batch=batch, + directory=directory, + incremental=incremental, + source=source, + root_id=root_id, + fast=fast, + stats=stats, + to_enrich=to_enrich, + ) + batch = [] + await asyncio.sleep(0) + if batch: + await self._process_scan_batch( + batch=batch, + directory=directory, + incremental=incremental, + source=source, + root_id=root_id, + fast=fast, + stats=stats, + to_enrich=to_enrich, + ) + except asyncio.CancelledError: + stop_event.set() + raise + + async def _process_scan_batch( + self, + *, + batch: List[Path], + directory: str, + incremental: bool, + source: str, + root_id: Optional[str], + fast: bool, + stats: Dict[str, Any], + to_enrich: List[str], + ) -> None: + await self._scan_stream_batch( + batch=batch, + base_dir=directory, + incremental=incremental, + source=source, + root_id=root_id, + fast=fast, + stats=stats, + to_enrich=to_enrich, + ) + async def _scan_stream_batch( self, *, @@ -421,19 +516,7 @@ class IndexScanner: filepaths = [self._normalize_filepath_str(p) for p in batch] journal_map = (await self._get_journal_entries(filepaths)) if incremental and filepaths else {} - existing_map: Dict[str, Dict[str, Any]] = {} - - if filepaths: - existing_rows = await self.db.aquery_in( - "SELECT filepath, id, mtime FROM assets WHERE {IN_CLAUSE}", - "filepath", - filepaths, - ) - if existing_rows.ok and existing_rows.data and len(existing_rows.data) > 0: - for row in existing_rows.data: - fp = row.get("filepath") - if fp: - existing_map[str(fp)] = row + existing_map = await self._existing_map_for_filepaths(filepaths) await self._index_batch( batch=batch, @@ -448,6 +531,9 @@ class IndexScanner: to_enrich=to_enrich, ) + async def _existing_map_for_filepaths(self, filepaths: List[str]) -> Dict[str, Dict[str, Any]]: + return await self._existing_map_for_batch(filepaths) + async def index_paths( self, paths: List[Path], @@ -470,111 +556,20 @@ class IndexScanner: Result with indexing statistics """ async with self._index_lock: - # Filter unsupported files to prevent indexing internal files (DBs, etc.) - filtered_paths = [] - for p in paths: - try: - ext = p.suffix.lower() if p.suffix else "" - if ext and _EXT_TO_KIND: - if _EXT_TO_KIND.get(ext, "unknown") != "unknown": - filtered_paths.append(p) - continue - if classify_file(str(p)) != "unknown": - filtered_paths.append(p) - except Exception: - pass - paths = filtered_paths - + paths = self._filter_indexable_paths(paths) if not paths: - return Result.Ok({ - "scanned": 0, "added": 0, "updated": 0, "skipped": 0, "errors": 0, - "start_time": datetime.now().isoformat(), - "end_time": datetime.now().isoformat() - }) + return Result.Ok(self._empty_index_stats()) scan_id = str(uuid4()) self._current_scan_id = scan_id scan_start = time.perf_counter() - - # Reduce log noise for single-file indexing (common in incremental updates) - start_log_level = logging.DEBUG if len(paths) == 1 else logging.INFO - self._log_scan_event( - start_log_level, - "Starting file list index", - file_count=len(paths), - base_dir=base_dir, - incremental=incremental - ) - - stats: Dict[str, Any] = { - "scanned": len(paths), - "added": 0, - "updated": 0, - "skipped": 0, - "errors": 0, - "batch_fallbacks": 0, - "skipped_state_changed": 0, - "start_time": datetime.now().isoformat() - } + self._log_index_paths_start(paths, base_dir, incremental) + stats = self._new_index_stats(len(paths)) added_ids: List[int] = [] try: - for batch in self._chunk_file_batches(paths): - if not batch: - continue - - filepaths = [self._normalize_filepath_str(p) for p in batch] - journal_map = (await self._get_journal_entries(filepaths)) if incremental and filepaths else {} - existing_map: Dict[str, Dict[str, Any]] = {} - - if filepaths: - existing_rows = await self.db.aquery_in( - "SELECT filepath, id, mtime FROM assets WHERE {IN_CLAUSE}", - "filepath", - filepaths, - ) - if existing_rows.ok and existing_rows.data and len(existing_rows.data) > 0: - for row in existing_rows.data: - fp = row.get("filepath") - if fp: - existing_map[str(fp)] = row - - await self._index_batch( - batch=batch, - base_dir=base_dir, - incremental=incremental, - source=source, - root_id=root_id, - fast=False, - journal_map=journal_map, - existing_map=existing_map, - stats=stats, - to_enrich=None, - added_ids=added_ids, - ) - await asyncio.sleep(0) + await self._index_paths_batches(paths, base_dir, incremental, source, root_id, stats, added_ids) finally: - stats["end_time"] = datetime.now().isoformat() - duration = time.perf_counter() - scan_start - - # Do not update last_scan_end for targeted indexing (DnD, executed events, etc.). - # The UI uses last_scan_end to detect *full scans* and reload the grid; updating it - # for single-file indexing causes unnecessary grid reloads/flicker. - await MetadataHelpers.set_metadata_value(self.db, "last_index_end", stats["end_time"]) - - # Only log completion at INFO if meaningful changes occurred - has_changes = stats["added"] > 0 or stats["updated"] > 0 or stats["errors"] > 0 - complete_log_level = logging.INFO if has_changes else logging.DEBUG - - self._log_scan_event( - complete_log_level, - "File list index complete", - duration_seconds=duration, - scanned=stats["scanned"], - added=stats["added"], - updated=stats["updated"], - skipped=stats["skipped"], - errors=stats["errors"] - ) + await self._finalize_index_paths(scan_start, stats) self._current_scan_id = None # Include added/updated asset IDs for WebSocket notification @@ -582,6 +577,127 @@ class IndexScanner: stats["added_ids"] = added_ids return Result.Ok(stats) + @staticmethod + def _filter_indexable_paths(paths: List[Path]) -> List[Path]: + filtered_paths: List[Path] = [] + for p in paths: + try: + ext = p.suffix.lower() if p.suffix else "" + if ext and _EXT_TO_KIND and _EXT_TO_KIND.get(ext, "unknown") != "unknown": + filtered_paths.append(p) + continue + if classify_file(str(p)) != "unknown": + filtered_paths.append(p) + except Exception: + continue + return filtered_paths + + @staticmethod + def _empty_index_stats() -> Dict[str, Any]: + now = datetime.now().isoformat() + return { + "scanned": 0, + "added": 0, + "updated": 0, + "skipped": 0, + "errors": 0, + "start_time": now, + "end_time": now, + } + + def _log_index_paths_start(self, paths: List[Path], base_dir: str, incremental: bool) -> None: + start_log_level = logging.DEBUG if len(paths) == 1 else logging.INFO + self._log_scan_event( + start_log_level, + "Starting file list index", + file_count=len(paths), + base_dir=base_dir, + incremental=incremental, + ) + + @staticmethod + def _new_index_stats(scanned: int) -> Dict[str, Any]: + return { + "scanned": scanned, + "added": 0, + "updated": 0, + "skipped": 0, + "errors": 0, + "batch_fallbacks": 0, + "skipped_state_changed": 0, + "start_time": datetime.now().isoformat(), + } + + async def _index_paths_batches( + self, + paths: List[Path], + base_dir: str, + incremental: bool, + source: str, + root_id: Optional[str], + stats: Dict[str, Any], + added_ids: List[int], + ) -> None: + for batch in self._chunk_file_batches(paths): + if not batch: + continue + filepaths = [self._normalize_filepath_str(p) for p in batch] + journal_map = await self._journal_map_for_batch(filepaths, incremental) + existing_map = await self._existing_map_for_batch(filepaths) + await self._index_batch( + batch=batch, + base_dir=base_dir, + incremental=incremental, + source=source, + root_id=root_id, + fast=False, + journal_map=journal_map, + existing_map=existing_map, + stats=stats, + to_enrich=None, + added_ids=added_ids, + ) + await asyncio.sleep(0) + + async def _journal_map_for_batch(self, filepaths: List[str], incremental: bool) -> Dict[str, Dict[str, Any]]: + if not incremental or not filepaths: + return {} + return await self._get_journal_entries(filepaths) + + async def _existing_map_for_batch(self, filepaths: List[str]) -> Dict[str, Dict[str, Any]]: + existing_map: Dict[str, Dict[str, Any]] = {} + if not filepaths: + return existing_map + existing_rows = await self.db.aquery_in( + "SELECT filepath, id, mtime FROM assets WHERE {IN_CLAUSE}", + "filepath", + filepaths, + ) + if not existing_rows.ok or not existing_rows.data: + return existing_map + for row in existing_rows.data: + fp = row.get("filepath") + if fp: + existing_map[str(fp)] = row + return existing_map + + async def _finalize_index_paths(self, scan_start: float, stats: Dict[str, Any]) -> None: + stats["end_time"] = datetime.now().isoformat() + duration = time.perf_counter() - scan_start + await MetadataHelpers.set_metadata_value(self.db, "last_index_end", stats["end_time"]) + has_changes = stats["added"] > 0 or stats["updated"] > 0 or stats["errors"] > 0 + complete_log_level = logging.INFO if has_changes else logging.DEBUG + self._log_scan_event( + complete_log_level, + "File list index complete", + duration_seconds=duration, + scanned=stats["scanned"], + added=stats["added"], + updated=stats["updated"], + skipped=stats["skipped"], + errors=stats["errors"], + ) + async def _index_batch( self, batch: List[Path], @@ -596,192 +712,440 @@ class IndexScanner: to_enrich: Optional[List[str]] = None, added_ids: Optional[List[int]] = None, ) -> None: - """ - Index a batch of files using a single transaction for all DB writes. - - This drastically reduces SQLite transaction overhead on large scans (10k+ files) - while keeping the "do not crash UI" and Result patterns intact. - """ - + # Index a batch of files using one DB transaction for writes. batch_start = time.perf_counter() - - # Phase 1: Stat files and determine which need metadata extraction - prepared: List[Dict[str, Any]] = [] - needs_metadata: List[tuple[Path, str, int, int, int, str, Optional[int]]] = [] # (path, filepath, mtime_ns, mtime, size, state_hash, existing_id) - - # Prefetch metadata cache and asset_metadata entries for the entire batch to avoid N+1 queries filepaths = [self._normalize_filepath_str(p) for p in batch] - cache_map = {} - has_rich_meta_set = set() - if filepaths: - # Prefetch metadata_cache entries - cache_rows = await self.db.aquery_in( - "SELECT filepath, state_hash, metadata_raw FROM metadata_cache WHERE {IN_CLAUSE}", - "filepath", - filepaths, + cache_map, has_rich_meta_set = await self._prefetch_batch_cache_and_rich_meta( + filepaths, + existing_map, + ) + prepared, needs_metadata = await self._prepare_batch_entries( + batch=batch, + base_dir=base_dir, + incremental=incremental, + fast=fast, + journal_map=journal_map, + existing_map=existing_map, + cache_map=cache_map, + has_rich_meta_set=has_rich_meta_set, + stats=stats, + ) + await self._append_batch_metadata_entries( + prepared=prepared, + needs_metadata=needs_metadata, + base_dir=base_dir, + fast=fast, + ) + + if not prepared: + return + + await self._persist_prepared_entries( + prepared=prepared, + base_dir=base_dir, + source=source, + root_id=root_id, + stats=stats, + to_enrich=to_enrich, + added_ids=added_ids, + ) + + duration = time.perf_counter() - batch_start + if duration > 0.2: + logger.debug( + "_index_batch timing: %.3fs (batch=%s)", + duration, + len(batch) if batch is not None else 0, ) - if cache_rows.ok and cache_rows.data: - for row in cache_rows.data: - fp = row.get("filepath") - state_hash = row.get("state_hash") - if fp and state_hash: - cache_map[(fp, state_hash)] = row.get("metadata_raw") - # Prefetch asset_metadata entries - asset_ids = [] - for fp in filepaths: - existing_row = existing_map.get(fp) - if existing_row and existing_row.get("id"): - try: - existing_id = int(existing_row.get("id") or 0) - if existing_id: - asset_ids.append(existing_id) - except (ValueError, TypeError): - pass + async def _prefetch_batch_cache_and_rich_meta( + self, + filepaths: List[str], + existing_map: Dict[str, Dict[str, Any]], + ) -> tuple[Dict[tuple[str, str], Any], set[int]]: + cache_map: Dict[tuple[str, str], Any] = {} + has_rich_meta_set: set[int] = set() + if not filepaths: + return cache_map, has_rich_meta_set - if asset_ids: - meta_rows = await self.db.aquery_in( - "SELECT asset_id, metadata_quality, metadata_raw FROM asset_metadata WHERE {IN_CLAUSE}", - "asset_id", - asset_ids, - ) - if meta_rows.ok and meta_rows.data: - for row in meta_rows.data: - asset_id = row.get("asset_id") - if asset_id: - try: - metadata_quality = str(row.get("metadata_quality") or "").strip().lower() - metadata_raw = str(row.get("metadata_raw") or "").strip() - has_rich = ( - metadata_quality not in ("", "none") - or metadata_raw not in ("", "{}", "null") - ) - if has_rich: - has_rich_meta_set.add(int(asset_id)) - except Exception: - pass + await self._prefetch_metadata_cache_rows(filepaths, cache_map) + asset_ids = self._asset_ids_from_existing_rows(filepaths, existing_map) + if not asset_ids: + return cache_map, has_rich_meta_set - for file_path in batch: - fp = self._normalize_filepath_str(file_path) - existing_state: Optional[Dict[str, Any]] - if incremental and fp in journal_map: - existing_state = dict(existing_map.get(fp) or {}) - existing_state["journal_state_hash"] = journal_map.get(fp) - else: - existing_state = existing_map.get(fp) - if IS_WINDOWS and not existing_state: - try: - existing_ci = await self.db.aquery( - "SELECT id, mtime, filepath FROM assets WHERE filepath = ? COLLATE NOCASE LIMIT 1", - (fp,), - ) - if existing_ci.ok and existing_ci.data: - existing_state = existing_ci.data[0] - existing_map[fp] = existing_state - except Exception: - pass + await self._prefetch_rich_metadata_rows(asset_ids, has_rich_meta_set) + return cache_map, has_rich_meta_set - # Stat the file - stat_result = await self._stat_with_retry(file_path) - if not stat_result[0]: - stats["errors"] += 1 - logger.warning("Failed to stat %s: %s", str(file_path), stat_result[1]) + async def _prefetch_metadata_cache_rows( + self, + filepaths: List[str], + cache_map: Dict[tuple[str, str], Any], + ) -> None: + cache_rows = await self.db.aquery_in( + "SELECT filepath, state_hash, metadata_raw FROM metadata_cache WHERE {IN_CLAUSE}", + "filepath", + filepaths, + ) + if not cache_rows.ok or not cache_rows.data: + return + for row in cache_rows.data: + fp = row.get("filepath") + state_hash = row.get("state_hash") + if fp and state_hash: + cache_map[(str(fp), str(state_hash))] = row.get("metadata_raw") + + @staticmethod + def _asset_ids_from_existing_rows(filepaths: List[str], existing_map: Dict[str, Dict[str, Any]]) -> List[int]: + asset_ids: List[int] = [] + for fp in filepaths: + existing_row = existing_map.get(fp) + if not existing_row or not existing_row.get("id"): + continue + try: + existing_id = int(existing_row.get("id") or 0) + except (TypeError, ValueError): + continue + if existing_id: + asset_ids.append(existing_id) + return asset_ids + + async def _prefetch_rich_metadata_rows(self, asset_ids: List[int], has_rich_meta_set: set[int]) -> None: + meta_rows = await self.db.aquery_in( + "SELECT asset_id, metadata_quality, metadata_raw FROM asset_metadata WHERE {IN_CLAUSE}", + "asset_id", + asset_ids, + ) + if not meta_rows.ok or not meta_rows.data: + return + for row in meta_rows.data: + asset_id = row.get("asset_id") + if not asset_id: + continue + try: + metadata_quality = str(row.get("metadata_quality") or "").strip().lower() + metadata_raw = str(row.get("metadata_raw") or "").strip() + has_rich = metadata_quality not in ("", "none") or metadata_raw not in ("", "{}", "null") + if has_rich: + has_rich_meta_set.add(int(asset_id)) + except Exception: continue - stat = stat_result[1] - mtime_ns = stat.st_mtime_ns - mtime = int(stat.st_mtime) - size = int(stat.st_size) - filepath = self._normalize_filepath_str(file_path) - state_hash = self._compute_state_hash(filepath, mtime_ns, size) + async def _prepare_batch_entries( + self, + batch: List[Path], + base_dir: str, + incremental: bool, + fast: bool, + journal_map: Dict[str, str], + existing_map: Dict[str, Dict[str, Any]], + cache_map: Dict[tuple[str, str], Any], + has_rich_meta_set: set[int], + stats: Dict[str, Any], + ) -> tuple[List[Dict[str, Any]], List[tuple[Path, str, int, int, int, str, Optional[int]]]]: + prepared: List[Dict[str, Any]] = [] + needs_metadata: List[tuple[Path, str, int, int, int, str, Optional[int]]] = [] - # Check journal skip - journal_state_hash = None - if isinstance(existing_state, dict) and "journal_state_hash" in existing_state: - journal_state_hash = existing_state.get("journal_state_hash") + for file_path in batch: + prepared_entry, metadata_item = await self._prepare_single_batch_entry( + file_path=file_path, + base_dir=base_dir, + incremental=incremental, + fast=fast, + journal_map=journal_map, + existing_map=existing_map, + cache_map=cache_map, + has_rich_meta_set=has_rich_meta_set, + stats=stats, + ) + if prepared_entry is not None: + prepared.append(prepared_entry) + if metadata_item is not None: + needs_metadata.append(metadata_item) - existing_id = 0 - existing_mtime = 0 - if isinstance(existing_state, dict) and existing_state.get("id") is not None: - try: - existing_id = int(existing_state.get("id") or 0) - existing_mtime = int(existing_state.get("mtime") or 0) - except Exception: - pass + return prepared, needs_metadata - # Journal fast-path is safe only when we are not asking for richer metadata. - if incremental and journal_state_hash and str(journal_state_hash) == state_hash: - if fast or (existing_id and existing_id in has_rich_meta_set): - prepared.append({"action": "skipped_journal"}) - continue + def _build_fast_batch_entry( + self, + *, + existing_id: int, + file_path: Path, + base_dir: str, + filepath: str, + mtime_ns: int, + mtime: int, + size: int, + state_hash: str, + ) -> Dict[str, Any]: + if existing_id: + return self._build_updated_entry( + asset_id=existing_id, + metadata_result=Result.Ok({}), + filepath=filepath, + file_path=file_path, + state_hash=state_hash, + mtime=mtime, + size=size, + fast=True, + cache_store=False, + mtime_ns=mtime_ns, + ) - # Check cache skip using prefetched data - # Optimization: Skip processing if mtime matches, even for non-incremental scans. - # This prevents massive redundant DB updates when performing full directory scans. - if existing_id and existing_mtime == mtime: - # Check if we have cached metadata for this filepath and state_hash - cached_raw = cache_map.get((filepath, state_hash)) - if cached_raw: - cached_result = Result.Ok({"metadata_raw": cached_raw}, source="cache") - prepared.append({ - "action": "refresh", - "asset_id": existing_id, - "metadata_result": cached_result, - "filepath": filepath, - "file_path": file_path, - "state_hash": state_hash, - "mtime": mtime, - "size": size, - "fast": fast, - "cache_store": False, - }) - continue + rel_path = self._safe_relative_path(file_path, base_dir) + return self._build_added_entry( + filename=file_path.name, + subfolder=str(rel_path.parent) if rel_path.parent != Path(".") else "", + filepath=filepath, + kind=classify_file(file_path.name), + metadata_result=Result.Ok({}), + file_path=file_path, + state_hash=state_hash, + mtime=mtime, + size=size, + fast=True, + cache_store=False, + mtime_ns=mtime_ns, + ) - # Check if asset has metadata using prefetched data - if existing_id in has_rich_meta_set: - prepared.append({"action": "skipped"}) - continue + @staticmethod + def _batch_stat_to_values(stat_obj: os.stat_result) -> tuple[int, int, int]: + return int(stat_obj.st_mtime_ns), int(stat_obj.st_mtime), int(stat_obj.st_size) - # This file needs processing - if not fast: - needs_metadata.append((file_path, filepath, int(mtime_ns), mtime, size, state_hash, existing_id if existing_id else None)) - else: - # Fast mode - no metadata - rel_path = self._safe_relative_path(file_path, base_dir) - prepared.append({ - "action": "updated" if existing_id else "added", - "asset_id": existing_id if existing_id else None, - "metadata_result": Result.Ok({}), - "filepath": filepath, - "file_path": file_path, - "mtime_ns": int(mtime_ns), - "state_hash": state_hash, - "mtime": mtime, - "size": size, - "fast": True, - "cache_store": False, - "filename": file_path.name if not existing_id else None, - "subfolder": str(rel_path.parent) if not existing_id and rel_path.parent != Path(".") else ("" if not existing_id else None), - "kind": classify_file(file_path.name) if not existing_id else None, - }) + @staticmethod + def _should_skip_by_journal( + *, + incremental: bool, + journal_state_hash: Optional[str], + state_hash: str, + fast: bool, + existing_id: int, + has_rich_meta_set: set[int], + ) -> bool: + if not (incremental and journal_state_hash and str(journal_state_hash) == state_hash): + return False + return bool(fast or (existing_id and existing_id in has_rich_meta_set)) - # Phase 2: Batch metadata extraction - if needs_metadata: - paths_to_extract = [item[0] for item in needs_metadata] - batch_metadata = await self.metadata.get_metadata_batch([str(p) for p in paths_to_extract], scan_id=self._current_scan_id) + @staticmethod + def _cached_refresh_payload( + *, + existing_id: int, + existing_mtime: int, + mtime: int, + cache_map: Dict[tuple[str, str], Any], + filepath: str, + state_hash: str, + ) -> tuple[Optional[Any], bool]: + if not (existing_id and existing_mtime == mtime): + return None, False + cached_raw = cache_map.get((filepath, state_hash)) + if cached_raw: + return cached_raw, False + return None, True - for file_path, filepath, mtime_ns, mtime, size, state_hash, existing_id_opt in needs_metadata: - metadata_result: Result[Dict[str, Any]] | None = batch_metadata.get(str(file_path)) - if not metadata_result: - metadata_result = MetadataHelpers.metadata_error_payload(Result.Err("METADATA_MISSING", "No metadata returned"), filepath) - elif not metadata_result.ok: - metadata_result = MetadataHelpers.metadata_error_payload(metadata_result, filepath) + def _build_cached_refresh_entry( + self, + *, + existing_id: int, + cached_raw: Any, + filepath: str, + file_path: Path, + state_hash: str, + mtime: int, + size: int, + fast: bool, + ) -> Dict[str, Any]: + return self._build_refresh_entry( + asset_id=existing_id, + metadata_result=Result.Ok({"metadata_raw": cached_raw}, source="cache"), + filepath=filepath, + file_path=file_path, + state_hash=state_hash, + mtime=mtime, + size=size, + fast=fast, + cache_store=False, + ) - cache_store = metadata_result.ok + @staticmethod + def _journal_state_hash(existing_state: Optional[Dict[str, Any]]) -> Optional[str]: + if isinstance(existing_state, dict): + return existing_state.get("journal_state_hash") + return None - if existing_id_opt: - prepared.append({ + @staticmethod + def _metadata_queue_item( + file_path: Path, + filepath: str, + mtime_ns: int, + mtime: int, + size: int, + state_hash: str, + existing_id: int, + ) -> tuple[Path, str, int, int, int, str, Optional[int]]: + return ( + file_path, + filepath, + mtime_ns, + mtime, + size, + state_hash, + existing_id if existing_id else None, + ) + + async def _prepare_single_batch_entry( + self, + file_path: Path, + base_dir: str, + incremental: bool, + fast: bool, + journal_map: Dict[str, str], + existing_map: Dict[str, Dict[str, Any]], + cache_map: Dict[tuple[str, str], Any], + has_rich_meta_set: set[int], + stats: Dict[str, Any], + ) -> tuple[Optional[Dict[str, Any]], Optional[tuple[Path, str, int, int, int, str, Optional[int]]]]: + fp = self._normalize_filepath_str(file_path) + existing_state = await self._resolve_existing_state_for_batch( + fp=fp, + incremental=incremental, + journal_map=journal_map, + existing_map=existing_map, + ) + + stat_result = await self._stat_with_retry(file_path) + if not stat_result[0]: + stats["errors"] += 1 + logger.warning("Failed to stat %s: %s", str(file_path), stat_result[1]) + return None, None + + stat = stat_result[1] + mtime_ns, mtime, size = self._batch_stat_to_values(stat) + filepath = self._normalize_filepath_str(file_path) + state_hash = self._compute_state_hash(filepath, mtime_ns, size) + + journal_state_hash = self._journal_state_hash(existing_state if isinstance(existing_state, dict) else None) + + existing_id, existing_mtime = self._extract_existing_asset_state( + existing_state if isinstance(existing_state, dict) else None + ) + + if self._should_skip_by_journal( + incremental=incremental, + journal_state_hash=journal_state_hash, + state_hash=state_hash, + fast=fast, + existing_id=existing_id, + has_rich_meta_set=has_rich_meta_set, + ): + return {"action": "skipped_journal"}, None + + cached_raw, has_cached_miss = self._cached_refresh_payload( + existing_id=existing_id, + existing_mtime=existing_mtime, + mtime=mtime, + cache_map=cache_map, + filepath=filepath, + state_hash=state_hash, + ) + if cached_raw is not None: + return ( + self._build_cached_refresh_entry( + existing_id=existing_id, + cached_raw=cached_raw, + filepath=filepath, + file_path=file_path, + state_hash=state_hash, + mtime=mtime, + size=size, + fast=fast, + ), + None, + ) + if has_cached_miss and existing_id in has_rich_meta_set: + return {"action": "skipped"}, None + + if fast: + return ( + self._build_fast_batch_entry( + existing_id=existing_id, + file_path=file_path, + base_dir=base_dir, + filepath=filepath, + mtime_ns=mtime_ns, + mtime=mtime, + size=size, + state_hash=state_hash, + ), + None, + ) + + return None, self._metadata_queue_item( + file_path, + filepath, + mtime_ns, + mtime, + size, + state_hash, + existing_id, + ) + + async def _resolve_existing_state_for_batch( + self, + fp: str, + incremental: bool, + journal_map: Dict[str, str], + existing_map: Dict[str, Dict[str, Any]], + ) -> Optional[Dict[str, Any]]: + existing_state: Optional[Dict[str, Any]] + if incremental and fp in journal_map: + existing_state = dict(existing_map.get(fp) or {}) + existing_state["journal_state_hash"] = journal_map.get(fp) + else: + existing_state = existing_map.get(fp) + + if IS_WINDOWS and not existing_state: + try: + existing_ci = await self.db.aquery( + "SELECT id, mtime, filepath FROM assets WHERE filepath = ? COLLATE NOCASE LIMIT 1", + (fp,), + ) + if existing_ci.ok and existing_ci.data: + existing_state = existing_ci.data[0] + existing_map[fp] = existing_state + except Exception: + pass + return existing_state + + async def _append_batch_metadata_entries( + self, + prepared: List[Dict[str, Any]], + needs_metadata: List[tuple[Path, str, int, int, int, str, Optional[int]]], + base_dir: str, + fast: bool, + ) -> None: + if not needs_metadata: + return + + paths_to_extract = [item[0] for item in needs_metadata] + batch_metadata = await self.metadata.get_metadata_batch( + [str(p) for p in paths_to_extract], + scan_id=self._current_scan_id, + ) + + for file_path, filepath, mtime_ns, mtime, size, state_hash, existing_id_opt in needs_metadata: + metadata_result: Optional[Result[Dict[str, Any]]] = batch_metadata.get(str(file_path)) + if not metadata_result: + metadata_result = MetadataHelpers.metadata_error_payload( + Result.Err("METADATA_MISSING", "No metadata returned"), + filepath, + ) + elif not metadata_result.ok: + metadata_result = MetadataHelpers.metadata_error_payload(metadata_result, filepath) + + cache_store = bool(metadata_result.ok) + if existing_id_opt: + prepared.append( + { "action": "updated", "asset_id": existing_id_opt, "metadata_result": metadata_result, @@ -793,209 +1157,52 @@ class IndexScanner: "size": size, "fast": fast, "cache_store": cache_store, - }) - else: - rel_path = self._safe_relative_path(file_path, base_dir) - filename = file_path.name - subfolder = str(rel_path.parent) if rel_path.parent != Path(".") else "" - kind = classify_file(filename) + } + ) + continue - prepared.append({ - "action": "added", - "filename": filename, - "subfolder": subfolder, - "filepath": filepath, - "kind": kind, - "metadata_result": metadata_result, - "file_path": file_path, - "mtime_ns": int(mtime_ns), - "state_hash": state_hash, - "mtime": mtime, - "size": size, - "fast": fast, - "cache_store": cache_store, - }) + rel_path = self._safe_relative_path(file_path, base_dir) + filename = file_path.name + subfolder = str(rel_path.parent) if rel_path.parent != Path(".") else "" + kind = classify_file(filename) + prepared.append( + { + "action": "added", + "filename": filename, + "subfolder": subfolder, + "filepath": filepath, + "kind": kind, + "metadata_result": metadata_result, + "file_path": file_path, + "mtime_ns": int(mtime_ns), + "state_hash": state_hash, + "mtime": mtime, + "size": size, + "fast": fast, + "cache_store": cache_store, + } + ) - if not prepared: - return - - # Apply DB writes for the whole batch in one transaction. - # If the batch fails, fall back to processing items individually + async def _persist_prepared_entries( + self, + prepared: List[Dict[str, Any]], + base_dir: str, + source: str, + root_id: Optional[str], + stats: Dict[str, Any], + to_enrich: Optional[List[str]], + added_ids: Optional[List[int]], + ) -> None: try: - async with self.db.atransaction(mode="immediate") as tx: - if not tx.ok: - raise RuntimeError(tx.error or "Failed to begin transaction") - for entry in prepared: - action = entry.get("action") - - if action in ("skipped", "skipped_journal"): - stats["skipped"] += 1 - continue - file_path_value = entry.get("file_path") - if isinstance(file_path_value, Path): - expected_mtime_ns = int(entry.get("mtime_ns") or 0) - expected_size = int(entry.get("size") or 0) - if self._file_state_drifted(file_path_value, expected_mtime_ns, expected_size): - stats["skipped"] += 1 - stats["skipped_state_changed"] = int(stats.get("skipped_state_changed") or 0) + 1 - continue - - if action == "refresh": - asset_id = entry.get("asset_id") - metadata_result = entry.get("metadata_result") - if not asset_id or not isinstance(metadata_result, Result): - stats["skipped"] += 1 - continue - try: - refreshed = await MetadataHelpers.refresh_metadata_if_needed( - self.db, - int(asset_id), - metadata_result, - entry.get("filepath") or "", - base_dir, - entry.get("state_hash") or "", - int(entry.get("mtime") or 0), - int(entry.get("size") or 0), - self._write_scan_journal_entry, - ) - stats["skipped"] += 1 - if refreshed and entry.get("fast") and to_enrich is not None: - if len(to_enrich) < MAX_TO_ENRICH_ITEMS: - to_enrich.append(entry.get("filepath") or str(entry.get("file_path") or "")) - except Exception as exc: - stats["errors"] += 1 - logger.warning("Metadata refresh failed for asset_id=%s: %s", asset_id, exc) - continue - - if action == "updated": - asset_id = entry.get("asset_id") - metadata_result = entry.get("metadata_result") - if not asset_id or not isinstance(metadata_result, Result): - stats["errors"] += 1 - continue - file_path_value = entry.get("file_path") - if not isinstance(file_path_value, Path): - stats["errors"] += 1 - continue - cache_store = bool(entry.get("cache_store")) - if cache_store: - try: - await MetadataHelpers.store_metadata_cache( - self.db, - entry.get("filepath") or "", - entry.get("state_hash") or "", - metadata_result, - ) - except Exception: - pass - res = await self._update_asset( - int(asset_id), - file_path_value, - int(entry.get("mtime") or 0), - int(entry.get("size") or 0), - metadata_result, - source=source, - root_id=root_id, - write_metadata=not bool(entry.get("fast")), - ) - if res.ok: - await self._write_scan_journal_entry( - entry.get("filepath") or "", - base_dir, - entry.get("state_hash") or "", - int(entry.get("mtime") or 0), - int(entry.get("size") or 0), - ) - stats["updated"] += 1 - # Collect updated asset ID for WebSocket notification - if added_ids is not None and asset_id: - try: - added_ids.append(int(asset_id)) - except Exception: - pass - if entry.get("fast") and to_enrich is not None: - if len(to_enrich) < MAX_TO_ENRICH_ITEMS: - to_enrich.append(entry.get("filepath") or str(entry.get("file_path") or "")) - try: - logger.debug( - "Indexed file (updated): %s [asset_id=%s]", - entry.get("filepath") or "", - asset_id, - ) - except Exception: - pass - else: - stats["errors"] += 1 - continue - - if action == "added": - metadata_result = entry.get("metadata_result") - if not isinstance(metadata_result, Result): - stats["errors"] += 1 - continue - kind_value = entry.get("kind") - file_path_value = entry.get("file_path") - if not isinstance(kind_value, str) or not isinstance(file_path_value, Path): - stats["errors"] += 1 - continue - cache_store = bool(entry.get("cache_store")) - res = await self._add_asset( - entry.get("filename") or "", - entry.get("subfolder") or "", - entry.get("filepath") or "", - cast(FileKind, kind_value), - int(entry.get("mtime") or 0), - int(entry.get("size") or 0), - file_path_value, - metadata_result, - source=source, - root_id=root_id, - write_metadata=True, - ) - if res.ok: - if cache_store: - try: - await MetadataHelpers.store_metadata_cache( - self.db, - entry.get("filepath") or "", - entry.get("state_hash") or "", - metadata_result, - ) - except Exception: - pass - await self._write_scan_journal_entry( - entry.get("filepath") or "", - base_dir, - entry.get("state_hash") or "", - int(entry.get("mtime") or 0), - int(entry.get("size") or 0), - ) - stats["added"] += 1 - # Collect added asset ID for WebSocket notification - if added_ids is not None and res.data and res.data.get("asset_id"): - try: - added_ids.append(int(res.data["asset_id"])) - except Exception: - pass - if entry.get("fast") and to_enrich is not None: - if len(to_enrich) < MAX_TO_ENRICH_ITEMS: - to_enrich.append(entry.get("filepath") or str(entry.get("file_path") or "")) - try: - logger.debug( - "Indexed file (added): %s [asset_id=%s]", - entry.get("filepath") or "", - (res.data or {}).get("asset_id") if isinstance(res.data, dict) else None, - ) - except Exception: - pass - else: - stats["errors"] += 1 - continue - - # Unknown action: count as skipped to be safe. - stats["skipped"] += 1 - if not tx.ok: - raise RuntimeError(tx.error or "Commit failed") + await self._persist_prepared_entries_tx( + prepared=prepared, + base_dir=base_dir, + source=source, + root_id=root_id, + stats=stats, + to_enrich=to_enrich, + added_ids=added_ids, + ) except Exception as batch_error: if _is_fatal_db_error(batch_error): logger.error( @@ -1004,7 +1211,7 @@ class IndexScanner: exc_info=batch_error, ) raise - # If the entire batch transaction fails, fall back to processing items individually + suspect_fp, suspect_reason = self._diagnose_batch_failure(prepared, batch_error) if suspect_fp: logger.warning( @@ -1017,205 +1224,638 @@ class IndexScanner: "Batch transaction failed (%s). Falling back to individual processing.", suspect_reason, ) - logger.warning("Batch transaction failed: %s. Falling back to individual processing.", str(batch_error)) - stats["batch_fallbacks"] = int(stats.get("batch_fallbacks") or 0) + 1 - with self._batch_fallback_lock: - self._batch_fallback_count += 1 - stats["errors"] += len(prepared) # Temporarily count all as errors, will be corrected below - failed_entries: list[str] = [] + logger.warning( + "Batch transaction failed: %s. Falling back to individual processing.", + str(batch_error), + ) + await self._persist_prepared_entries_fallback( + prepared=prepared, + base_dir=base_dir, + source=source, + root_id=root_id, + stats=stats, + to_enrich=to_enrich, + ) - # Process each item individually to isolate failures + async def _persist_prepared_entries_tx( + self, + prepared: List[Dict[str, Any]], + base_dir: str, + source: str, + root_id: Optional[str], + stats: Dict[str, Any], + to_enrich: Optional[List[str]], + added_ids: Optional[List[int]], + ) -> None: + async with self.db.atransaction(mode="immediate") as tx: + if not tx.ok: + raise RuntimeError(tx.error or "Failed to begin transaction") for entry in prepared: - action = entry.get("action") - filepath_value = str(entry.get("filepath") or entry.get("file_path") or "unknown") + await self._process_prepared_entry_tx( + entry=entry, + base_dir=base_dir, + source=source, + root_id=root_id, + stats=stats, + to_enrich=to_enrich, + added_ids=added_ids, + ) + if not tx.ok: + raise RuntimeError(tx.error or "Commit failed") - if action in ("skipped", "skipped_journal"): - stats["skipped"] += 1 - stats["errors"] = max(0, stats["errors"] - 1) # Correct the error count + async def _process_prepared_entry_tx( + self, + entry: Dict[str, Any], + base_dir: str, + source: str, + root_id: Optional[str], + stats: Dict[str, Any], + to_enrich: Optional[List[str]], + added_ids: Optional[List[int]], + ) -> None: + action = entry.get("action") + if action in ("skipped", "skipped_journal"): + stats["skipped"] += 1 + return + if self._entry_state_drifted(entry, stats): + return + if action == "refresh": + await self._process_refresh_entry( + entry=entry, + base_dir=base_dir, + stats=stats, + to_enrich=to_enrich, + fallback_mode=False, + respect_enrich_limit=True, + ) + return + if action == "updated": + await self._process_updated_entry( + entry=entry, + base_dir=base_dir, + source=source, + root_id=root_id, + stats=stats, + to_enrich=to_enrich, + added_ids=added_ids, + fallback_mode=False, + respect_enrich_limit=True, + ) + return + if action == "added": + await self._process_added_entry( + entry=entry, + base_dir=base_dir, + source=source, + root_id=root_id, + stats=stats, + to_enrich=to_enrich, + added_ids=added_ids, + fallback_mode=False, + respect_enrich_limit=True, + ) + return + stats["skipped"] += 1 + + async def _persist_prepared_entries_fallback( + self, + prepared: List[Dict[str, Any]], + base_dir: str, + source: str, + root_id: Optional[str], + stats: Dict[str, Any], + to_enrich: Optional[List[str]], + ) -> None: + stats["batch_fallbacks"] = int(stats.get("batch_fallbacks") or 0) + 1 + with self._batch_fallback_lock: + self._batch_fallback_count += 1 + stats["errors"] += len(prepared) + failed_entries: List[str] = [] + + for entry in prepared: + action = entry.get("action") + filepath_value = self._entry_display_path(entry) + + if action in ("skipped", "skipped_journal"): + stats["skipped"] += 1 + self._fallback_correct_error(stats) + continue + try: + if self._entry_state_drifted(entry, stats): + self._fallback_correct_error(stats) continue - try: - file_path_obj = entry.get("file_path") - if isinstance(file_path_obj, Path): - expected_mtime_ns = int(entry.get("mtime_ns") or 0) - expected_size = int(entry.get("size") or 0) - if self._file_state_drifted(file_path_obj, expected_mtime_ns, expected_size): - stats["skipped"] += 1 - stats["skipped_state_changed"] = int(stats.get("skipped_state_changed") or 0) + 1 - stats["errors"] = max(0, stats["errors"] - 1) - continue - async with self.db.atransaction(mode="immediate") as tx: - if not tx.ok: - failed_entries.append(filepath_value) - continue - if action == "refresh": - asset_id = entry.get("asset_id") - metadata_result = entry.get("metadata_result") - if not asset_id or not isinstance(metadata_result, Result): - failed_entries.append(filepath_value) - continue - try: - refreshed = await MetadataHelpers.refresh_metadata_if_needed( - self.db, - int(asset_id), - metadata_result, - entry.get("filepath") or "", - base_dir, - entry.get("state_hash") or "", - int(entry.get("mtime") or 0), - int(entry.get("size") or 0), - self._write_scan_journal_entry, - ) - stats["skipped"] += 1 - stats["errors"] = max(0, stats["errors"] - 1) # Correct the error count - if refreshed and entry.get("fast") and to_enrich is not None: - if len(to_enrich) < MAX_TO_ENRICH_ITEMS: - to_enrich.append(entry.get("filepath") or str(entry.get("file_path") or "")) - except Exception as exc: - failed_entries.append(filepath_value) - logger.warning("Metadata refresh failed for asset_id=%s: %s", asset_id, exc) - continue - - if action == "updated": - asset_id = entry.get("asset_id") - metadata_result = entry.get("metadata_result") - if not asset_id or not isinstance(metadata_result, Result): - failed_entries.append(filepath_value) - continue - file_path_value = entry.get("file_path") - if not isinstance(file_path_value, Path): - failed_entries.append(filepath_value) - continue - cache_store = bool(entry.get("cache_store")) - if cache_store: - try: - await MetadataHelpers.store_metadata_cache( - self.db, - entry.get("filepath") or "", - entry.get("state_hash") or "", - metadata_result, - ) - except Exception: - pass - res = await self._update_asset( - int(asset_id), - file_path_value, - int(entry.get("mtime") or 0), - int(entry.get("size") or 0), - metadata_result, - source=source, - root_id=root_id, - write_metadata=not bool(entry.get("fast")), - ) - if res.ok: - await self._write_scan_journal_entry( - entry.get("filepath") or "", - base_dir, - entry.get("state_hash") or "", - int(entry.get("mtime") or 0), - int(entry.get("size") or 0), - ) - stats["updated"] += 1 - stats["errors"] = max(0, stats["errors"] - 1) # Correct the error count - if entry.get("fast") and to_enrich is not None: - to_enrich.append(entry.get("filepath") or str(entry.get("file_path") or "")) - try: - logger.debug( - "Indexed file (updated/fallback): %s [asset_id=%s]", - entry.get("filepath") or "", - asset_id, - ) - except Exception: - pass - else: - failed_entries.append(filepath_value) - continue - - if action == "added": - metadata_result = entry.get("metadata_result") - if not isinstance(metadata_result, Result): - failed_entries.append(filepath_value) - continue - kind_value = entry.get("kind") - file_path_value = entry.get("file_path") - if not isinstance(kind_value, str) or not isinstance(file_path_value, Path): - failed_entries.append(filepath_value) - continue - cache_store = bool(entry.get("cache_store")) - res = await self._add_asset( - entry.get("filename") or "", - entry.get("subfolder") or "", - entry.get("filepath") or "", - cast(FileKind, kind_value), - int(entry.get("mtime") or 0), - int(entry.get("size") or 0), - file_path_value, - metadata_result, - source=source, - root_id=root_id, - write_metadata=True, - ) - if res.ok: - if cache_store: - try: - await MetadataHelpers.store_metadata_cache( - self.db, - entry.get("filepath") or "", - entry.get("state_hash") or "", - metadata_result, - ) - except Exception: - pass - await self._write_scan_journal_entry( - entry.get("filepath") or "", - base_dir, - entry.get("state_hash") or "", - int(entry.get("mtime") or 0), - int(entry.get("size") or 0), - ) - stats["added"] += 1 - stats["errors"] = max(0, stats["errors"] - 1) # Correct the error count - try: - logger.debug( - "Indexed file (added/fallback): %s [asset_id=%s]", - entry.get("filepath") or "", - (res.data or {}).get("asset_id") if isinstance(res.data, dict) else None, - ) - except Exception: - pass - if entry.get("fast") and to_enrich is not None: - to_enrich.append(entry.get("filepath") or str(entry.get("file_path") or "")) - else: - failed_entries.append(filepath_value) - continue - - # Unknown action: count as skipped to be safe. - stats["skipped"] += 1 - stats["errors"] = max(0, stats["errors"] - 1) # Correct the error count + async with self.db.atransaction(mode="immediate") as tx: if not tx.ok: failed_entries.append(filepath_value) - except Exception as individual_error: - # If individual processing also fails, log and keep as error - failed_entries.append(filepath_value) - logger.warning( - "Individual processing failed for entry: %s. Error: %s", - filepath_value, - str(individual_error), - exc_info=individual_error, + continue + ok = await self._process_prepared_entry_fallback( + entry=entry, + base_dir=base_dir, + source=source, + root_id=root_id, + stats=stats, + to_enrich=to_enrich, ) - # Error count is already correct - if failed_entries: - sample = failed_entries[:5] - logger.warning("Batch fallback completed with %s failures (sample: %s)", len(failed_entries), sample) + if not ok: + failed_entries.append(filepath_value) + continue + if not tx.ok: + failed_entries.append(filepath_value) + except Exception as individual_error: + failed_entries.append(filepath_value) + logger.warning( + "Individual processing failed for entry: %s. Error: %s", + filepath_value, + str(individual_error), + exc_info=individual_error, + ) - duration = time.perf_counter() - batch_start - if duration > 0.2: - logger.debug( - "_index_batch timing: %.3fs (batch=%s)", - duration, - len(batch) if batch is not None else 0, + if failed_entries: + sample = failed_entries[:5] + logger.warning( + "Batch fallback completed with %s failures (sample: %s)", + len(failed_entries), + sample, ) + async def _process_prepared_entry_fallback( + self, + entry: Dict[str, Any], + base_dir: str, + source: str, + root_id: Optional[str], + stats: Dict[str, Any], + to_enrich: Optional[List[str]], + ) -> bool: + action = entry.get("action") + if action == "refresh": + return await self._process_refresh_entry( + entry=entry, + base_dir=base_dir, + stats=stats, + to_enrich=to_enrich, + fallback_mode=True, + respect_enrich_limit=True, + ) + if action == "updated": + return await self._process_updated_entry( + entry=entry, + base_dir=base_dir, + source=source, + root_id=root_id, + stats=stats, + to_enrich=to_enrich, + added_ids=None, + fallback_mode=True, + respect_enrich_limit=False, + ) + if action == "added": + return await self._process_added_entry( + entry=entry, + base_dir=base_dir, + source=source, + root_id=root_id, + stats=stats, + to_enrich=to_enrich, + added_ids=None, + fallback_mode=True, + respect_enrich_limit=False, + ) + + stats["skipped"] += 1 + self._fallback_correct_error(stats) + return True + + async def _process_refresh_entry( + self, + entry: Dict[str, Any], + base_dir: str, + stats: Dict[str, Any], + to_enrich: Optional[List[str]], + fallback_mode: bool, + respect_enrich_limit: bool, + ) -> bool: + asset_id = entry.get("asset_id") + metadata_result = entry.get("metadata_result") + invalid_refresh = self._invalid_refresh_entry(asset_id, metadata_result, stats, fallback_mode) + if invalid_refresh is not None: + return invalid_refresh + + try: + refreshed = await MetadataHelpers.refresh_metadata_if_needed( + self.db, + int(asset_id), + metadata_result, + *self._refresh_entry_context(entry, base_dir), + self._write_scan_journal_entry, + ) + self._record_refresh_outcome( + stats=stats, + fallback_mode=fallback_mode, + refreshed=refreshed, + entry=entry, + to_enrich=to_enrich, + respect_enrich_limit=respect_enrich_limit, + ) + return True + except Exception as exc: + if not fallback_mode: + stats["errors"] += 1 + logger.warning("Metadata refresh failed for asset_id=%s: %s", asset_id, exc) + return not fallback_mode + + @staticmethod + def _invalid_refresh_entry( + asset_id: Any, + metadata_result: Any, + stats: Dict[str, Any], + fallback_mode: bool, + ) -> Optional[bool]: + if asset_id and isinstance(metadata_result, Result): + return None + if fallback_mode: + return False + stats["skipped"] += 1 + return True + + @staticmethod + def _refresh_entry_context(entry: Dict[str, Any], base_dir: str) -> tuple[str, str, int, int]: + return ( + entry.get("filepath") or "", + base_dir, + entry.get("state_hash") or "", + int(entry.get("mtime") or 0), + int(entry.get("size") or 0), + ) + + def _record_refresh_outcome( + self, + *, + stats: Dict[str, Any], + fallback_mode: bool, + refreshed: bool, + entry: Dict[str, Any], + to_enrich: Optional[List[str]], + respect_enrich_limit: bool, + ) -> None: + stats["skipped"] += 1 + if fallback_mode: + self._fallback_correct_error(stats) + if refreshed: + self._append_to_enrich( + entry=entry, + to_enrich=to_enrich, + respect_limit=respect_enrich_limit, + ) + + async def _process_updated_entry( + self, + entry: Dict[str, Any], + base_dir: str, + source: str, + root_id: Optional[str], + stats: Dict[str, Any], + to_enrich: Optional[List[str]], + added_ids: Optional[List[int]], + fallback_mode: bool, + respect_enrich_limit: bool, + ) -> bool: + ctx = self._extract_update_entry_context(entry) + if ctx is None: + return self._handle_invalid_prepared_entry(stats, fallback_mode) + + await self._maybe_store_entry_cache(entry, ctx["metadata_result"]) + if not await self._apply_update_entry(entry, ctx, source, root_id): + return self._handle_update_or_add_failure(stats, fallback_mode) + + await self._write_entry_scan_journal(entry, base_dir) + self._record_index_entry_success( + stats=stats, + fallback_mode=fallback_mode, + added_ids=added_ids, + added_asset_id=ctx["asset_id"], + entry=entry, + action="updated", + to_enrich=to_enrich, + respect_enrich_limit=respect_enrich_limit, + ) + return True + + async def _process_added_entry( + self, + entry: Dict[str, Any], + base_dir: str, + source: str, + root_id: Optional[str], + stats: Dict[str, Any], + to_enrich: Optional[List[str]], + added_ids: Optional[List[int]], + fallback_mode: bool, + respect_enrich_limit: bool, + ) -> bool: + ctx = self._extract_add_entry_context(entry) + if ctx is None: + return self._handle_invalid_prepared_entry(stats, fallback_mode) + + add_result = await self._apply_add_entry(entry, ctx, source, root_id) + if not add_result.ok: + return self._handle_update_or_add_failure(stats, fallback_mode) + + await self._maybe_store_entry_cache(entry, ctx["metadata_result"]) + await self._write_entry_scan_journal(entry, base_dir) + added_asset_id = self._added_asset_id_from_result(add_result) + self._record_index_entry_success( + stats=stats, + fallback_mode=fallback_mode, + added_ids=added_ids, + added_asset_id=added_asset_id, + entry=entry, + action="added", + to_enrich=to_enrich, + respect_enrich_limit=respect_enrich_limit, + ) + return True + + def _extract_update_entry_context(self, entry: Dict[str, Any]) -> Optional[Dict[str, Any]]: + asset_id = entry.get("asset_id") + metadata_result = entry.get("metadata_result") + file_path_value = entry.get("file_path") + if not asset_id or not isinstance(metadata_result, Result) or not isinstance(file_path_value, Path): + return None + return { + "asset_id": int(asset_id), + "metadata_result": metadata_result, + "file_path": file_path_value, + } + + def _extract_add_entry_context(self, entry: Dict[str, Any]) -> Optional[Dict[str, Any]]: + metadata_result = entry.get("metadata_result") + kind_value = entry.get("kind") + file_path_value = entry.get("file_path") + if ( + not isinstance(metadata_result, Result) + or not isinstance(kind_value, str) + or not isinstance(file_path_value, Path) + ): + return None + return { + "metadata_result": metadata_result, + "kind": cast(FileKind, kind_value), + "file_path": file_path_value, + } + + @staticmethod + def _handle_invalid_prepared_entry(stats: Dict[str, Any], fallback_mode: bool) -> bool: + if fallback_mode: + return False + stats["errors"] += 1 + return True + + @staticmethod + def _handle_update_or_add_failure(stats: Dict[str, Any], fallback_mode: bool) -> bool: + if fallback_mode: + return False + stats["errors"] += 1 + return True + + async def _maybe_store_entry_cache(self, entry: Dict[str, Any], metadata_result: Result[Any]) -> None: + if not bool(entry.get("cache_store")): + return + try: + await MetadataHelpers.store_metadata_cache( + self.db, + entry.get("filepath") or "", + entry.get("state_hash") or "", + metadata_result, + ) + except Exception: + return + + async def _apply_update_entry( + self, + entry: Dict[str, Any], + ctx: Dict[str, Any], + source: str, + root_id: Optional[str], + ) -> bool: + res = await self._update_asset( + ctx["asset_id"], + ctx["file_path"], + int(entry.get("mtime") or 0), + int(entry.get("size") or 0), + ctx["metadata_result"], + source=source, + root_id=root_id, + write_metadata=not bool(entry.get("fast")), + ) + return bool(res.ok) + + async def _apply_add_entry( + self, + entry: Dict[str, Any], + ctx: Dict[str, Any], + source: str, + root_id: Optional[str], + ) -> Result[Dict[str, Any]]: + return await self._add_asset( + entry.get("filename") or "", + entry.get("subfolder") or "", + entry.get("filepath") or "", + ctx["kind"], + int(entry.get("mtime") or 0), + int(entry.get("size") or 0), + ctx["file_path"], + ctx["metadata_result"], + source=source, + root_id=root_id, + write_metadata=True, + ) + + async def _write_entry_scan_journal(self, entry: Dict[str, Any], base_dir: str) -> None: + await self._write_scan_journal_entry( + entry.get("filepath") or "", + base_dir, + entry.get("state_hash") or "", + int(entry.get("mtime") or 0), + int(entry.get("size") or 0), + ) + + @staticmethod + def _added_asset_id_from_result(add_result: Result[Dict[str, Any]]) -> Optional[int]: + try: + if add_result.data and add_result.data.get("asset_id"): + return int(add_result.data["asset_id"]) + except Exception: + return None + return None + + def _record_index_entry_success( + self, + *, + stats: Dict[str, Any], + fallback_mode: bool, + added_ids: Optional[List[int]], + added_asset_id: Optional[int], + entry: Dict[str, Any], + action: str, + to_enrich: Optional[List[str]], + respect_enrich_limit: bool, + ) -> None: + stats[action] += 1 + if fallback_mode: + self._fallback_correct_error(stats) + elif added_ids is not None and added_asset_id: + try: + added_ids.append(int(added_asset_id)) + except Exception: + pass + self._append_to_enrich(entry=entry, to_enrich=to_enrich, respect_limit=respect_enrich_limit) + try: + logger.debug( + "Indexed file (%s%s): %s [asset_id=%s]", + action, + "/fallback" if fallback_mode else "", + entry.get("filepath") or "", + added_asset_id, + ) + except Exception: + pass + + def _entry_state_drifted(self, entry: Dict[str, Any], stats: Dict[str, Any]) -> bool: + file_path_value = entry.get("file_path") + if not isinstance(file_path_value, Path): + return False + expected_mtime_ns = int(entry.get("mtime_ns") or 0) + expected_size = int(entry.get("size") or 0) + if not self._file_state_drifted(file_path_value, expected_mtime_ns, expected_size): + return False + stats["skipped"] += 1 + stats["skipped_state_changed"] = int(stats.get("skipped_state_changed") or 0) + 1 + return True + + def _append_to_enrich( + self, + entry: Dict[str, Any], + to_enrich: Optional[List[str]], + respect_limit: bool, + ) -> None: + if to_enrich is None or not entry.get("fast"): + return + if respect_limit and len(to_enrich) >= MAX_TO_ENRICH_ITEMS: + return + to_enrich.append(self._entry_display_path(entry)) + + @staticmethod + def _fallback_correct_error(stats: Dict[str, Any]) -> None: + stats["errors"] = max(0, int(stats.get("errors") or 0) - 1) + + @staticmethod + def _entry_display_path(entry: Dict[str, Any]) -> str: + return str(entry.get("filepath") or entry.get("file_path") or "unknown") + + @staticmethod + def _build_refresh_entry( + *, + asset_id: int, + metadata_result: Result[Dict[str, Any]], + filepath: str, + file_path: Path, + state_hash: str, + mtime: int, + size: int, + fast: bool, + cache_store: bool, + ) -> Dict[str, Any]: + return { + "action": "refresh", + "asset_id": asset_id, + "metadata_result": metadata_result, + "filepath": filepath, + "file_path": file_path, + "state_hash": state_hash, + "mtime": mtime, + "size": size, + "fast": fast, + "cache_store": cache_store, + } + + @staticmethod + def _build_updated_entry( + *, + asset_id: int, + metadata_result: Result[Dict[str, Any]], + filepath: str, + file_path: Path, + state_hash: str, + mtime: int, + size: int, + fast: bool, + cache_store: bool, + mtime_ns: Optional[int] = None, + ) -> Dict[str, Any]: + entry: Dict[str, Any] = { + "action": "updated", + "asset_id": asset_id, + "metadata_result": metadata_result, + "filepath": filepath, + "file_path": file_path, + "state_hash": state_hash, + "mtime": mtime, + "size": size, + "fast": fast, + "cache_store": cache_store, + } + if mtime_ns is not None: + entry["mtime_ns"] = int(mtime_ns) + return entry + + @staticmethod + def _build_added_entry( + *, + filename: str, + subfolder: str, + filepath: str, + kind: FileKind, + metadata_result: Result[Dict[str, Any]], + file_path: Path, + state_hash: str, + mtime: int, + size: int, + fast: bool, + cache_store: bool, + mtime_ns: Optional[int] = None, + ) -> Dict[str, Any]: + entry: Dict[str, Any] = { + "action": "added", + "filename": filename, + "subfolder": subfolder, + "filepath": filepath, + "kind": kind, + "metadata_result": metadata_result, + "file_path": file_path, + "state_hash": state_hash, + "mtime": mtime, + "size": size, + "fast": fast, + "cache_store": cache_store, + } + if mtime_ns is not None: + entry["mtime_ns"] = int(mtime_ns) + return entry + + async def _prepare_metadata_for_entry( + self, + filepath: str, + fast: bool, + ) -> tuple[Result[Dict[str, Any]], bool]: + if fast: + return Result.Ok({}), False + metadata_result = await self.metadata.get_metadata(filepath, scan_id=self._current_scan_id) + if metadata_result.ok: + return metadata_result, True + return MetadataHelpers.metadata_error_payload(metadata_result, filepath), False + async def _prepare_index_entry( self, file_path: Path, @@ -1230,129 +1870,168 @@ class IndexScanner: Prepare indexing work for a file (stat + incremental decision + metadata extraction), but do not write to the DB. DB writes are applied in a batch transaction. """ + prepare_ctx, stat_error = await self._prepare_index_entry_context(file_path, existing_state, incremental) + if prepare_ctx is None: + return Result.Err("STAT_FAILED", f"Failed to stat file: {stat_error}") + + skip = await self._maybe_skip_prepare_for_incremental( + prepare_ctx=prepare_ctx, + incremental=incremental, + fast=fast, + file_path=file_path, + ) + if skip is not None: + return skip + + metadata_result, cache_store = await self._prepare_metadata_for_entry(prepare_ctx["filepath"], fast) + if self._is_incremental_unchanged(prepare_ctx, incremental): + return Result.Ok( + self._build_refresh_entry( + asset_id=prepare_ctx["existing_id"], + metadata_result=metadata_result, + filepath=prepare_ctx["filepath"], + file_path=file_path, + state_hash=prepare_ctx["state_hash"], + mtime=prepare_ctx["mtime"], + size=prepare_ctx["size"], + fast=fast, + cache_store=cache_store, + ) + ) + if prepare_ctx["existing_id"]: + return Result.Ok( + self._build_updated_entry( + asset_id=prepare_ctx["existing_id"], + metadata_result=metadata_result, + filepath=prepare_ctx["filepath"], + file_path=file_path, + state_hash=prepare_ctx["state_hash"], + mtime=prepare_ctx["mtime"], + size=prepare_ctx["size"], + fast=fast, + cache_store=cache_store, + ) + ) + return Result.Ok(self._build_added_entry_from_prepare_ctx(prepare_ctx, file_path, base_dir, metadata_result, cache_store, fast)) + + async def _prepare_index_entry_context( + self, + file_path: Path, + existing_state: Optional[Dict[str, Any]], + incremental: bool, + ) -> tuple[Optional[Dict[str, Any]], Optional[str]]: stat_result = await self._stat_with_retry(file_path) if not stat_result[0]: - return Result.Err("STAT_FAILED", f"Failed to stat file: {stat_result[1]}") + return None, str(stat_result[1]) stat = stat_result[1] - mtime_ns = stat.st_mtime_ns + mtime_ns = int(stat.st_mtime_ns) mtime = int(stat.st_mtime) size = int(stat.st_size) - filepath = self._normalize_filepath_str(file_path) state_hash = self._compute_state_hash(filepath, mtime_ns, size) - - journal_state_hash = None - if isinstance(existing_state, dict) and "journal_state_hash" in existing_state: - journal_state_hash = existing_state.get("journal_state_hash") - elif incremental: - journal_entry = await self._get_journal_entry(filepath) - journal_state_hash = journal_entry.get("state_hash") if isinstance(journal_entry, dict) else None - + journal_state_hash = await self._get_journal_state_hash_for_index_file(filepath, existing_state, incremental) existing_asset = existing_state if isinstance(existing_state, dict) and existing_state.get("id") is not None else None - existing_id = 0 - existing_mtime = 0 - if existing_asset is not None: - try: - existing_id = int(existing_asset.get("id") or 0) - existing_mtime = int(existing_asset.get("mtime") or 0) - except Exception: - existing_id = 0 - existing_mtime = 0 + existing_id, existing_mtime = self._extract_existing_asset_state(existing_asset) + return { + "filepath": filepath, + "state_hash": state_hash, + "journal_state_hash": journal_state_hash, + "existing_id": existing_id, + "existing_mtime": existing_mtime, + "mtime": mtime, + "size": size, + }, None - if incremental and journal_state_hash and str(journal_state_hash) == state_hash: - if fast: - return Result.Ok({"action": "skipped_journal"}) - if existing_id: - if await self._asset_has_rich_metadata(existing_id): - return Result.Ok({"action": "skipped_journal"}) + async def _maybe_skip_prepare_for_incremental( + self, + *, + prepare_ctx: Dict[str, Any], + incremental: bool, + fast: bool, + file_path: Path, + ) -> Optional[Result[Dict[str, Any]]]: + if await self._should_skip_by_journal_state(prepare_ctx, incremental=incremental, fast=fast): + return Result.Ok({"action": "skipped_journal"}) + if not self._is_incremental_unchanged(prepare_ctx, incremental): + return None + cached_result = await self._refresh_entry_from_cached_metadata(prepare_ctx, file_path=file_path, fast=fast) + if cached_result is not None: + return cached_result + if await self._asset_has_rich_metadata(prepare_ctx["existing_id"]): + return Result.Ok({"action": "skipped"}) + return None - # If incremental and unchanged, try cached metadata first. - if incremental and existing_id and existing_mtime == mtime: - cached_metadata = await MetadataHelpers.retrieve_cached_metadata(self.db, filepath, state_hash) - if cached_metadata and cached_metadata.ok: - return Result.Ok( - { - "action": "refresh", - "asset_id": existing_id, - "metadata_result": cached_metadata, - "filepath": filepath, - "file_path": file_path, - "state_hash": state_hash, - "mtime": mtime, - "size": size, - "fast": fast, - "cache_store": False, - } - ) - - if await self._asset_has_rich_metadata(existing_id): - return Result.Ok({"action": "skipped"}) - - # Compute metadata only when needed. - cache_store = False + async def _should_skip_by_journal_state( + self, + prepare_ctx: Dict[str, Any], + *, + incremental: bool, + fast: bool, + ) -> bool: + if not incremental: + return False + if not prepare_ctx["journal_state_hash"]: + return False + if str(prepare_ctx["journal_state_hash"]) != prepare_ctx["state_hash"]: + return False if fast: - metadata_result: Result[Dict[str, Any]] = Result.Ok({}) - else: - metadata_result = await self.metadata.get_metadata(filepath, scan_id=self._current_scan_id) - if metadata_result.ok: - cache_store = True - if not metadata_result.ok: - metadata_result = MetadataHelpers.metadata_error_payload(metadata_result, filepath) + return True + existing_id = int(prepare_ctx.get("existing_id") or 0) + return bool(existing_id and await self._asset_has_rich_metadata(existing_id)) - # For unchanged files, this is a metadata refresh attempt. - if incremental and existing_id and existing_mtime == mtime: - return Result.Ok( - { - "action": "refresh", - "asset_id": existing_id, - "metadata_result": metadata_result, - "filepath": filepath, - "file_path": file_path, - "state_hash": state_hash, - "mtime": mtime, - "size": size, - "fast": fast, - "cache_store": cache_store, - } + async def _refresh_entry_from_cached_metadata( + self, + prepare_ctx: Dict[str, Any], + *, + file_path: Path, + fast: bool, + ) -> Optional[Result[Dict[str, Any]]]: + cached = await MetadataHelpers.retrieve_cached_metadata(self.db, prepare_ctx["filepath"], prepare_ctx["state_hash"]) + if not (cached and cached.ok): + return None + return Result.Ok( + self._build_refresh_entry( + asset_id=prepare_ctx["existing_id"], + metadata_result=cached, + filepath=prepare_ctx["filepath"], + file_path=file_path, + state_hash=prepare_ctx["state_hash"], + mtime=prepare_ctx["mtime"], + size=prepare_ctx["size"], + fast=fast, + cache_store=False, ) + ) - # Calculate relative path components (needed for inserts). + @staticmethod + def _is_incremental_unchanged(prepare_ctx: Dict[str, Any], incremental: bool) -> bool: + return bool(incremental and prepare_ctx["existing_id"] and prepare_ctx["existing_mtime"] == prepare_ctx["mtime"]) + + def _build_added_entry_from_prepare_ctx( + self, + prepare_ctx: Dict[str, Any], + file_path: Path, + base_dir: str, + metadata_result: Result[Dict[str, Any]], + cache_store: bool, + fast: bool, + ) -> Dict[str, Any]: rel_path = self._safe_relative_path(file_path, base_dir) filename = file_path.name subfolder = str(rel_path.parent) if rel_path.parent != Path(".") else "" kind = classify_file(filename) - - if existing_id: - return Result.Ok( - { - "action": "updated", - "asset_id": existing_id, - "metadata_result": metadata_result, - "filepath": filepath, - "file_path": file_path, - "state_hash": state_hash, - "mtime": mtime, - "size": size, - "fast": fast, - "cache_store": cache_store, - } - ) - - return Result.Ok( - { - "action": "added", - "filename": filename, - "subfolder": subfolder, - "filepath": filepath, - "kind": kind, - "metadata_result": metadata_result, - "file_path": file_path, - "state_hash": state_hash, - "mtime": mtime, - "size": size, - "fast": fast, - "cache_store": cache_store, - } + return self._build_added_entry( + filename=filename, + subfolder=subfolder, + filepath=prepare_ctx["filepath"], + kind=kind, + metadata_result=metadata_result, + file_path=file_path, + state_hash=prepare_ctx["state_hash"], + mtime=prepare_ctx["mtime"], + size=prepare_ctx["size"], + fast=fast, + cache_store=cache_store, ) @@ -1368,15 +2047,6 @@ class IndexScanner: Yields: File paths one by one """ - def _is_supported(p: Path) -> bool: - try: - ext = p.suffix.lower() - except Exception: - ext = "" - if ext and _EXT_TO_KIND: - return _EXT_TO_KIND.get(ext, "unknown") != "unknown" - return classify_file(str(p)) != "unknown" - if recursive: # Iterative scandir is generally faster than os.walk on large trees/NAS shares. stack: list[Path] = [directory] @@ -1386,34 +2056,55 @@ class IndexScanner: with os.scandir(current) as it: for entry in it: self._scan_iops_wait() - try: - if entry.is_dir(follow_symlinks=False): - stack.append(Path(entry.path)) - continue - # Keep historical behavior: index symlinks to files, but do not recurse - # into symlinked directories. - if not entry.is_file(follow_symlinks=True): - continue - name = entry.name - try: - ext = os.path.splitext(name)[1].lower() - except Exception: - ext = "" - if ext and _EXT_TO_KIND and _EXT_TO_KIND.get(ext, "unknown") == "unknown": - continue - file_path = Path(entry.path) - if _is_supported(file_path): - yield file_path - except (OSError, PermissionError): + next_dir = self._iter_files_next_dir(entry) + if isinstance(next_dir, Path): + stack.append(next_dir) continue + file_path = self._iter_files_candidate(entry) + if file_path is not None: + yield file_path except (OSError, PermissionError): continue else: for item in directory.iterdir(): self._scan_iops_wait() - if item.is_file() and _is_supported(item): + if item.is_file() and self._is_supported_scan_file(item): yield item + @staticmethod + def _is_supported_scan_file(path: Path) -> bool: + try: + ext = path.suffix.lower() + except Exception: + ext = "" + if ext and _EXT_TO_KIND: + return _EXT_TO_KIND.get(ext, "unknown") != "unknown" + return classify_file(str(path)) != "unknown" + + @staticmethod + def _iter_files_next_dir(entry) -> Optional[Path]: + try: + if entry.is_dir(follow_symlinks=False): + return Path(entry.path) + except (OSError, PermissionError): + return None + return None + + def _iter_files_candidate(self, entry) -> Optional[Path]: + try: + # Keep historical behavior: index symlinks to files, but do not recurse into symlinked dirs. + if not entry.is_file(follow_symlinks=True): + return None + ext = os.path.splitext(entry.name)[1].lower() + if ext and _EXT_TO_KIND and _EXT_TO_KIND.get(ext, "unknown") == "unknown": + return None + file_path = Path(entry.path) + if self._is_supported_scan_file(file_path): + return file_path + except (OSError, PermissionError): + return None + return None + @staticmethod def _file_state_drifted(file_path: Path, expected_mtime_ns: int, expected_size: int) -> bool: """Return True when file changed between initial stat and DB write time.""" @@ -1497,11 +2188,9 @@ class IndexScanner: Batch lookup scan_journal state_hash for a list of filepaths. Returns {filepath: state_hash}. """ - cleaned = [str(p) for p in (filepaths or []) if p] + cleaned = self._clean_journal_lookup_paths(filepaths) if not cleaned: return {} - if len(cleaned) > MAX_SCAN_JOURNAL_LOOKUP: - cleaned = cleaned[:MAX_SCAN_JOURNAL_LOOKUP] res = await self.db.aquery_in( "SELECT filepath, state_hash FROM scan_journal WHERE {IN_CLAUSE}", "filepath", @@ -1509,8 +2198,19 @@ class IndexScanner: ) if not res.ok: return {} + return self._journal_rows_to_map(res.data or []) + + @staticmethod + def _clean_journal_lookup_paths(filepaths: List[str]) -> List[str]: + cleaned = [str(p) for p in (filepaths or []) if p] + if len(cleaned) > MAX_SCAN_JOURNAL_LOOKUP: + return cleaned[:MAX_SCAN_JOURNAL_LOOKUP] + return cleaned + + @staticmethod + def _journal_rows_to_map(rows: List[Any]) -> Dict[str, str]: out: Dict[str, str] = {} - for row in res.data or []: + for row in rows: if not isinstance(row, dict): continue fp = row.get("filepath") @@ -1592,222 +2292,444 @@ class IndexScanner: Returns: Result with action taken (added, updated, skipped) """ - stat_result = await self._stat_with_retry(file_path) - if not stat_result[0]: - return Result.Err("STAT_FAILED", f"Failed to stat file: {stat_result[1]}") - stat = stat_result[1] - mtime_ns = stat.st_mtime_ns - mtime = int(stat.st_mtime) - size = stat.st_size + state = await self._build_index_file_state(file_path) + if isinstance(state, Result): + return state + file_path, filepath, state_hash, mtime, size = state - # Normalize Windows path casing to avoid duplicate rows differing only by case. - if IS_WINDOWS: - try: - file_path = Path(str(file_path)).resolve(strict=True) - except Exception: - pass + journal_state_hash = await self._get_journal_state_hash_for_index_file( + filepath, + existing_state, + incremental, + ) + existing_asset = await self._resolve_existing_asset_for_index_file(filepath, existing_state) + existing_id, existing_mtime = self._extract_existing_asset_state(existing_asset) - # Compute file fingerprint before any caching or journal logic (prevents use-before-set bugs) - filepath = self._normalize_filepath_str(file_path) - state_hash = self._compute_state_hash(filepath, mtime_ns, size) + skip_journal = await self._try_skip_by_journal( + incremental=incremental, + journal_state_hash=journal_state_hash, + state_hash=state_hash, + fast=fast, + existing_id=existing_id, + ) + if skip_journal is not None: + return skip_journal - # Skip unchanged files as early as possible (before running ExifTool/FFProbe). - journal_state_hash = None - if isinstance(existing_state, dict) and "journal_state_hash" in existing_state: - journal_state_hash = existing_state.get("journal_state_hash") - else: - journal_entry = await self._get_journal_entry(filepath) - journal_state_hash = journal_entry.get("state_hash") if isinstance(journal_entry, dict) else None + cached_refresh = await self._try_cached_incremental_refresh( + incremental=incremental, + existing_id=existing_id, + existing_mtime=existing_mtime, + mtime=mtime, + filepath=filepath, + state_hash=state_hash, + base_dir=base_dir, + size=size, + ) + if cached_refresh is not None: + return cached_refresh - # Check if file already exists (allow caller to prefetch to avoid N+1 queries) - existing_asset = None - if isinstance(existing_state, dict) and existing_state.get("id") is not None: - existing_asset = existing_state - else: - existing = await self.db.aquery( - "SELECT id, mtime, filepath FROM assets WHERE filepath = ?", - (filepath,) - ) + metadata_result = await self._extract_metadata_for_index_file( + file_path=file_path, + filepath=filepath, + state_hash=state_hash, + fast=fast, + ) - if existing.ok and existing.data and len(existing.data) > 0: - existing_asset = existing.data[0] - elif IS_WINDOWS: - try: - existing_ci = await self.db.aquery( - "SELECT id, mtime, filepath FROM assets WHERE filepath = ? COLLATE NOCASE", - (filepath,) - ) - if existing_ci.ok and existing_ci.data and len(existing_ci.data) > 0: - existing_asset = existing_ci.data[0] - except Exception: - pass + existing_id, existing_mtime = self._extract_existing_asset_state(existing_asset) + incremental_refresh = await self._try_incremental_refresh_with_metadata( + incremental=incremental, + existing_id=existing_id, + existing_mtime=existing_mtime, + mtime=mtime, + filepath=filepath, + state_hash=state_hash, + base_dir=base_dir, + size=size, + metadata_result=metadata_result, + ) + if incremental_refresh is not None: + return incremental_refresh - existing_id = 0 - if existing_asset is not None: - try: - existing_id = int(existing_asset.get("id") or 0) - except Exception: - existing_id = 0 - - if incremental and journal_state_hash and str(journal_state_hash) == state_hash: - if fast: - return Result.Ok({"action": "skipped_journal"}) - if existing_id: - if await self._asset_has_rich_metadata(existing_id): - return Result.Ok({"action": "skipped_journal"}) - - if existing_asset is not None: - try: - existing_mtime = int(existing_asset.get("mtime") or 0) - existing_id = int(existing_asset.get("id") or 0) - except Exception: - existing_mtime = 0 - existing_id = 0 - - if incremental and existing_mtime == mtime and existing_id: - # Prefer cached metadata; it's cheap and avoids rerunning tools. - cached_metadata = await MetadataHelpers.retrieve_cached_metadata(self.db, filepath, state_hash) - if cached_metadata and cached_metadata.ok: - refreshed = False - async with self.db.atransaction(mode="immediate") as tx: - if not tx.ok: - logger.warning("Metadata refresh skipped (transaction begin failed) for %s: %s", filepath, tx.error) - return Result.Ok({"action": "skipped"}) - refreshed = await MetadataHelpers.refresh_metadata_if_needed( - self.db, - existing_id, - cached_metadata, - filepath, - base_dir, - state_hash, - mtime, - size, - self._write_scan_journal_entry, - ) - if refreshed: - await self._write_scan_journal_entry(filepath, base_dir, state_hash, mtime, size) - if not tx.ok: - logger.warning("Metadata refresh commit failed for %s: %s", filepath, tx.error) - return Result.Ok({"action": "skipped"}) - action = "skipped_refresh" if refreshed else "skipped" - return Result.Ok({"action": action}) - - # If we already have a metadata row, don't rerun tools just for an incremental pass. - if await self._asset_has_rich_metadata(existing_id): - return Result.Ok({"action": "skipped"}) - - # Compute metadata (may run tools) only when needed. - if fast: - metadata_result: Result[Dict[str, Any]] = Result.Ok({}) - else: - metadata_result = await self.metadata.get_metadata(filepath, scan_id=self._current_scan_id) - if metadata_result.ok: - quality = metadata_result.meta.get("quality", "full") - if quality in ("degraded", "partial"): - logger.debug( - "Metadata extraction completed with degraded quality for %s: %s", - file_path, - quality, - ) - await MetadataHelpers.store_metadata_cache(self.db, filepath, state_hash, metadata_result) - - if not metadata_result.ok: - if metadata_result.code == ErrorCode.FFPROBE_ERROR: - self._log_scan_event( - logging.WARNING, - "FFprobe error during metadata extraction", - filepath=str(file_path), - tool="ffprobe", - error=metadata_result.error, - code=metadata_result.code - ) - else: - self._log_scan_event( - logging.WARNING, - "Metadata extraction issue", - filepath=str(file_path), - error=metadata_result.error, - code=metadata_result.code - ) - metadata_result = MetadataHelpers.metadata_error_payload(metadata_result, str(file_path)) - - # Calculate relative path components rel_path = self._safe_relative_path(file_path, base_dir) filename = file_path.name subfolder = str(rel_path.parent) if rel_path.parent != Path(".") else "" kind = classify_file(filename) - if existing_asset is not None: - # Skip if incremental and unchanged - try: - existing_mtime = int(existing_asset.get("mtime") or 0) - existing_id = int(existing_asset.get("id") or 0) - except Exception: - existing_mtime = 0 - existing_id = 0 + return await self._insert_new_asset_for_index_file( + filename=filename, + subfolder=subfolder, + filepath=filepath, + kind=kind, + mtime=mtime, + size=size, + file_path=file_path, + metadata_result=metadata_result, + base_dir=base_dir, + state_hash=state_hash, + source=source, + root_id=root_id, + ) - if incremental and existing_mtime == mtime and existing_id: - refreshed = False - tx_state = None - try: - async with self.db.lock_for_asset(existing_id): - try: - async with self.db.atransaction(mode="immediate") as tx: - tx_state = tx - if not tx_state.ok: - return Result.Err("DB_ERROR", tx_state.error or "Failed to begin transaction") - refreshed = await MetadataHelpers.refresh_metadata_if_needed( - self.db, - existing_id, - metadata_result, - filepath, - base_dir, - state_hash, - mtime, - size, - self._write_scan_journal_entry, - ) - if refreshed: - await self._write_scan_journal_entry(filepath, base_dir, state_hash, mtime, size) - except sqlite3.OperationalError as exc: - if self.db._is_locked_error(exc): - logger.warning("Database busy while refreshing metadata (asset=%s): %s", existing_id, exc) - raise - except sqlite3.OperationalError: - return Result.Err("DB_BUSY", "Database busy while refreshing metadata") - if not tx_state or not tx_state.ok: - return Result.Err("DB_ERROR", tx_state.error or "Commit failed") - action = "skipped_refresh" if refreshed else "skipped" - return Result.Ok({"action": action}) + async def _build_index_file_state( + self, + file_path: Path, + ) -> Result[Dict[str, str]] | tuple[Path, str, str, int, int]: + stat_result = await self._stat_with_retry(file_path) + if not stat_result[0]: + return Result.Err("STAT_FAILED", f"Failed to stat file: {stat_result[1]}") + stat = stat_result[1] + mtime_ns = int(stat.st_mtime_ns) + mtime = int(stat.st_mtime) + size = int(stat.st_size) + resolved_path = self._resolve_index_file_path(file_path) + filepath = self._normalize_filepath_str(resolved_path) + state_hash = self._compute_state_hash(filepath, mtime_ns, size) + return resolved_path, filepath, state_hash, mtime, size - # Add new asset - tx_state = None + @staticmethod + def _resolve_index_file_path(file_path: Path) -> Path: + if not IS_WINDOWS: + return file_path try: - async with self.db.atransaction(mode="immediate") as tx: - tx_state = tx - if not tx_state.ok: - return Result.Err("DB_ERROR", tx_state.error or "Failed to begin transaction") - result = await self._add_asset( - filename, - subfolder, + return Path(str(file_path)).resolve(strict=True) + except Exception: + return file_path + + async def _try_skip_by_journal( + self, + *, + incremental: bool, + journal_state_hash: Optional[str], + state_hash: str, + fast: bool, + existing_id: int, + ) -> Optional[Result[Dict[str, str]]]: + if not (incremental and journal_state_hash and str(journal_state_hash) == state_hash): + return None + if fast: + return Result.Ok({"action": "skipped_journal"}) + if existing_id and await self._asset_has_rich_metadata(existing_id): + return Result.Ok({"action": "skipped_journal"}) + return None + + async def _get_journal_state_hash_for_index_file( + self, + filepath: str, + existing_state: Optional[Dict[str, Any]], + incremental: bool, + ) -> Optional[str]: + if isinstance(existing_state, dict) and "journal_state_hash" in existing_state: + return cast(Optional[str], existing_state.get("journal_state_hash")) + if not incremental: + return None + journal_entry = await self._get_journal_entry(filepath) + if not isinstance(journal_entry, dict): + return None + value = journal_entry.get("state_hash") + return str(value) if value else None + + async def _resolve_existing_asset_for_index_file( + self, + filepath: str, + existing_state: Optional[Dict[str, Any]], + ) -> Optional[Dict[str, Any]]: + if isinstance(existing_state, dict) and existing_state.get("id") is not None: + return existing_state + + existing = await self.db.aquery( + "SELECT id, mtime, filepath FROM assets WHERE filepath = ?", + (filepath,), + ) + if existing.ok and existing.data: + return existing.data[0] + + if not IS_WINDOWS: + return None + + try: + existing_ci = await self.db.aquery( + "SELECT id, mtime, filepath FROM assets WHERE filepath = ? COLLATE NOCASE", + (filepath,), + ) + except Exception: + return None + if existing_ci.ok and existing_ci.data: + return existing_ci.data[0] + return None + + @staticmethod + def _extract_existing_asset_state(existing_asset: Optional[Dict[str, Any]]) -> tuple[int, int]: + if not isinstance(existing_asset, dict): + return 0, 0 + try: + existing_id = int(existing_asset.get("id") or 0) + existing_mtime = int(existing_asset.get("mtime") or 0) + except Exception: + return 0, 0 + return existing_id, existing_mtime + + async def _try_cached_incremental_refresh( + self, + incremental: bool, + existing_id: int, + existing_mtime: int, + mtime: int, + filepath: str, + state_hash: str, + base_dir: str, + size: int, + ) -> Optional[Result[Dict[str, str]]]: + if not (incremental and existing_id and existing_mtime == mtime): + return None + + cached_metadata = await MetadataHelpers.retrieve_cached_metadata(self.db, filepath, state_hash) + if cached_metadata and cached_metadata.ok: + return await self._refresh_from_cached_metadata( + existing_id=existing_id, + cached_metadata=cached_metadata, + filepath=filepath, + base_dir=base_dir, + state_hash=state_hash, + mtime=mtime, + size=size, + ) + + if await self._asset_has_rich_metadata(existing_id): + return Result.Ok({"action": "skipped"}) + return None + + async def _refresh_from_cached_metadata( + self, + *, + existing_id: int, + cached_metadata: Result[Dict[str, Any]], + filepath: str, + base_dir: str, + state_hash: str, + mtime: int, + size: int, + ) -> Result[Dict[str, str]]: + tx_state, refreshed = await self._refresh_from_cached_metadata_tx( + existing_id=existing_id, + cached_metadata=cached_metadata, + filepath=filepath, + base_dir=base_dir, + state_hash=state_hash, + mtime=mtime, + size=size, + ) + if not tx_state: + return Result.Ok({"action": "skipped"}) + if not tx_state.ok: + logger.warning("Metadata refresh commit failed for %s: %s", filepath, tx_state.error) + return Result.Ok({"action": "skipped"}) + return Result.Ok({"action": "skipped_refresh" if refreshed else "skipped"}) + + async def _refresh_from_cached_metadata_tx( + self, + *, + existing_id: int, + cached_metadata: Result[Dict[str, Any]], + filepath: str, + base_dir: str, + state_hash: str, + mtime: int, + size: int, + ) -> tuple[Optional[Any], bool]: + refreshed = False + async with self.db.atransaction(mode="immediate") as tx: + if not tx.ok: + logger.warning( + "Metadata refresh skipped (transaction begin failed) for %s: %s", filepath, - kind, + tx.error, + ) + return None, False + refreshed = await MetadataHelpers.refresh_metadata_if_needed( + self.db, + existing_id, + cached_metadata, + filepath, + base_dir, + state_hash, + mtime, + size, + self._write_scan_journal_entry, + ) + if refreshed: + await self._write_scan_journal_entry(filepath, base_dir, state_hash, mtime, size) + return tx, refreshed + + async def _extract_metadata_for_index_file( + self, + file_path: Path, + filepath: str, + state_hash: str, + fast: bool, + ) -> Result[Dict[str, Any]]: + if fast: + return Result.Ok({}) + + metadata_result = await self.metadata.get_metadata(filepath, scan_id=self._current_scan_id) + if metadata_result.ok: + quality = metadata_result.meta.get("quality", "full") + if quality in ("degraded", "partial"): + logger.debug( + "Metadata extraction completed with degraded quality for %s: %s", + file_path, + quality, + ) + await MetadataHelpers.store_metadata_cache(self.db, filepath, state_hash, metadata_result) + return metadata_result + + if metadata_result.code == ErrorCode.FFPROBE_ERROR: + self._log_scan_event( + logging.WARNING, + "FFprobe error during metadata extraction", + filepath=str(file_path), + tool="ffprobe", + error=metadata_result.error, + code=metadata_result.code, + ) + else: + self._log_scan_event( + logging.WARNING, + "Metadata extraction issue", + filepath=str(file_path), + error=metadata_result.error, + code=metadata_result.code, + ) + return MetadataHelpers.metadata_error_payload(metadata_result, str(file_path)) + + async def _try_incremental_refresh_with_metadata( + self, + incremental: bool, + existing_id: int, + existing_mtime: int, + mtime: int, + filepath: str, + state_hash: str, + base_dir: str, + size: int, + metadata_result: Result[Dict[str, Any]], + ) -> Optional[Result[Dict[str, str]]]: + if not (incremental and existing_id and existing_mtime == mtime): + return None + + refreshed = False + try: + tx_state, refreshed = await self._run_incremental_metadata_refresh_locked( + existing_id, + metadata_result, + filepath, + base_dir, + state_hash, + mtime, + size, + ) + except sqlite3.OperationalError: + return Result.Err("DB_BUSY", "Database busy while refreshing metadata") + + if not tx_state or not tx_state.ok: + return Result.Err("DB_ERROR", tx_state.error or "Commit failed") + action = "skipped_refresh" if refreshed else "skipped" + return Result.Ok({"action": action}) + + async def _run_incremental_metadata_refresh_locked( + self, + existing_id: int, + metadata_result: Result[Dict[str, Any]], + filepath: str, + base_dir: str, + state_hash: str, + mtime: int, + size: int, + ) -> tuple[Any, bool]: + async with self.db.lock_for_asset(existing_id): + try: + return await self._run_incremental_metadata_refresh_tx( + existing_id, + metadata_result, + filepath, + base_dir, + state_hash, mtime, size, - file_path, - metadata_result, - source=source, - root_id=root_id, - write_metadata=False, - skip_lock=True, ) - if result.ok: - await self._write_scan_journal_entry(filepath, base_dir, state_hash, mtime, size) + except sqlite3.OperationalError as exc: + if self.db._is_locked_error(exc): + logger.warning( + "Database busy while refreshing metadata (asset=%s): %s", + existing_id, + exc, + ) + raise + + async def _run_incremental_metadata_refresh_tx( + self, + existing_id: int, + metadata_result: Result[Dict[str, Any]], + filepath: str, + base_dir: str, + state_hash: str, + mtime: int, + size: int, + ) -> tuple[Any, bool]: + refreshed = False + async with self.db.atransaction(mode="immediate") as tx: + if not tx.ok: + return tx, refreshed + refreshed = await MetadataHelpers.refresh_metadata_if_needed( + self.db, + existing_id, + metadata_result, + filepath, + base_dir, + state_hash, + mtime, + size, + self._write_scan_journal_entry, + ) + if refreshed: + await self._write_scan_journal_entry(filepath, base_dir, state_hash, mtime, size) + return tx, refreshed + + async def _insert_new_asset_for_index_file( + self, + filename: str, + subfolder: str, + filepath: str, + kind: FileKind, + mtime: int, + size: int, + file_path: Path, + metadata_result: Result[Dict[str, Any]], + base_dir: str, + state_hash: str, + source: str, + root_id: Optional[str], + ) -> Result[Dict[str, str]]: + try: + tx_state, result = await self._insert_new_asset_tx( + filename, + subfolder, + filepath, + kind, + mtime, + size, + file_path, + metadata_result, + base_dir, + state_hash, + source, + root_id, + ) except sqlite3.OperationalError as exc: if self.db._is_locked_error(exc): logger.warning("Database busy while inserting asset %s: %s", filepath, exc) return Result.Err("DB_BUSY", "Database busy while inserting asset") raise + if not tx_state or not tx_state.ok: return Result.Err("DB_ERROR", tx_state.error or "Commit failed") if result.ok and isinstance(result.data, dict): @@ -1816,6 +2738,42 @@ class IndexScanner: await self._write_metadata_row(int(asset_id), metadata_result, filepath=filepath) return result + async def _insert_new_asset_tx( + self, + filename: str, + subfolder: str, + filepath: str, + kind: FileKind, + mtime: int, + size: int, + file_path: Path, + metadata_result: Result[Dict[str, Any]], + base_dir: str, + state_hash: str, + source: str, + root_id: Optional[str], + ) -> tuple[Any, Result[Dict[str, str]]]: + async with self.db.atransaction(mode="immediate") as tx: + if not tx.ok: + return tx, Result.Err("DB_ERROR", tx.error or "Failed to begin transaction") + result = await self._add_asset( + filename, + subfolder, + filepath, + kind, + mtime, + size, + file_path, + metadata_result, + source=source, + root_id=root_id, + write_metadata=False, + skip_lock=True, + ) + if result.ok: + await self._write_scan_journal_entry(filepath, base_dir, state_hash, mtime, size) + return tx, result + async def _add_asset( self, filename: str, @@ -1834,15 +2792,7 @@ class IndexScanner: """ Add new asset to database. """ - width = None - height = None - duration = None - - if metadata_result.ok and metadata_result.data: - meta = metadata_result.data - width = meta.get("width") - height = meta.get("height") - duration = meta.get("duration") + width, height, duration = self._asset_dimensions_from_metadata(metadata_result) # Insert into assets table (NO workflow fields here) insert_result = await self.db.aexecute( @@ -1873,24 +2823,50 @@ class IndexScanner: asset_id = insert_result.data if insert_result.ok else None if not asset_id: return Result.Err("INSERT_FAILED", "Failed to get inserted asset ID") - if write_metadata and not skip_lock: - metadata_write = await MetadataHelpers.write_asset_metadata_row( - self.db, - asset_id, - metadata_result, - filepath=filepath, - ) - if not metadata_write.ok: - self._log_scan_event( - logging.WARNING, - "Failed to insert metadata row", - asset_id=asset_id, - error=metadata_write.error, - stage="metadata_write" - ) + await self._write_asset_metadata_if_needed( + asset_id, + metadata_result, + filepath=filepath, + write_metadata=write_metadata, + skip_lock=skip_lock, + ) return Result.Ok({"action": "added", "asset_id": asset_id}) + @staticmethod + def _asset_dimensions_from_metadata(metadata_result: Result[Dict[str, Any]]) -> tuple[Any, Any, Any]: + if metadata_result.ok and metadata_result.data: + meta = metadata_result.data + return meta.get("width"), meta.get("height"), meta.get("duration") + return None, None, None + + async def _write_asset_metadata_if_needed( + self, + asset_id: Any, + metadata_result: Result[Dict[str, Any]], + *, + filepath: str, + write_metadata: bool, + skip_lock: bool, + ) -> None: + if not write_metadata or skip_lock: + return + metadata_write = await MetadataHelpers.write_asset_metadata_row( + self.db, + asset_id, + metadata_result, + filepath=filepath, + ) + if metadata_write.ok: + return + self._log_scan_event( + logging.WARNING, + "Failed to insert metadata row", + asset_id=asset_id, + error=metadata_write.error, + stage="metadata_write" + ) + async def _update_asset( self, asset_id: int, diff --git a/mjr_am_backend/features/index/searcher.py b/mjr_am_backend/features/index/searcher.py index da33d04d..7187b6fd 100644 --- a/mjr_am_backend/features/index/searcher.py +++ b/mjr_am_backend/features/index/searcher.py @@ -69,6 +69,22 @@ def _build_filter_clauses(filters: Optional[Dict[str, Any]], alias: str = "a") - params: List[Any] = [] if not filters: return clauses, params + _append_kind_and_source_filters(filters, alias, clauses, params) + _append_extension_filter(filters, alias, clauses, params) + _append_numeric_range_filters(filters, alias, clauses, params) + _append_workflow_type_filter(filters, clauses, params) + _append_has_workflow_filter(filters, clauses) + _append_mtime_filters(filters, alias, clauses, params) + _append_exclude_root_filter(filters, alias, clauses, params) + return clauses, params + + +def _append_kind_and_source_filters( + filters: Dict[str, Any], + alias: str, + clauses: List[str], + params: List[Any], +) -> None: kind = filters.get("kind") if isinstance(kind, str) and kind: clauses.append(f"AND {alias}.kind = ?") @@ -77,103 +93,119 @@ def _build_filter_clauses(filters: Optional[Dict[str, Any]], alias: str = "a") - if isinstance(source, str) and source.strip(): clauses.append(f"AND LOWER({alias}.source) = ?") params.append(source.strip().lower()) + + +def _append_extension_filter(filters: Dict[str, Any], alias: str, clauses: List[str], params: List[Any]) -> None: extensions = filters.get("extensions") - if isinstance(extensions, list): - normalized_exts = [] - for ext in extensions: - norm = _normalize_extension(ext) - if norm: - normalized_exts.append(norm) - if normalized_exts: - placeholders = ", ".join("?" for _ in normalized_exts) - clauses.append(f"AND LOWER({alias}.ext) IN ({placeholders})") - params.extend(normalized_exts) + if not isinstance(extensions, list): + return + normalized_exts = [_normalize_extension(ext) for ext in extensions] + normalized_exts = [ext for ext in normalized_exts if ext] + if not normalized_exts: + return + placeholders = ", ".join("?" for _ in normalized_exts) + clauses.append(f"AND LOWER({alias}.ext) IN ({placeholders})") + params.extend(normalized_exts) + + +def _append_numeric_range_filters(filters: Dict[str, Any], alias: str, clauses: List[str], params: List[Any]) -> None: + int_specs = [ + ("min_size_bytes", f"AND COALESCE({alias}.size, 0) >= ?"), + ("max_size_bytes", f"AND COALESCE({alias}.size, 0) <= ?"), + ("min_width", f"AND COALESCE({alias}.width, 0) >= ?"), + ("min_height", f"AND COALESCE({alias}.height, 0) >= ?"), + ("max_width", f"AND COALESCE({alias}.width, 0) <= ?"), + ("max_height", f"AND COALESCE({alias}.height, 0) <= ?"), + ] if "min_rating" in filters: clauses.append("AND COALESCE(m.rating, 0) >= ?") params.append(filters["min_rating"]) - if "min_size_bytes" in filters: - clauses.append(f"AND COALESCE({alias}.size, 0) >= ?") - params.append(int(filters["min_size_bytes"])) - if "max_size_bytes" in filters: - clauses.append(f"AND COALESCE({alias}.size, 0) <= ?") - params.append(int(filters["max_size_bytes"])) - if "min_width" in filters: - clauses.append(f"AND COALESCE({alias}.width, 0) >= ?") - params.append(int(filters["min_width"])) - if "min_height" in filters: - clauses.append(f"AND COALESCE({alias}.height, 0) >= ?") - params.append(int(filters["min_height"])) - if "max_width" in filters: - clauses.append(f"AND COALESCE({alias}.width, 0) <= ?") - params.append(int(filters["max_width"])) - if "max_height" in filters: - clauses.append(f"AND COALESCE({alias}.height, 0) <= ?") - params.append(int(filters["max_height"])) - if "workflow_type" in filters: - raw = str(filters.get("workflow_type") or "").strip().upper() - alias_map = { - "T2I": ["T2I"], - "I2I": ["I2I"], - "T2V": ["T2V"], - "I2V": ["I2V"], - "V2V": ["V2V"], - "A2A": ["A2A"], - "TTS": ["TTS", "T2A"], - "UPSCL": ["UPSCL", "UPSCALE", "UPSCALER"], - "INPT": ["INPT", "INPUT", "LOAD_INPUT"], - "FLF": ["FLF", "FIRST_LAST_FRAME", "FIRST_FRAME_LAST_FRAME"], - } - variants = alias_map.get(raw, [raw] if raw else []) - if variants: - placeholders = ", ".join("?" for _ in variants) - clauses.append( - "AND UPPER(COALESCE(" - "json_extract(m.metadata_raw, '$.workflow_type'), " - "json_extract(m.metadata_raw, '$.geninfo.engine.type'), " - "json_extract(m.metadata_raw, '$.engine.type'), " - "''" - f")) IN ({placeholders})" - ) - params.extend(variants) - if "has_workflow" in filters: - # Backward-compatible: older/stale rows may have has_workflow=0 even though metadata_raw - # contains a workflow/prompt. Prefer not to hide such assets when filtering. - want_workflow = bool(filters.get("has_workflow")) - if want_workflow: - clauses.append( - "AND (" - "COALESCE(m.has_workflow, 0) = 1 " - "OR json_extract(m.metadata_raw, '$.workflow') IS NOT NULL " - "OR json_extract(m.metadata_raw, '$.prompt') IS NOT NULL " - "OR json_extract(m.metadata_raw, '$.parameters') IS NOT NULL" - ")" - ) - else: - clauses.append( - "AND (" - "COALESCE(m.has_workflow, 0) = 0 " - "AND json_extract(m.metadata_raw, '$.workflow') IS NULL " - "AND json_extract(m.metadata_raw, '$.prompt') IS NULL " - "AND json_extract(m.metadata_raw, '$.parameters') IS NULL" - ")" - ) + for key, clause in int_specs: + if key in filters: + clauses.append(clause) + params.append(int(filters[key])) + + +def _append_workflow_type_filter(filters: Dict[str, Any], clauses: List[str], params: List[Any]) -> None: + if "workflow_type" not in filters: + return + raw = str(filters.get("workflow_type") or "").strip().upper() + variants = _workflow_type_variants(raw) + if not variants: + return + placeholders = ", ".join("?" for _ in variants) + clauses.append( + "AND UPPER(COALESCE(" + "json_extract(m.metadata_raw, '$.workflow_type'), " + "json_extract(m.metadata_raw, '$.geninfo.engine.type'), " + "json_extract(m.metadata_raw, '$.engine.type'), " + "''" + f")) IN ({placeholders})" + ) + params.extend(variants) + + +def _workflow_type_variants(raw: str) -> List[str]: + alias_map = { + "T2I": ["T2I"], + "I2I": ["I2I"], + "T2V": ["T2V"], + "I2V": ["I2V"], + "V2V": ["V2V"], + "A2A": ["A2A"], + "TTS": ["TTS", "T2A"], + "UPSCL": ["UPSCL", "UPSCALE", "UPSCALER"], + "INPT": ["INPT", "INPUT", "LOAD_INPUT"], + "FLF": ["FLF", "FIRST_LAST_FRAME", "FIRST_FRAME_LAST_FRAME"], + } + return alias_map.get(raw, [raw] if raw else []) + + +def _append_has_workflow_filter(filters: Dict[str, Any], clauses: List[str]) -> None: + if "has_workflow" not in filters: + return + want_workflow = bool(filters.get("has_workflow")) + if want_workflow: + clauses.append( + "AND (" + "COALESCE(m.has_workflow, 0) = 1 " + "OR json_extract(m.metadata_raw, '$.workflow') IS NOT NULL " + "OR json_extract(m.metadata_raw, '$.prompt') IS NOT NULL " + "OR json_extract(m.metadata_raw, '$.parameters') IS NOT NULL" + ")" + ) + return + clauses.append( + "AND (" + "COALESCE(m.has_workflow, 0) = 0 " + "AND json_extract(m.metadata_raw, '$.workflow') IS NULL " + "AND json_extract(m.metadata_raw, '$.prompt') IS NULL " + "AND json_extract(m.metadata_raw, '$.parameters') IS NULL" + ")" + ) + + +def _append_mtime_filters(filters: Dict[str, Any], alias: str, clauses: List[str], params: List[Any]) -> None: if "mtime_start" in filters: clauses.append(f"AND {alias}.mtime >= ?") params.append(filters["mtime_start"]) if "mtime_end" in filters: clauses.append(f"AND {alias}.mtime < ?") params.append(filters["mtime_end"]) + + +def _append_exclude_root_filter(filters: Dict[str, Any], alias: str, clauses: List[str], params: List[Any]) -> None: exclude_root = filters.get("exclude_root") - if isinstance(exclude_root, str) and exclude_root.strip(): - try: - root = str(Path(exclude_root).resolve(strict=False)) - prefix = root.rstrip(os.sep) + os.sep - esc = prefix.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_") - clauses.append(f"AND NOT ({alias}.filepath = ? OR {alias}.filepath LIKE ? ESCAPE '\\')") - params.extend([root, f"{esc}%"]) - except Exception: - pass - return clauses, params + if not isinstance(exclude_root, str) or not exclude_root.strip(): + return + try: + root = str(Path(exclude_root).resolve(strict=False)) + prefix = root.rstrip(os.sep) + os.sep + esc = prefix.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_") + clauses.append(f"AND NOT ({alias}.filepath = ? OR {alias}.filepath LIKE ? ESCAPE '\\')") + params.extend([root, f"{esc}%"]) + except Exception: + return def _assert_safe_sql_fragment(fragment: str, *, label: str = "sql_fragment") -> None: @@ -187,6 +219,195 @@ def _assert_safe_sql_fragment(fragment: str, *, label: str = "sql_fragment") -> raise ValueError(f"Unsafe {label}") +def _normalize_pagination(limit: int, offset: int) -> tuple[int, int]: + limit_i = max(0, min(SEARCH_MAX_LIMIT, int(limit))) + offset_i = max(0, min(SEARCH_MAX_OFFSET, int(offset))) + return limit_i, offset_i + + +def _escape_like_pattern(pattern: str) -> str: + return pattern.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_") + + +def _resolve_search_roots(roots: List[str]) -> List[str]: + cleaned_roots: List[str] = [] + for raw in roots or []: + if not raw: + continue + try: + cleaned_roots.append(str(Path(raw).resolve())) + except Exception: + continue + return cleaned_roots + + +def _build_roots_where_clause(roots: List[str], alias: str = "a") -> Tuple[str, List[Any]]: + parts: List[str] = [] + params: List[Any] = [] + for root in roots: + prefix = root.rstrip(os.sep) + os.sep + escaped_prefix = _escape_like_pattern(prefix) + parts.append(f"({alias}.filepath = ? OR {alias}.filepath LIKE ? ESCAPE '\\')") + params.extend([root, f"{escaped_prefix}%"]) + clause = "(" + " OR ".join(parts) + ")" + _assert_safe_sql_fragment(clause, label="roots_where") + return clause, params + + +def _normalize_month_range(month_start: int, month_end: int) -> Optional[Tuple[int, int]]: + try: + start_i = int(month_start) + end_i = int(month_end) + except Exception: + return None + if start_i <= 0 or end_i <= 0 or end_i <= start_i: + return None + return start_i, end_i + + +def _sanitize_histogram_filters(filters: Optional[Dict[str, Any]]) -> Dict[str, Any]: + safe_filters = dict(filters or {}) + safe_filters.pop("mtime_start", None) + safe_filters.pop("mtime_end", None) + return safe_filters + + +def _build_histogram_query( + roots_clause: str, + roots_params: List[Any], + start_i: int, + end_i: int, + filters: Optional[Dict[str, Any]], +) -> Tuple[str, Tuple[Any, ...]]: + sql_parts = [ + """ + SELECT + strftime('%Y-%m-%d', a.mtime, 'unixepoch', 'localtime') AS day, + COUNT(*) AS count + FROM assets a + LEFT JOIN asset_metadata m ON a.id = m.asset_id + WHERE + """, + roots_clause, + "AND a.mtime >= ? AND a.mtime < ?", + ] + params: List[Any] = [] + params.extend(roots_params) + params.extend([start_i, end_i]) + filter_clauses, filter_params = _build_filter_clauses(_sanitize_histogram_filters(filters)) + sql_parts.extend(filter_clauses) + params.extend(filter_params) + sql_parts.append("GROUP BY day ORDER BY day ASC") + return " ".join(sql_parts), tuple(params) + + +def _coerce_histogram_days(rows: Any) -> Dict[str, int]: + days: Dict[str, int] = {} + for row in rows or []: + try: + day = str(row.get("day") or "").strip() + count = int(row.get("count") or 0) + except Exception: + continue + if day: + days[day] = max(0, count) + return days + + +def _normalize_asset_ids(asset_ids: List[int]) -> List[int]: + cleaned: List[int] = [] + seen = set() + for raw in asset_ids or []: + try: + n = int(raw) + except Exception: + continue + if n <= 0 or n in seen: + continue + seen.add(n) + cleaned.append(n) + if len(cleaned) >= MAX_ASSET_BATCH_IDS: + break + return cleaned + + +def _map_assets_by_id(rows: Any, hydrator) -> Dict[int, Dict[str, Any]]: + by_id: Dict[int, Dict[str, Any]] = {} + for row in rows or []: + try: + asset = hydrator(dict(row)) + raw_id = asset.get("id") + if raw_id is None: + continue + by_id[int(raw_id)] = asset + except Exception: + continue + return by_id + + +def _assets_in_requested_order(cleaned: List[int], by_id: Dict[int, Dict[str, Any]]) -> List[Dict[str, Any]]: + out: List[Dict[str, Any]] = [] + for aid in cleaned: + item = by_id.get(aid) + if item: + out.append(item) + return out + + +def _normalize_lookup_filepaths(filepaths: List[str]) -> List[str]: + cleaned = [str(p) for p in (filepaths or []) if p] + if len(cleaned) > MAX_FILEPATH_LOOKUP: + return cleaned[:MAX_FILEPATH_LOOKUP] + return cleaned + + +def _hydrate_lookup_row(row: Dict[str, Any]) -> Optional[tuple[str, Dict[str, Any]]]: + fp = row.get("filepath") + if not fp: + return None + item = dict(row) + tags_raw = item.get("tags") + if tags_raw: + try: + item["tags"] = json.loads(tags_raw) if isinstance(tags_raw, str) else tags_raw + except Exception: + item["tags"] = [] + else: + item["tags"] = [] + return str(fp), item + + +def _map_lookup_rows(rows: Any) -> Dict[str, Dict[str, Any]]: + out: Dict[str, Dict[str, Any]] = {} + for row in rows or []: + if not isinstance(row, dict): + continue + hydrated = _hydrate_lookup_row(row) + if hydrated is None: + continue + fp, item = hydrated + out[fp] = item + return out + + +def _hydrate_search_rows(rows: List[Dict[str, Any]], *, include_highlight: bool) -> List[Dict[str, Any]]: + assets: List[Dict[str, Any]] = [] + for row in rows: + asset = dict(row) + if asset.get("tags"): + try: + asset["tags"] = json.loads(asset["tags"]) + except (ValueError, json.JSONDecodeError, TypeError): + asset["tags"] = [] + else: + asset["tags"] = [] + asset.setdefault("tags_text", "") + if include_highlight: + asset["highlight"] = asset.get("highlight") or None + assets.append(asset) + return assets + + class IndexSearcher: """ Handles asset search and retrieval operations. @@ -239,11 +460,399 @@ class IndexSearcher: ) if not res.ok: return Result.Ok([]) - + terms = [str(r["term"]) for r in res.data or []] return Result.Ok(terms) except Exception: return Result.Ok([]) + + def _validate_search_query(self, query: str) -> Optional[Result[Any]]: + if not query or not query.strip(): + return Result.Err("EMPTY_QUERY", "Search query cannot be empty") + validation = self._validate_search_input(query) + if validation and not validation.ok: + return validation + return None + + def _filter_clauses( + self, + filters: Optional[Dict[str, Any]], + *, + assert_safe: bool = False, + ) -> Tuple[List[str], List[Any]]: + filter_clauses, filter_params = _build_filter_clauses(filters) + if assert_safe: + for clause in filter_clauses: + _assert_safe_sql_fragment(clause, label="filter_clause") + return filter_clauses, filter_params + + async def _run_search_query_rows( + self, + sql: str, + params: List[Any], + *, + failure_message: str, + ) -> Result[List[Dict[str, Any]]]: + result = await self.db.aquery(sql, tuple(params)) + if result.ok: + return Result.Ok(result.data or []) + if self._is_malformed_match_error(result.error): + return Result.Err("INVALID_INPUT", "Invalid search query syntax") + return Result.Err("SEARCH_FAILED", result.error or failure_message) + + @staticmethod + def _search_rows_total(data: Any) -> tuple[List[Dict[str, Any]], Optional[int]]: + if not isinstance(data, dict): + return [], None + rows = data.get("rows") + total = data.get("total") + if not isinstance(rows, list): + rows = [] + return rows, total + + def _build_search_payload( + self, + *, + assets: List[Dict[str, Any]], + limit: int, + offset: int, + query: str, + include_total: bool, + total: Optional[int], + sort: Optional[str] = None, + ) -> Dict[str, Any]: + payload: Dict[str, Any] = { + "assets": assets, + "limit": limit, + "offset": offset, + "query": query, + "total": int(total or 0) if include_total else None, + } + if sort is not None: + payload["sort"] = _normalize_sort_key(sort) + return payload + + async def _search_global_browse_rows( + self, + *, + limit: int, + offset: int, + filters: Optional[Dict[str, Any]], + include_total: bool, + metadata_tags_text_clause: str, + ) -> Result[Dict[str, Any]]: + sql_parts = [ + f""" + SELECT + a.id, a.filename, a.subfolder, a.filepath, a.kind, + a.source, a.root_id, + a.width, a.height, a.duration, a.size, a.mtime, + COALESCE(m.rating, 0) as rating, + COALESCE(m.tags, '[]') as tags, +{metadata_tags_text_clause} m.has_workflow as has_workflow, + m.has_generation_data as has_generation_data, + NULL as generation_time_ms, + NULL as file_creation_time, + NULL as file_birth_time + FROM assets a + LEFT JOIN asset_metadata m ON a.id = m.asset_id + WHERE 1=1 + """ + ] + params: List[Any] = [] + + filter_clauses, filter_params = self._filter_clauses(filters) + sql_parts.extend(filter_clauses) + params.extend(filter_params) + sql_parts.append("ORDER BY a.mtime DESC") + sql_parts.append("LIMIT ? OFFSET ?") + params.extend([limit, offset]) + + rows_res = await self._run_search_query_rows( + " ".join(sql_parts), + params, + failure_message="Search query failed", + ) + if not rows_res.ok: + return Result.Err(rows_res.code, rows_res.error or "Search query failed") + rows = rows_res.data or [] + + total: Optional[int] = None + if include_total: + count_sql = "SELECT COUNT(*) as total FROM assets a LEFT JOIN asset_metadata m ON a.id = m.asset_id WHERE 1=1" + count_params: List[Any] = [] + if filter_clauses: + count_sql += " " + " ".join(filter_clauses) + count_params.extend(filter_params) + count_result = await self.db.aquery(count_sql, tuple(count_params)) + total = count_result.data[0]["total"] if count_result.ok and count_result.data else 0 + return Result.Ok({"rows": rows, "total": total}) + + async def _search_global_fts_rows( + self, + *, + fts_query: str, + limit: int, + offset: int, + filters: Optional[Dict[str, Any]], + include_total: bool, + metadata_tags_text_clause: str, + ) -> Result[Dict[str, Any]]: + total_field = "COUNT(*) OVER() as _total," if include_total else "" + sql_parts = [self._global_fts_select_sql(total_field, metadata_tags_text_clause)] + params: List[Any] = [fts_query, fts_query] + + filter_clauses, filter_params = self._filter_clauses(filters) + sql_parts.extend(filter_clauses) + params.extend(filter_params) + sql_parts.append("ORDER BY rank") + sql_parts.append("LIMIT ? OFFSET ?") + params.extend([limit, offset]) + + rows_res = await self._run_search_query_rows( + " ".join(sql_parts), + params, + failure_message="Search query failed", + ) + if not rows_res.ok: + return Result.Err(rows_res.code, rows_res.error or "Search query failed") + rows = rows_res.data or [] + + total: Optional[int] = None + if include_total and rows and "_total" in rows[0]: + total = rows[0]["_total"] + if include_total and total is None: + total = await self._global_fts_total_count(fts_query, filter_clauses, filter_params) + return Result.Ok({"rows": rows, "total": total}) + + @staticmethod + def _global_fts_select_sql(total_field: str, metadata_tags_text_clause: str) -> str: + return f""" + WITH matches AS ( + SELECT rowid AS asset_id, bm25(assets_fts) AS rank + FROM assets_fts + WHERE assets_fts MATCH ? + + UNION ALL + + SELECT rowid AS asset_id, (bm25(asset_metadata_fts) + 8.0) AS rank + FROM asset_metadata_fts + WHERE asset_metadata_fts MATCH ? + ), + best AS ( + SELECT asset_id, MIN(rank) AS rank + FROM matches + GROUP BY asset_id + ) + SELECT + {total_field} + a.id, a.filename, a.subfolder, a.filepath, a.kind, + a.source, a.root_id, + a.width, a.height, a.duration, a.size, a.mtime, + COALESCE(m.rating, 0) as rating, + COALESCE(m.tags, '[]') as tags, +{metadata_tags_text_clause} m.has_workflow as has_workflow, + m.has_generation_data as has_generation_data, + NULL as generation_time_ms, + NULL as file_creation_time, + NULL as file_birth_time, + best.rank as rank + FROM best + JOIN assets a ON best.asset_id = a.id + LEFT JOIN asset_metadata m ON a.id = m.asset_id + WHERE 1=1 + """ + + async def _global_fts_total_count( + self, + fts_query: str, + filter_clauses: List[str], + filter_params: List[Any], + ) -> int: + count_sql = """ + WITH matches AS ( + SELECT rowid AS asset_id + FROM assets_fts + WHERE assets_fts MATCH ? + + UNION + + SELECT rowid AS asset_id + FROM asset_metadata_fts + WHERE asset_metadata_fts MATCH ? + ) + SELECT COUNT(*) as total + FROM (SELECT DISTINCT asset_id FROM matches) t + JOIN assets a ON t.asset_id = a.id + LEFT JOIN asset_metadata m ON a.id = m.asset_id + WHERE 1=1 + """ + count_params: List[Any] = [fts_query, fts_query] + if filter_clauses: + count_sql += " " + " ".join(filter_clauses) + count_params.extend(filter_params) + count_result = await self.db.aquery(count_sql, tuple(count_params)) + if count_result.ok and count_result.data: + return count_result.data[0]["total"] + return 0 + + async def _search_scoped_browse_rows( + self, + *, + roots_clause: str, + roots_params: List[Any], + limit: int, + offset: int, + filters: Optional[Dict[str, Any]], + include_total: bool, + metadata_tags_text_clause: str, + sort: Optional[str], + ) -> Result[Dict[str, Any]]: + sql_parts = [ + f""" + SELECT + a.id, a.filename, a.subfolder, a.filepath, a.kind, + a.width, a.height, a.duration, a.size, a.mtime, + COALESCE(m.rating, 0) as rating, + COALESCE(m.tags, '[]') as tags, +{metadata_tags_text_clause} m.has_workflow as has_workflow, + m.has_generation_data as has_generation_data, + NULL as generation_time_ms, + NULL as file_creation_time, + NULL as file_birth_time + FROM assets a + LEFT JOIN asset_metadata m ON a.id = m.asset_id + WHERE {roots_clause} + """ + ] + params: List[Any] = list(roots_params) + + filter_clauses, filter_params = self._filter_clauses(filters, assert_safe=True) + sql_parts.extend(filter_clauses) + params.extend(filter_params) + sql_parts.append(_build_sort_sql(sort, table_alias="a")) + sql_parts.append("LIMIT ? OFFSET ?") + params.extend([limit, offset]) + + rows_res = await self._run_search_query_rows( + " ".join(sql_parts), + params, + failure_message="Scoped search query failed", + ) + if not rows_res.ok: + return Result.Err(rows_res.code, rows_res.error or "Scoped search query failed") + rows = rows_res.data or [] + + total: Optional[int] = None + if include_total: + count_sql = f""" + SELECT COUNT(*) as total + FROM assets a + LEFT JOIN asset_metadata m ON a.id = m.asset_id + WHERE {roots_clause} + """ + count_params: List[Any] = list(roots_params) + if filter_clauses: + count_sql += " " + " ".join(filter_clauses) + count_params.extend(filter_params) + count_result = await self.db.aquery(count_sql, tuple(count_params)) + total = count_result.data[0]["total"] if count_result.ok and count_result.data else 0 + return Result.Ok({"rows": rows, "total": total}) + + async def _search_scoped_fts_rows( + self, + *, + fts_query: str, + roots_clause: str, + roots_params: List[Any], + limit: int, + offset: int, + filters: Optional[Dict[str, Any]], + include_total: bool, + metadata_tags_text_clause: str, + sort: Optional[str], + ) -> Result[Dict[str, Any]]: + sql_parts = [ + f""" + WITH matches AS ( + SELECT rowid AS asset_id, bm25(assets_fts) AS rank + FROM assets_fts + WHERE assets_fts MATCH ? + + UNION ALL + + SELECT rowid AS asset_id, (bm25(asset_metadata_fts) + 8.0) AS rank + FROM asset_metadata_fts + WHERE asset_metadata_fts MATCH ? + ), + best AS ( + SELECT asset_id, MIN(rank) AS rank + FROM matches + GROUP BY asset_id + ) + SELECT + a.id, a.filename, a.subfolder, a.filepath, a.kind, + a.width, a.height, a.duration, a.size, a.mtime, + COALESCE(m.rating, 0) as rating, + COALESCE(m.tags, '[]') as tags, +{metadata_tags_text_clause} m.has_workflow as has_workflow, + m.has_generation_data as has_generation_data, + NULL as generation_time_ms, + NULL as file_creation_time, + NULL as file_birth_time, best.rank as rank + FROM best + JOIN assets a ON best.asset_id = a.id + LEFT JOIN asset_metadata m ON a.id = m.asset_id + WHERE {roots_clause} + """ + ] + params: List[Any] = [fts_query, fts_query] + params.extend(roots_params) + + filter_clauses, filter_params = self._filter_clauses(filters, assert_safe=True) + sql_parts.extend(filter_clauses) + params.extend(filter_params) + sql_parts.append(_build_sort_sql(sort, table_alias="a", rank_alias="best.rank")) + sql_parts.append("LIMIT ? OFFSET ?") + params.extend([limit, offset]) + + rows_res = await self._run_search_query_rows( + " ".join(sql_parts), + params, + failure_message="Scoped search query failed", + ) + if not rows_res.ok: + return Result.Err(rows_res.code, rows_res.error or "Scoped search query failed") + rows = rows_res.data or [] + + total: Optional[int] = None + if include_total: + count_sql = f""" + WITH matches AS ( + SELECT rowid AS asset_id + FROM assets_fts + WHERE assets_fts MATCH ? + + UNION + + SELECT rowid AS asset_id + FROM asset_metadata_fts + WHERE asset_metadata_fts MATCH ? + ) + SELECT COUNT(*) as total + FROM (SELECT DISTINCT asset_id FROM matches) t + JOIN assets a ON t.asset_id = a.id + LEFT JOIN asset_metadata m ON a.id = m.asset_id + WHERE {roots_clause} + """ + count_params: List[Any] = [fts_query, fts_query] + count_params.extend(roots_params) + if filter_clauses: + count_sql += " " + " ".join(filter_clauses) + count_params.extend(filter_params) + count_result = await self.db.aquery(count_sql, tuple(count_params)) + total = count_result.data[0]["total"] if count_result.ok and count_result.data else 0 + return Result.Ok({"rows": rows, "total": total}) async def search( self, query: str, @@ -255,173 +864,52 @@ class IndexSearcher: """ Search assets using FTS5 or browse mode when query is '*'. """ - limit = max(0, min(SEARCH_MAX_LIMIT, int(limit))) - offset = max(0, min(SEARCH_MAX_OFFSET, int(offset))) - - if not query or not query.strip(): - return Result.Err("EMPTY_QUERY", "Search query cannot be empty") - - validation = self._validate_search_input(query) - if validation and not validation.ok: + limit, offset = _normalize_pagination(limit, offset) + validation = self._validate_search_query(query) + if validation: return validation include_total = bool(include_total) logger.debug("Searching for: %s (limit=%s, offset=%s)", query, limit, offset) metadata_tags_text_clause = self._build_tags_text_clause() - is_browse_all = query.strip() == "*" - if is_browse_all: - sql_parts = [ - f""" - SELECT - a.id, a.filename, a.subfolder, a.filepath, a.kind, - a.source, a.root_id, - a.width, a.height, a.duration, a.size, a.mtime, - COALESCE(m.rating, 0) as rating, - COALESCE(m.tags, '[]') as tags, -{metadata_tags_text_clause} m.has_workflow as has_workflow, - m.has_generation_data as has_generation_data, - NULL as generation_time_ms, - NULL as file_creation_time, - NULL as file_birth_time - FROM assets a - LEFT JOIN asset_metadata m ON a.id = m.asset_id - WHERE 1=1 - """ - ] - params: List[Any] = [] - - filter_clauses, filter_params = _build_filter_clauses(filters) - sql_parts.extend(filter_clauses) - params.extend(filter_params) - - sql_parts.append("ORDER BY a.mtime DESC") - sql_parts.append("LIMIT ? OFFSET ?") - params.extend([limit, offset]) - - result = await self.db.aquery(" ".join(sql_parts), tuple(params)) - if not result.ok: - if self._is_malformed_match_error(result.error): - return Result.Err("INVALID_INPUT", "Invalid search query syntax") - return Result.Err("SEARCH_FAILED", result.error or "Search query failed") - rows = result.data or [] - - total = None - if include_total: - count_sql = "SELECT COUNT(*) as total FROM assets a LEFT JOIN asset_metadata m ON a.id = m.asset_id WHERE 1=1" - count_params: List[Any] = [] - if filter_clauses: - count_sql += " " + " ".join(filter_clauses) - count_params.extend(filter_params) - count_result = await self.db.aquery(count_sql, tuple(count_params)) - total = count_result.data[0]["total"] if count_result.ok and count_result.data else 0 + if query.strip() == "*": + rows_total_res = await self._search_global_browse_rows( + limit=limit, + offset=offset, + filters=filters, + include_total=include_total, + metadata_tags_text_clause=metadata_tags_text_clause, + ) else: fts_query = self._sanitize_fts_query(query) if not fts_query: return Result.Err("INVALID_INPUT", "Invalid search query syntax") - total_field = "COUNT(*) OVER() as _total," if include_total else "" - sql_parts = [ - f""" - WITH matches AS ( - SELECT rowid AS asset_id, bm25(assets_fts) AS rank - FROM assets_fts - WHERE assets_fts MATCH ? + rows_total_res = await self._search_global_fts_rows( + fts_query=fts_query, + limit=limit, + offset=offset, + filters=filters, + include_total=include_total, + metadata_tags_text_clause=metadata_tags_text_clause, + ) - UNION ALL - - SELECT rowid AS asset_id, (bm25(asset_metadata_fts) + 8.0) AS rank - FROM asset_metadata_fts - WHERE asset_metadata_fts MATCH ? - ), - best AS ( - SELECT asset_id, MIN(rank) AS rank - FROM matches - GROUP BY asset_id - ) - SELECT - {total_field} - a.id, a.filename, a.subfolder, a.filepath, a.kind, - a.source, a.root_id, - a.width, a.height, a.duration, a.size, a.mtime, - COALESCE(m.rating, 0) as rating, - COALESCE(m.tags, '[]') as tags, -{metadata_tags_text_clause} m.has_workflow as has_workflow, - m.has_generation_data as has_generation_data, - NULL as generation_time_ms, - NULL as file_creation_time, - NULL as file_birth_time, - best.rank as rank - FROM best - JOIN assets a ON best.asset_id = a.id - LEFT JOIN asset_metadata m ON a.id = m.asset_id - WHERE 1=1 - """ - ] - params = [fts_query, fts_query] - - filter_clauses, filter_params = _build_filter_clauses(filters) - sql_parts.extend(filter_clauses) - params.extend(filter_params) - - sql_parts.append("ORDER BY rank") - sql_parts.append("LIMIT ? OFFSET ?") - params.extend([limit, offset]) - - result = await self.db.aquery(" ".join(sql_parts), tuple(params)) - if not result.ok: - if self._is_malformed_match_error(result.error): - return Result.Err("INVALID_INPUT", "Invalid search query syntax") - return Result.Err("SEARCH_FAILED", result.error or "Search query failed") - rows = result.data or [] - - total = None - if include_total and rows and "_total" in rows[0]: - total = rows[0]["_total"] - if include_total and total is None: - count_sql = """ - WITH matches AS ( - SELECT rowid AS asset_id - FROM assets_fts - WHERE assets_fts MATCH ? - - UNION - - SELECT rowid AS asset_id - FROM asset_metadata_fts - WHERE asset_metadata_fts MATCH ? - ) - SELECT COUNT(*) as total - FROM (SELECT DISTINCT asset_id FROM matches) t - JOIN assets a ON t.asset_id = a.id - LEFT JOIN asset_metadata m ON a.id = m.asset_id - WHERE 1=1 - """ - count_params2 = [fts_query, fts_query] - if filter_clauses: - count_sql += " " + " ".join(filter_clauses) - count_params2.extend(filter_params) - count_result = await self.db.aquery(count_sql, tuple(count_params2)) - total = count_result.data[0]["total"] if count_result.ok and count_result.data else 0 - - assets = [] - for row in rows: - asset = dict(row) - if asset.get("tags"): - try: - asset["tags"] = json.loads(asset["tags"]) - except (ValueError, json.JSONDecodeError, TypeError): - asset["tags"] = [] - else: - asset["tags"] = [] - asset.setdefault("tags_text", "") - asset["highlight"] = asset.get("highlight") or None - assets.append(asset) + if not rows_total_res.ok: + return Result.Err(rows_total_res.code, rows_total_res.error or "Search query failed") + rows, total = self._search_rows_total(rows_total_res.data) + assets = _hydrate_search_rows(rows, include_highlight=True) logger.debug("Found %s results (total=%s)", len(assets), total if include_total else "skipped") - payload: Dict[str, Any] = {"assets": assets, "limit": limit, "offset": offset, "query": query} - payload["total"] = int(total or 0) if include_total else None - - return Result.Ok(payload) + return Result.Ok( + self._build_search_payload( + assets=assets, + limit=limit, + offset=offset, + query=query, + include_total=include_total, + total=total, + ) + ) async def search_scoped( self, query: str, @@ -437,218 +925,64 @@ class IndexSearcher: This is used for UI scopes like Outputs / Inputs / All without breaking the existing DB structure. """ - cleaned_roots: List[str] = [] - for r in roots or []: - if not r: - continue - try: - cleaned_roots.append(str(Path(r).resolve())) - except Exception: - continue - + cleaned_roots = _resolve_search_roots(roots) if not cleaned_roots: return Result.Err("INVALID_INPUT", "Missing or invalid roots") - # Reuse the existing search logic but inject a filepath prefix constraint. - limit = max(0, min(SEARCH_MAX_LIMIT, int(limit))) - offset = max(0, min(SEARCH_MAX_OFFSET, int(offset))) - - if not query or not query.strip(): - return Result.Err("EMPTY_QUERY", "Search query cannot be empty") - - validation = self._validate_search_input(query) - if validation and not validation.ok: + limit, offset = _normalize_pagination(limit, offset) + validation = self._validate_search_query(query) + if validation: return validation include_total = bool(include_total) logger.debug(f"Searching (scoped) for: {query} (limit={limit}, offset={offset}, roots={len(cleaned_roots)})") metadata_tags_text_clause = self._build_tags_text_clause() - def _escape_like_pattern(pattern: str) -> str: - """Escape LIKE special characters (% and _) for safe prefix matching.""" - return pattern.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_") - - def _roots_where() -> Tuple[str, List[Any]]: - parts = [] - params: List[Any] = [] - for root in cleaned_roots: - prefix = root.rstrip(os.sep) + os.sep - escaped_prefix = _escape_like_pattern(prefix) - parts.append("(a.filepath = ? OR a.filepath LIKE ? ESCAPE '\\')") - params.extend([root, f"{escaped_prefix}%"]) - clause = "(" + " OR ".join(parts) + ")" - _assert_safe_sql_fragment(clause, label="roots_where") - return clause, params - is_browse_all = query.strip() == "*" - - roots_clause, roots_params = _roots_where() - + roots_clause, roots_params = _build_roots_where_clause(cleaned_roots) if is_browse_all: - sql_parts = [ - f""" - SELECT - a.id, a.filename, a.subfolder, a.filepath, a.kind, - a.width, a.height, a.duration, a.size, a.mtime, - COALESCE(m.rating, 0) as rating, - COALESCE(m.tags, '[]') as tags, -{metadata_tags_text_clause} m.has_workflow as has_workflow, - m.has_generation_data as has_generation_data, - NULL as generation_time_ms, - NULL as file_creation_time, - NULL as file_birth_time - FROM assets a - LEFT JOIN asset_metadata m ON a.id = m.asset_id - WHERE {roots_clause} - """ - ] - params: List[Any] = [] - params.extend(roots_params) - - filter_clauses, filter_params = _build_filter_clauses(filters) - for c in filter_clauses: - _assert_safe_sql_fragment(c, label="filter_clause") - sql_parts.extend(filter_clauses) - params.extend(filter_params) - - sql_parts.append(_build_sort_sql(sort, table_alias="a")) - sql_parts.append("LIMIT ? OFFSET ?") - params.extend([limit, offset]) - - sql = " ".join(sql_parts) - result = await self.db.aquery(sql, tuple(params)) - if not result.ok: - if self._is_malformed_match_error(result.error): - return Result.Err("INVALID_INPUT", "Invalid search query syntax") - return Result.Err("SEARCH_FAILED", result.error or "Scoped search query failed") - - rows = result.data or [] - - total = None - if include_total: - count_sql = f""" - SELECT COUNT(*) as total - FROM assets a - LEFT JOIN asset_metadata m ON a.id = m.asset_id - WHERE {roots_clause} - """ - count_params: List[Any] = [] - count_params.extend(roots_params) - - if filter_clauses: - count_sql += " " + " ".join(filter_clauses) - count_params.extend(filter_params) - - count_result = await self.db.aquery(count_sql, tuple(count_params)) - total = count_result.data[0]["total"] if count_result.ok and count_result.data else 0 + rows_total_res = await self._search_scoped_browse_rows( + roots_clause=roots_clause, + roots_params=roots_params, + limit=limit, + offset=offset, + filters=filters, + include_total=include_total, + metadata_tags_text_clause=metadata_tags_text_clause, + sort=sort, + ) else: fts_query = self._sanitize_fts_query(query) if not fts_query: return Result.Err("INVALID_INPUT", "Invalid search query syntax") - sql_parts = [ - f""" - WITH matches AS ( - SELECT rowid AS asset_id, bm25(assets_fts) AS rank - FROM assets_fts - WHERE assets_fts MATCH ? + rows_total_res = await self._search_scoped_fts_rows( + fts_query=fts_query, + roots_clause=roots_clause, + roots_params=roots_params, + limit=limit, + offset=offset, + filters=filters, + include_total=include_total, + metadata_tags_text_clause=metadata_tags_text_clause, + sort=sort, + ) - UNION ALL + if not rows_total_res.ok: + return Result.Err(rows_total_res.code, rows_total_res.error or "Scoped search query failed") - SELECT rowid AS asset_id, (bm25(asset_metadata_fts) + 8.0) AS rank - FROM asset_metadata_fts - WHERE asset_metadata_fts MATCH ? - ), - best AS ( - SELECT asset_id, MIN(rank) AS rank - FROM matches - GROUP BY asset_id - ) - SELECT - a.id, a.filename, a.subfolder, a.filepath, a.kind, - a.width, a.height, a.duration, a.size, a.mtime, - COALESCE(m.rating, 0) as rating, - COALESCE(m.tags, '[]') as tags, -{metadata_tags_text_clause} m.has_workflow as has_workflow, - m.has_generation_data as has_generation_data, - NULL as generation_time_ms, - NULL as file_creation_time, - NULL as file_birth_time, best.rank as rank - FROM best - JOIN assets a ON best.asset_id = a.id - LEFT JOIN asset_metadata m ON a.id = m.asset_id - WHERE {roots_clause} - """ - ] - - params = [fts_query, fts_query] - params.extend(roots_params) - - filter_clauses, filter_params = _build_filter_clauses(filters) - for c in filter_clauses: - _assert_safe_sql_fragment(c, label="filter_clause") - sql_parts.extend(filter_clauses) - params.extend(filter_params) - - sql_parts.append(_build_sort_sql(sort, table_alias="a", rank_alias="best.rank")) - sql_parts.append("LIMIT ? OFFSET ?") - params.extend([limit, offset]) - - sql = " ".join(sql_parts) - result = await self.db.aquery(sql, tuple(params)) - if not result.ok: - if self._is_malformed_match_error(result.error): - return Result.Err("INVALID_INPUT", "Invalid search query syntax") - return Result.Err("SEARCH_FAILED", result.error or "Scoped search query failed") - - rows = result.data or [] - - total = None - if include_total: - count_sql = f""" - WITH matches AS ( - SELECT rowid AS asset_id - FROM assets_fts - WHERE assets_fts MATCH ? - - UNION - - SELECT rowid AS asset_id - FROM asset_metadata_fts - WHERE asset_metadata_fts MATCH ? - ) - SELECT COUNT(*) as total - FROM (SELECT DISTINCT asset_id FROM matches) t - JOIN assets a ON t.asset_id = a.id - LEFT JOIN asset_metadata m ON a.id = m.asset_id - WHERE {roots_clause} - """ - count_params2: List[Any] = [fts_query, fts_query] - count_params2.extend(roots_params) - - if filter_clauses: - count_sql += " " + " ".join(filter_clauses) - count_params2.extend(filter_params) - - count_result = await self.db.aquery(count_sql, tuple(count_params2)) - total = count_result.data[0]["total"] if count_result.ok and count_result.data else 0 - - assets = [] - for row in rows: - asset = dict(row) - if asset.get("tags"): - try: - asset["tags"] = json.loads(asset["tags"]) - except Exception: - asset["tags"] = [] - else: - asset["tags"] = [] - asset.setdefault("tags_text", "") - assets.append(asset) - - payload: Dict[str, Any] = {"assets": assets, "limit": limit, "offset": offset, "query": query} - payload["total"] = int(total or 0) if include_total else None - payload["sort"] = _normalize_sort_key(sort) - return Result.Ok(payload) + rows, total = self._search_rows_total(rows_total_res.data) + assets = _hydrate_search_rows(rows, include_highlight=False) + return Result.Ok( + self._build_search_payload( + assets=assets, + limit=limit, + offset=offset, + query=query, + include_total=include_total, + total=total, + sort=sort, + ) + ) async def has_assets_under_root(self, root: str) -> Result[bool]: """ @@ -690,87 +1024,20 @@ class IndexSearcher: - Uses localtime conversion to match the UI's date filters (which use local time). - Intended for calendar "days with assets" indicators (no query/FTS). """ - cleaned_roots: List[str] = [] - for r in roots or []: - if not r: - continue - try: - cleaned_roots.append(str(Path(r).resolve())) - except Exception: - continue - + cleaned_roots = _resolve_search_roots(roots) if not cleaned_roots: return Result.Err("INVALID_INPUT", "Missing or invalid roots") - try: - start_i = int(month_start) - end_i = int(month_end) - except Exception: + month_range = _normalize_month_range(month_start, month_end) + if month_range is None: return Result.Err("INVALID_INPUT", "Invalid month range") - - if start_i <= 0 or end_i <= 0 or end_i <= start_i: - return Result.Err("INVALID_INPUT", "Invalid month range") - - def _escape_like_pattern(pattern: str) -> str: - return pattern.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_") - - def _roots_where() -> Tuple[str, List[Any]]: - parts = [] - params: List[Any] = [] - for root in cleaned_roots: - prefix = root.rstrip(os.sep) + os.sep - escaped_prefix = _escape_like_pattern(prefix) - parts.append("(a.filepath = ? OR a.filepath LIKE ? ESCAPE '\\')") - params.extend([root, f"{escaped_prefix}%"]) - clause = "(" + " OR ".join(parts) + ")" - _assert_safe_sql_fragment(clause, label="roots_where") - return clause, params - - roots_clause, roots_params = _roots_where() - - # Do not allow callers to inject their own mtime window beyond the month we requested. - safe_filters = dict(filters or {}) - safe_filters.pop("mtime_start", None) - safe_filters.pop("mtime_end", None) - - sql_parts = [ - """ - SELECT - strftime('%Y-%m-%d', a.mtime, 'unixepoch', 'localtime') AS day, - COUNT(*) AS count - FROM assets a - LEFT JOIN asset_metadata m ON a.id = m.asset_id - WHERE - """, - roots_clause, - "AND a.mtime >= ? AND a.mtime < ?", - ] - params: List[Any] = [] - params.extend(roots_params) - params.extend([start_i, end_i]) - - filter_clauses, filter_params = _build_filter_clauses(safe_filters) - sql_parts.extend(filter_clauses) - params.extend(filter_params) - - sql_parts.append("GROUP BY day ORDER BY day ASC") - - sql = " ".join(sql_parts) - result = await self.db.aquery(sql, tuple(params)) + start_i, end_i = month_range + roots_clause, roots_params = _build_roots_where_clause(cleaned_roots, alias="a") + sql, params = _build_histogram_query(roots_clause, roots_params, start_i, end_i, filters) + result = await self.db.aquery(sql, params) if not result.ok: return Result.Err("DB_ERROR", result.error or "Histogram query failed") - - days: Dict[str, int] = {} - for row in result.data or []: - try: - day = str(row.get("day") or "").strip() - count = int(row.get("count") or 0) - except Exception: - continue - if day: - days[day] = max(0, count) - - return Result.Ok(days) + return Result.Ok(_coerce_histogram_days(result.data)) async def get_asset(self, asset_id: int) -> Result[Optional[Dict[str, Any]]]: """ @@ -816,22 +1083,7 @@ class IndexSearcher: Intended for UI batching (viewer preloading) without triggering per-asset metadata extraction. """ - cleaned: List[int] = [] - seen = set() - for raw in asset_ids or []: - try: - n = int(raw) - except Exception: - continue - if n <= 0: - continue - if n in seen: - continue - seen.add(n) - cleaned.append(n) - if len(cleaned) >= MAX_ASSET_BATCH_IDS: - break - + cleaned = _normalize_asset_ids(asset_ids) if not cleaned: return Result.Ok([]) @@ -858,28 +1110,8 @@ class IndexSearcher: if not result.ok: return Result.Err("QUERY_FAILED", result.error or "Batch asset lookup failed") - by_id: Dict[int, Dict[str, Any]] = {} - for row in result.data or []: - try: - asset = self._hydrate_asset_payload(dict(row)) - except Exception: - continue - try: - raw_id = asset.get("id") - if raw_id is None: - continue - aid = int(raw_id) - except Exception: - continue - by_id[aid] = asset - - # Preserve requested ordering. - out: List[Dict[str, Any]] = [] - for aid in cleaned: - item = by_id.get(aid) - if item: - out.append(item) - return Result.Ok(out) + by_id = _map_assets_by_id(result.data, self._hydrate_asset_payload) + return Result.Ok(_assets_in_requested_order(cleaned, by_id)) def _hydrate_asset_payload(self, asset: Dict[str, Any]) -> Dict[str, Any]: # Parse tags JSON (stored as string in DB) @@ -916,12 +1148,9 @@ class IndexSearcher: This is used to enrich filesystem listings (input/custom) without requiring a full directory scan on every request. """ - cleaned = [str(p) for p in (filepaths or []) if p] + cleaned = _normalize_lookup_filepaths(filepaths) if not cleaned: return Result.Ok({}) - # Guard against overly large IN clauses - if len(cleaned) > MAX_FILEPATH_LOOKUP: - cleaned = cleaned[:MAX_FILEPATH_LOOKUP] result = await self.db.aquery_in( """ @@ -943,25 +1172,7 @@ class IndexSearcher: ) if not result.ok: return Result.Err("DB_ERROR", result.error or "Filepath lookup failed") - - out: Dict[str, Dict[str, Any]] = {} - for row in result.data or []: - if not isinstance(row, dict): - continue - fp = row.get("filepath") - if not fp: - continue - item = dict(row) - tags_raw = item.get("tags") - if tags_raw: - try: - item["tags"] = json.loads(tags_raw) if isinstance(tags_raw, str) else tags_raw - except Exception: - item["tags"] = [] - else: - item["tags"] = [] - out[str(fp)] = item - return Result.Ok(out) + return Result.Ok(_map_lookup_rows(result.data)) def _sanitize_fts_query(self, query: str) -> str: """ @@ -1004,33 +1215,50 @@ class IndexSearcher: # Allow "browse all" queries. if trimmed == "*" or (trimmed and all(token == "*" for token in trimmed.split())): return None - if len(trimmed) > MAX_SEARCH_QUERY_LENGTH: - return Result.Err( - "QUERY_TOO_LONG", - f"Search queries must be at most {MAX_SEARCH_QUERY_LENGTH} characters" - ) - + length_error = self._validate_query_length(trimmed) + if length_error: + return length_error tokens = trimmed.split() + token_error = self._validate_query_tokens(tokens) + if token_error: + return token_error + wildcard_error = self._validate_wildcard_mix(tokens) + if wildcard_error: + return wildcard_error + + return None + + @staticmethod + def _validate_query_length(trimmed: str) -> Optional[Result[Any]]: + if len(trimmed) <= MAX_SEARCH_QUERY_LENGTH: + return None + return Result.Err( + "QUERY_TOO_LONG", + f"Search queries must be at most {MAX_SEARCH_QUERY_LENGTH} characters" + ) + + @staticmethod + def _validate_query_tokens(tokens: List[str]) -> Optional[Result[Any]]: if len(tokens) > MAX_SEARCH_TOKENS: return Result.Err( "QUERY_TOO_COMPLEX", f"Use at most {MAX_SEARCH_TOKENS} tokens for search queries" ) - if any(len(token) > MAX_TOKEN_LENGTH for token in tokens): return Result.Err( "TOKEN_TOO_LONG", f"Each search token must be under {MAX_TOKEN_LENGTH} characters" ) + return None - # Reject suspicious patterns (repeated wildcard tokens) + @staticmethod + def _validate_wildcard_mix(tokens: List[str]) -> Optional[Result[Any]]: wildcard_hits = sum(1 for token in tokens if token == "*") if wildcard_hits > 0 and wildcard_hits >= len(tokens) - 1: return Result.Err( "QUERY_TOO_GENERAL", "Search query must contain at least one non-wildcard term" ) - return None def _build_tags_text_clause(self) -> str: diff --git a/mjr_am_backend/features/index/service.py b/mjr_am_backend/features/index/service.py index 59c8939a..78f0eb35 100644 --- a/mjr_am_backend/features/index/service.py +++ b/mjr_am_backend/features/index/service.py @@ -1,4 +1,4 @@ -""" +""" Index Service - file scanning, incremental updates, and FTS search. This service coordinates multiple specialized components: @@ -26,6 +26,52 @@ from .metadata_helpers import MetadataHelpers logger = get_logger(__name__) +def _normalize_rename_paths(old_filepath: str, new_filepath: str) -> tuple[str, str]: + return str(old_filepath or ""), str(new_filepath or "") + + +def _extract_rename_target_info(new_fp: str) -> tuple[str, str, Optional[int]]: + new_path = Path(new_fp) + filename = new_path.name + subfolder = str(new_path.parent) + mtime: Optional[int] = None + try: + if new_path.exists(): + mtime = int(new_path.stat().st_mtime) + except Exception: + mtime = None + return filename, subfolder, mtime + + +async def _update_assets_filepath_row(db, old_fp: str, new_fp: str, filename: str, subfolder: str, mtime: Optional[int]) -> Result[Any]: + if mtime is None: + return await db.aexecute( + "UPDATE assets SET filepath = ?, filename = ?, subfolder = ?, updated_at = CURRENT_TIMESTAMP WHERE filepath = ?", + (new_fp, filename, subfolder, old_fp), + ) + return await db.aexecute( + "UPDATE assets SET filepath = ?, filename = ?, subfolder = ?, mtime = ?, updated_at = CURRENT_TIMESTAMP WHERE filepath = ?", + (new_fp, filename, subfolder, mtime, old_fp), + ) + + +async def _update_path_keyed_index_tables(db, old_fp: str, new_fp: str, subfolder: str, mtime: Optional[int]) -> Result[bool]: + sj = await db.aexecute( + "UPDATE scan_journal SET filepath = ?, dir_path = ?, mtime = COALESCE(?, mtime), last_seen = CURRENT_TIMESTAMP WHERE filepath = ?", + (new_fp, subfolder, mtime, old_fp), + ) + if not sj.ok: + return Result.Err("DB_ERROR", sj.error or "Failed to update scan_journal filepath") + + mc = await db.aexecute( + "UPDATE metadata_cache SET filepath = ?, last_updated = CURRENT_TIMESTAMP WHERE filepath = ?", + (new_fp, old_fp), + ) + if not mc.ok: + return Result.Err("DB_ERROR", mc.error or "Failed to update metadata_cache filepath") + return Result.Ok(True) + + class IndexService: """ Handles file indexing and search operations. @@ -176,55 +222,27 @@ class IndexService: This keeps DB metadata attached to the asset while updating path-keyed tables. """ - old_fp = str(old_filepath or "") - new_fp = str(new_filepath or "") + old_fp, new_fp = _normalize_rename_paths(old_filepath, new_filepath) if not old_fp or not new_fp: return Result.Err("INVALID_INPUT", "Missing old/new filepath") if old_fp == new_fp: return Result.Ok(True) try: - new_path = Path(new_fp) - filename = new_path.name - subfolder = str(new_path.parent) - mtime = None - try: - if new_path.exists(): - mtime = int(new_path.stat().st_mtime) - except Exception: - mtime = None + filename, subfolder, mtime = _extract_rename_target_info(new_fp) async with self.db.atransaction(mode="immediate"): defer_fk = await self.db.aexecute("PRAGMA defer_foreign_keys = ON") if not defer_fk.ok: return Result.Err("DB_ERROR", defer_fk.error or "Failed to defer foreign key checks") - if mtime is None: - upd = await self.db.aexecute( - "UPDATE assets SET filepath = ?, filename = ?, subfolder = ?, updated_at = CURRENT_TIMESTAMP WHERE filepath = ?", - (new_fp, filename, subfolder, old_fp), - ) - else: - upd = await self.db.aexecute( - "UPDATE assets SET filepath = ?, filename = ?, subfolder = ?, mtime = ?, updated_at = CURRENT_TIMESTAMP WHERE filepath = ?", - (new_fp, filename, subfolder, mtime, old_fp), - ) + upd = await _update_assets_filepath_row(self.db, old_fp, new_fp, filename, subfolder, mtime) if not upd.ok: return Result.Err("DB_ERROR", upd.error or "Failed to update asset filepath") - sj = await self.db.aexecute( - "UPDATE scan_journal SET filepath = ?, dir_path = ?, mtime = COALESCE(?, mtime), last_seen = CURRENT_TIMESTAMP WHERE filepath = ?", - (new_fp, subfolder, mtime, old_fp), - ) - if not sj.ok: - return Result.Err("DB_ERROR", sj.error or "Failed to update scan_journal filepath") - - mc = await self.db.aexecute( - "UPDATE metadata_cache SET filepath = ?, last_updated = CURRENT_TIMESTAMP WHERE filepath = ?", - (new_fp, old_fp), - ) - if not mc.ok: - return Result.Err("DB_ERROR", mc.error or "Failed to update metadata_cache filepath") + side_updates = await _update_path_keyed_index_tables(self.db, old_fp, new_fp, subfolder, mtime) + if not side_updates.ok: + return side_updates except Exception as exc: return Result.Err("DB_ERROR", str(exc)) @@ -388,4 +406,3 @@ class IndexService: return { "enrichment_queue_length": queue_len, } - diff --git a/mjr_am_backend/features/index/updater.py b/mjr_am_backend/features/index/updater.py index 7e38e0a5..0141bf9a 100644 --- a/mjr_am_backend/features/index/updater.py +++ b/mjr_am_backend/features/index/updater.py @@ -88,26 +88,37 @@ class AssetUpdater: if not check_result.ok or not check_result.data or len(check_result.data) == 0: return Result.Err("NOT_FOUND", f"Asset not found: {asset_id}") - # Sanitize tags - sanitized = [] + sanitized = self._sanitize_tags(tags) + tags_json, tags_text = self._build_tags_payload(sanitized) + result = await self._write_asset_tags(asset_id, tags_json, tags_text) + + if not result.ok: + return Result.Err("UPDATE_FAILED", result.error or "Failed to update tags") + + return Result.Ok({"asset_id": asset_id, "tags": sanitized}) + + def _sanitize_tags(self, tags: List[str]) -> List[str]: + sanitized: List[str] = [] for tag in tags: if not isinstance(tag, str): continue - tag = str(tag).strip() - if not tag or len(tag) > MAX_TAG_LENGTH: + cleaned = str(tag).strip() + if not cleaned or len(cleaned) > MAX_TAG_LENGTH: continue - if tag not in sanitized: # Deduplicate - sanitized.append(tag) + if cleaned in sanitized: + continue + sanitized.append(cleaned) + return sanitized + def _build_tags_payload(self, sanitized: List[str]) -> tuple[str, str | None]: tags_json = json.dumps(sanitized, ensure_ascii=False) - - # Build tags_text for FTS5 (if column exists) tags_text = " ".join(sanitized) if self._has_tags_text_column else None + return tags_json, tags_text - # Update or insert asset_metadata (serialize per-asset to avoid race with enrichers) + async def _write_asset_tags(self, asset_id: int, tags_json: str, tags_text: str | None) -> Result[Any]: async with self.db.lock_for_asset(asset_id): if self._has_tags_text_column: - result = await self.db.aexecute( + return await self.db.aexecute( """ INSERT INTO asset_metadata (asset_id, rating, tags, tags_text) SELECT ?, 0, ?, ? @@ -119,23 +130,17 @@ class AssetUpdater: """, (asset_id, tags_json, tags_text, asset_id) ) - else: - result = await self.db.aexecute( - """ - INSERT INTO asset_metadata (asset_id, rating, tags) - SELECT ?, 0, ? - WHERE EXISTS (SELECT 1 FROM assets WHERE id = ?) - ON CONFLICT(asset_id) DO UPDATE SET - tags = excluded.tags - WHERE EXISTS (SELECT 1 FROM assets WHERE id = excluded.asset_id) - """, - (asset_id, tags_json, asset_id) - ) - - if not result.ok: - return Result.Err("UPDATE_FAILED", result.error or "Failed to update tags") - - return Result.Ok({"asset_id": asset_id, "tags": sanitized}) + return await self.db.aexecute( + """ + INSERT INTO asset_metadata (asset_id, rating, tags) + SELECT ?, 0, ? + WHERE EXISTS (SELECT 1 FROM assets WHERE id = ?) + ON CONFLICT(asset_id) DO UPDATE SET + tags = excluded.tags + WHERE EXISTS (SELECT 1 FROM assets WHERE id = excluded.asset_id) + """, + (asset_id, tags_json, asset_id) + ) async def get_all_tags(self) -> Result[List[str]]: """ @@ -155,20 +160,29 @@ class AssetUpdater: if not result.ok: return Result.Err("DB_ERROR", result.error or "Failed to read tags") - # Collect all unique tags - all_tags = set() - for row in result.data or []: - tags_json = row.get("tags") - if not tags_json: - continue - try: - tags = json.loads(tags_json) if isinstance(tags_json, str) else tags_json - if isinstance(tags, list): - for tag in tags: - if isinstance(tag, str) and tag.strip(): - all_tags.add(tag.strip()) - except Exception: - continue + all_tags = self._collect_unique_tags(result.data or []) + return Result.Ok(sorted(all_tags)) - sorted_tags = sorted(all_tags) - return Result.Ok(sorted_tags) + def _collect_unique_tags(self, rows: List[Dict[str, Any]]) -> set[str]: + all_tags: set[str] = set() + for row in rows: + for tag in self._row_tags(row): + all_tags.add(tag) + return all_tags + + @staticmethod + def _row_tags(row: Dict[str, Any]) -> List[str]: + tags_json = row.get("tags") + if not tags_json: + return [] + try: + tags = json.loads(tags_json) if isinstance(tags_json, str) else tags_json + except Exception: + return [] + if not isinstance(tags, list): + return [] + out: List[str] = [] + for tag in tags: + if isinstance(tag, str) and tag.strip(): + out.append(tag.strip()) + return out diff --git a/mjr_am_backend/features/index/watcher.py b/mjr_am_backend/features/index/watcher.py index 38f226c7..91b946d4 100644 --- a/mjr_am_backend/features/index/watcher.py +++ b/mjr_am_backend/features/index/watcher.py @@ -11,7 +11,7 @@ import time from collections import deque from pathlib import Path from threading import Lock -from typing import Deque, Dict, Set, Tuple, Optional, Callable, Awaitable, Coroutine, Any, cast +from typing import Deque, Dict, Set, Tuple, Optional, Callable, Awaitable, Coroutine, Any, cast, List from watchdog.observers import Observer from watchdog.events import ( @@ -312,125 +312,150 @@ class DebouncedWatchHandler(FileSystemEventHandler): src_ok = self._is_supported(src_path) dst_ok = self._is_supported(dest_path) - if not src_ok and not dst_ok: - return - src_norm = self._normalize_path(src_path) dst_norm = self._normalize_path(dest_path) - - # If extension changed from supported->unsupported, treat as deletion. - if src_ok and not dst_ok and self._on_files_removed: - try: - coro = cast(Coroutine[Any, Any, None], self._on_files_removed([src_norm])) - asyncio.run_coroutine_threadsafe(coro, self._loop) - except Exception: - pass + mode = self._moved_file_mode(src_ok=src_ok, dst_ok=dst_ok) + if mode == "ignore": return - - # If extension changed from unsupported->supported, treat as creation. - if dst_ok and not src_ok: + if mode == "delete": + self._emit_removed_files([src_norm]) + return + if mode == "create": self._handle_file(dst_norm) return - - if self._on_files_moved: - try: - coro = cast(Coroutine[Any, Any, None], self._on_files_moved([(src_norm, dst_norm)])) - asyncio.run_coroutine_threadsafe(coro, self._loop) - return - except Exception: - pass + if mode == "move" and self._emit_moved_files([(src_norm, dst_norm)]): + return # Fallback: index new path and remove old path as separate operations. self._handle_file(dst_norm) - if self._on_files_removed: - try: - coro = cast(Coroutine[Any, Any, None], self._on_files_removed([src_norm])) - asyncio.run_coroutine_threadsafe(coro, self._loop) - except Exception: - pass + self._emit_removed_files([src_norm]) + + @staticmethod + def _moved_file_mode(*, src_ok: bool, dst_ok: bool) -> str: + if not src_ok and not dst_ok: + return "ignore" + if src_ok and not dst_ok: + return "delete" + if dst_ok and not src_ok: + return "create" + return "move" + + def _emit_removed_files(self, paths: List[str]) -> None: + if not self._on_files_removed: + return + try: + coro = cast(Coroutine[Any, Any, None], self._on_files_removed(paths)) + asyncio.run_coroutine_threadsafe(coro, self._loop) + except Exception: + return + + def _emit_moved_files(self, moves: List[Tuple[str, str]]) -> bool: + if not self._on_files_moved: + return False + try: + coro = cast(Coroutine[Any, Any, None], self._on_files_moved(moves)) + asyncio.run_coroutine_threadsafe(coro, self._loop) + return True + except Exception: + return False async def _flush(self): """Flush pending files to the indexer.""" self._refresh_runtime_settings() - with self._lock: - if not self._pending and not self._overflow: + candidates = self._drain_pending_candidates() + if not candidates: + return + + async with self._flush_semaphore: + _record_flush_volume(len(candidates)) + files = self._filter_flush_candidates(candidates) + if not files: return - now = time.time() + files = self._apply_flush_backpressure(files) + self._mark_recent_files(files) + try: + await self._on_files_ready(files) + except Exception as e: + logger.debug("Watcher flush error: %s", e) + + def _drain_pending_candidates(self) -> List[str]: + with self._lock: + if not self._pending and not self._overflow: + return [] candidates = list(self._pending.keys()) if self._overflow: candidates.extend(list(self._overflow.keys())) self._pending.clear() self._overflow.clear() self._flush_timer = None + return candidates - async with self._flush_semaphore: - _record_flush_volume(len(candidates)) + def _filter_flush_candidates(self, candidates: List[str]) -> List[str]: + files: List[str] = [] + for f in candidates: + if _is_recent_generated(f): + continue + if not self._is_supported_flush_file(f): + continue + size = self._flush_file_size_if_eligible(f) + if size is None: + continue + max_size = MAX_FILE_SIZE + if max_size > 0 and size > max_size: + logger.debug("Watcher skipping oversized file %s (%d bytes > %d)", f, size, max_size) + continue + files.append(f) + return files - # Filter out files that are too small (still being written), vanished, - # or recently indexed via the frontend executed-event path. - # TOCTOU: re-validate extension at flush time in case file was renamed. - files = [] - for f in candidates: - if _is_recent_generated(f): - continue - try: - ext = Path(f).suffix.lower() - if ext in EXCLUDED_EXTENSIONS or ext not in SUPPORTED_EXTENSIONS: - continue - except Exception: - continue - try: - size = os.path.getsize(f) - except OSError: - continue - if size < MIN_FILE_SIZE: - continue - max_size = MAX_FILE_SIZE - if max_size > 0 and size > max_size: - logger.debug("Watcher skipping oversized file %s (%d bytes > %d)", f, size, max_size) - continue - files.append(f) + @staticmethod + def _is_supported_flush_file(filepath: str) -> bool: + try: + ext = Path(filepath).suffix.lower() + return ext not in EXCLUDED_EXTENSIONS and ext in SUPPORTED_EXTENSIONS + except Exception: + return False - if not files: - return + @staticmethod + def _flush_file_size_if_eligible(filepath: str) -> Optional[int]: + try: + size = os.path.getsize(filepath) + except OSError: + return None + if size < MIN_FILE_SIZE: + return None + return size - # Backpressure handling: process a bounded chunk now and defer remainder. - if WATCHER_FLUSH_MAX_FILES > 0 and len(files) > WATCHER_FLUSH_MAX_FILES: - now = time.time() - deferred = files[WATCHER_FLUSH_MAX_FILES:] - files = files[:WATCHER_FLUSH_MAX_FILES] - with self._lock: - for f in deferred: - if f not in self._pending and f not in self._overflow: - self._overflow[f] = now - try: - logger.warning( - "Watcher flush capped to %d files (requested %d); deferring %d for next flush.", - WATCHER_FLUSH_MAX_FILES, - len(files) + len(deferred), - len(deferred), - ) - except Exception: - pass - # Keep draining deferred backlog without waiting for fresh fs events. - self._schedule_flush() + def _apply_flush_backpressure(self, files: List[str]) -> List[str]: + if WATCHER_FLUSH_MAX_FILES <= 0 or len(files) <= WATCHER_FLUSH_MAX_FILES: + return files + now = time.time() + deferred = files[WATCHER_FLUSH_MAX_FILES:] + files = files[:WATCHER_FLUSH_MAX_FILES] + with self._lock: + for f in deferred: + if f not in self._pending and f not in self._overflow: + self._overflow[f] = now + try: + logger.warning( + "Watcher flush capped to %d files (requested %d); deferring %d for next flush.", + WATCHER_FLUSH_MAX_FILES, + len(files) + len(deferred), + len(deferred), + ) + except Exception: + pass + self._schedule_flush() + return files - # Mark as recently processed - with self._lock: - now = time.time() - for f in files: - self._recent[f] = now - - # Prune old entries from recent - to_remove = [k for k, v in self._recent.items() if (now - v) > self._dedupe_ttl_s * 2] - for k in to_remove: - self._recent.pop(k, None) - - try: - await self._on_files_ready(files) - except Exception as e: - logger.debug("Watcher flush error: %s", e) + def _mark_recent_files(self, files: List[str]) -> None: + with self._lock: + now = time.time() + for f in files: + self._recent[f] = now + to_remove = [k for k, v in self._recent.items() if (now - v) > self._dedupe_ttl_s * 2] + for k in to_remove: + self._recent.pop(k, None) def _normalize_recent_key(path: str) -> str: @@ -444,19 +469,28 @@ def _is_recent_generated(path: str) -> bool: key = _normalize_recent_key(path) if not key: return False - now = time.time() try: with _RECENT_GENERATED_LOCK: - if _RECENT_GENERATED: - to_remove = [k for k, ts in _RECENT_GENERATED.items() if (now - float(ts or 0)) > _RECENT_GENERATED_TTL_S] - for k in to_remove: - _RECENT_GENERATED.pop(k, None) - ts = _RECENT_GENERATED.get(key) - if ts and (now - float(ts or 0)) <= _RECENT_GENERATED_TTL_S: - return True + now = time.time() + _prune_recent_generated(now) + return _is_recent_generated_key_locked(key, now) except Exception: return False - return False + + +def _prune_recent_generated(now: float) -> None: + if not _RECENT_GENERATED: + return + to_remove = [k for k, ts in _RECENT_GENERATED.items() if (now - float(ts or 0)) > _RECENT_GENERATED_TTL_S] + for key in to_remove: + _RECENT_GENERATED.pop(key, None) + + +def _is_recent_generated_key_locked(key: str, now: float) -> bool: + ts = _RECENT_GENERATED.get(key) + if not ts: + return False + return (now - float(ts or 0)) <= _RECENT_GENERATED_TTL_S def is_recent_generated(path: str) -> bool: @@ -503,28 +537,39 @@ def _record_flush_volume(count: int) -> None: return now = time.time() window = WATCHER_STREAM_ALERT_WINDOW_SECONDS - cutoff_time = now - window with _STREAM_LOCK: _STREAM_EVENTS.append((now, count)) _STREAM_TOTAL_FILES += count - while _STREAM_EVENTS and _STREAM_EVENTS[0][0] < cutoff_time: - _, stale = _STREAM_EVENTS.popleft() - _STREAM_TOTAL_FILES -= stale - if _STREAM_TOTAL_FILES < 0: - _STREAM_TOTAL_FILES = 0 - cooldown = max(WATCHER_STREAM_ALERT_COOLDOWN_SECONDS, 0.0) - if _STREAM_TOTAL_FILES >= WATCHER_STREAM_ALERT_THRESHOLD: - if cooldown <= 0 or (now - _LAST_STREAM_ALERT_TIME) >= cooldown: - _LAST_STREAM_ALERT_TIME = now - try: - logger.warning( - "Watcher ingest rate high: %d files in the last %.1f seconds (threshold=%d).", - _STREAM_TOTAL_FILES, - window, - WATCHER_STREAM_ALERT_THRESHOLD, - ) - except Exception: - pass + _prune_stream_events(now, window) + if not _should_emit_stream_alert(now): + return + _LAST_STREAM_ALERT_TIME = now + try: + logger.warning( + "Watcher ingest rate high: %d files in the last %.1f seconds (threshold=%d).", + _STREAM_TOTAL_FILES, + window, + WATCHER_STREAM_ALERT_THRESHOLD, + ) + except Exception: + pass + + +def _prune_stream_events(now: float, window: float) -> None: + global _STREAM_TOTAL_FILES + cutoff_time = now - window + while _STREAM_EVENTS and _STREAM_EVENTS[0][0] < cutoff_time: + _, stale = _STREAM_EVENTS.popleft() + _STREAM_TOTAL_FILES -= stale + if _STREAM_TOTAL_FILES < 0: + _STREAM_TOTAL_FILES = 0 + + +def _should_emit_stream_alert(now: float) -> bool: + if _STREAM_TOTAL_FILES < WATCHER_STREAM_ALERT_THRESHOLD: + return False + cooldown = max(WATCHER_STREAM_ALERT_COOLDOWN_SECONDS, 0.0) + return cooldown <= 0 or (now - _LAST_STREAM_ALERT_TIME) >= cooldown class OutputWatcher: @@ -565,144 +610,9 @@ class OutputWatcher: loop = loop or asyncio.get_event_loop() self._allowed_sources = set() - - async def on_files_ready(files: list): - if not files: - return - # Group by watched root for proper base_dir - by_dir: Dict[str, dict] = {} - for f in files: - try: - # Find the watched root this file belongs to - normalized_f = os.path.normcase(os.path.normpath(f)) - best = None - best_len = -1 - for entry in self._watched_paths.values(): - watched = entry.get("path") - if not watched: - continue - normalized_watched = os.path.normcase(os.path.normpath(watched)) - if _is_under_path(normalized_f, normalized_watched) and len(normalized_watched) > best_len: - best = entry - best_len = len(normalized_watched) - - if not best: - continue - - base = best.get("path") - if not base: - continue - key = str(base) - if key not in by_dir: - by_dir[key] = { - "files": [], - "source": best.get("source"), - "root_id": best.get("root_id"), - } - by_dir[key]["files"].append(f) - except Exception: - pass - - for base_dir, payload in by_dir.items(): - try: - await self._index_callback( - payload.get("files") or [], - base_dir, - payload.get("source"), - payload.get("root_id"), - ) - except Exception as e: - logger.debug("Watcher index error: %s", e) - - async def on_files_removed(files: list): - if not files or not callable(self._remove_callback): - return - by_dir: Dict[str, dict] = {} - for f in files: - try: - normalized_f = os.path.normcase(os.path.normpath(f)) - best = None - best_len = -1 - for entry in self._watched_paths.values(): - watched = entry.get("path") - if not watched: - continue - normalized_watched = os.path.normcase(os.path.normpath(watched)) - if _is_under_path(normalized_f, normalized_watched) and len(normalized_watched) > best_len: - best = entry - best_len = len(normalized_watched) - if not best: - continue - base = best.get("path") - if not base: - continue - key = str(base) - if key not in by_dir: - by_dir[key] = {"files": [], "source": best.get("source"), "root_id": best.get("root_id")} - by_dir[key]["files"].append(f) - except Exception: - continue - for base_dir, payload in by_dir.items(): - try: - await self._remove_callback( - payload.get("files") or [], - base_dir, - payload.get("source"), - payload.get("root_id"), - ) - except Exception as e: - logger.debug("Watcher remove error: %s", e) - - async def on_files_moved(moves: list): - if not moves: - return - if callable(self._move_callback): - by_dir: Dict[str, dict] = {} - for src, dst in moves: - try: - normalized_dst = os.path.normcase(os.path.normpath(dst)) - best = None - best_len = -1 - for entry in self._watched_paths.values(): - watched = entry.get("path") - if not watched: - continue - normalized_watched = os.path.normcase(os.path.normpath(watched)) - if _is_under_path(normalized_dst, normalized_watched) and len(normalized_watched) > best_len: - best = entry - best_len = len(normalized_watched) - if not best: - continue - base = best.get("path") - if not base: - continue - key = str(base) - if key not in by_dir: - by_dir[key] = {"moves": [], "source": best.get("source"), "root_id": best.get("root_id")} - by_dir[key]["moves"].append((src, dst)) - except Exception: - continue - for base_dir, payload in by_dir.items(): - try: - await self._move_callback( - payload.get("moves") or [], - base_dir, - payload.get("source"), - payload.get("root_id"), - ) - except Exception as e: - logger.debug("Watcher move error: %s", e) - return - # Fallback if no move callback: reindex destination and remove source. - try: - src_files = [m[0] for m in moves if isinstance(m, (list, tuple)) and len(m) == 2] - dst_files = [m[1] for m in moves if isinstance(m, (list, tuple)) and len(m) == 2] - if dst_files: - await on_files_ready(dst_files) - if src_files: - await on_files_removed(src_files) - except Exception: - return + on_files_ready = self._handle_ready_files + on_files_removed = self._handle_removed_files + on_files_moved = self._handle_moved_files self._handler = DebouncedWatchHandler( on_files_ready, @@ -717,30 +627,8 @@ class OutputWatcher: for path in paths: try: - if isinstance(path, dict): - raw_path = path.get("path") - source = path.get("source") - root_id = path.get("root_id") - else: - raw_path = path - source = None - root_id = None - - if not raw_path: - continue - normalized = os.path.normpath(raw_path) - if os.path.isdir(normalized): - watch = self._observer.schedule(self._handler, normalized, recursive=True) - source_norm = str(source or "").strip().lower() if source is not None else None - if source_norm: - self._allowed_sources.add(source_norm) - self._watched_paths[str(id(watch))] = { - "path": normalized, - "source": source_norm, - "root_id": root_id, - "watch": watch, - } - logger.info("Watcher started for: %s", normalized) + raw_path, source, root_id = self._resolve_watch_path(path) + self._register_watch_path(raw_path, source=source, root_id=root_id, log_label="started") except Exception as e: logger.warning("Failed to watch %s: %s", path, e) @@ -749,6 +637,83 @@ class OutputWatcher: self._running = True logger.info("File watcher started for %d directories", len(self._watched_paths)) + async def _handle_ready_files(self, files: list) -> None: + if not files: + return + grouped = self._group_files_by_watched_root(files) + await self._dispatch_file_groups(grouped, self._index_callback, "index") + + async def _handle_removed_files(self, files: list) -> None: + if not files or not callable(self._remove_callback): + return + grouped = self._group_files_by_watched_root(files) + await self._dispatch_file_groups(grouped, self._remove_callback, "remove") + + async def _handle_moved_files(self, moves: list) -> None: + if not moves: + return + if callable(self._move_callback): + await self._dispatch_move_groups(self._group_moves_by_watched_root(moves)) + return + await self._handle_move_fallback(moves) + + async def _dispatch_move_groups(self, grouped_moves: Dict[str, dict]) -> None: + if not callable(self._move_callback): + return + for base_dir, payload in grouped_moves.items(): + try: + await self._move_callback( + payload.get("moves") or [], + base_dir, + payload.get("source"), + payload.get("root_id"), + ) + except Exception as exc: + logger.debug("Watcher move error: %s", exc) + + async def _handle_move_fallback(self, moves: list) -> None: + # Fallback if no move callback: reindex destination and remove source. + try: + src_files, dst_files = self._split_move_pairs(moves) + if dst_files: + await self._handle_ready_files(dst_files) + if src_files: + await self._handle_removed_files(src_files) + except Exception: + return + + @staticmethod + def _resolve_watch_path(path: Any) -> tuple[Any, Optional[str], Optional[str]]: + if isinstance(path, dict): + return path.get("path"), path.get("source"), path.get("root_id") + return path, None, None + + def _register_watch_path( + self, + raw_path: Any, + *, + source: Optional[str], + root_id: Optional[str], + log_label: str, + ) -> bool: + if not self._observer or not self._handler or not raw_path: + return False + normalized = os.path.normpath(raw_path) + if not os.path.isdir(normalized): + return False + watch = self._observer.schedule(self._handler, normalized, recursive=True) + source_norm = self._normalize_source(source) + if source_norm: + self._allowed_sources.add(source_norm) + self._watched_paths[str(id(watch))] = { + "path": normalized, + "source": source_norm, + "root_id": root_id, + "watch": watch, + } + logger.info("Watcher %s for: %s", log_label, normalized) + return True + async def stop(self): """Stop watching all directories.""" if not self._running: @@ -774,29 +739,116 @@ class OutputWatcher: try: normalized = os.path.normpath(path) - if os.path.isdir(normalized): - # Check if already watched - for entry in self._watched_paths.values(): - if entry.get("path") == normalized: - return - source_norm = str(source or "").strip().lower() if source is not None else None - if not self._allows_source(source_norm): - try: - logger.debug("Watcher ignored path (scope mismatch): %s (source=%s)", normalized, source_norm) - except Exception: - pass - return - watch = self._observer.schedule(self._handler, normalized, recursive=True) - self._watched_paths[str(id(watch))] = { - "path": normalized, - "source": source_norm, - "root_id": root_id, - "watch": watch, - } - logger.info("Watcher added: %s", normalized) + if not os.path.isdir(normalized): + return + if self._is_already_watched(normalized): + return + source_norm = self._normalize_source(source) + if not self._allows_source(source_norm): + try: + logger.debug("Watcher ignored path (scope mismatch): %s (source=%s)", normalized, source_norm) + except Exception: + pass + return + watch = self._observer.schedule(self._handler, normalized, recursive=True) + self._watched_paths[str(id(watch))] = { + "path": normalized, + "source": source_norm, + "root_id": root_id, + "watch": watch, + } + logger.info("Watcher added: %s", normalized) except Exception as e: logger.warning("Failed to add watch for %s: %s", path, e) + def _normalize_source(self, source: Optional[str]) -> Optional[str]: + return str(source or "").strip().lower() if source is not None else None + + def _is_already_watched(self, normalized: str) -> bool: + for entry in self._watched_paths.values(): + if entry.get("path") == normalized: + return True + return False + + def _best_watched_entry_for_path(self, path_value: str) -> Optional[dict]: + normalized_path = os.path.normcase(os.path.normpath(path_value)) + best = None + best_len = -1 + for entry in self._watched_paths.values(): + watched = entry.get("path") + if not watched: + continue + normalized_watched = os.path.normcase(os.path.normpath(watched)) + if _is_under_path(normalized_path, normalized_watched) and len(normalized_watched) > best_len: + best = entry + best_len = len(normalized_watched) + return best + + def _group_files_by_watched_root(self, files: List[str]) -> Dict[str, dict]: + by_dir: Dict[str, dict] = {} + for file_path in files: + try: + best = self._best_watched_entry_for_path(file_path) + if not best: + continue + base = best.get("path") + if not base: + continue + key = str(base) + if key not in by_dir: + by_dir[key] = {"files": [], "source": best.get("source"), "root_id": best.get("root_id")} + by_dir[key]["files"].append(file_path) + except Exception: + continue + return by_dir + + def _group_moves_by_watched_root(self, moves: List[Any]) -> Dict[str, dict]: + by_dir: Dict[str, dict] = {} + for move in moves: + try: + src, dst = move + except Exception: + continue + best = self._best_watched_entry_for_path(str(dst)) + if not best: + continue + base = best.get("path") + if not base: + continue + key = str(base) + if key not in by_dir: + by_dir[key] = {"moves": [], "source": best.get("source"), "root_id": best.get("root_id")} + by_dir[key]["moves"].append((src, dst)) + return by_dir + + async def _dispatch_file_groups( + self, + grouped: Dict[str, dict], + callback: Callable[[list, str, Optional[str], Optional[str]], Awaitable[None]], + label: str, + ) -> None: + for base_dir, payload in grouped.items(): + try: + await callback( + payload.get("files") or [], + base_dir, + payload.get("source"), + payload.get("root_id"), + ) + except Exception as exc: + logger.debug("Watcher %s error: %s", label, exc) + + @staticmethod + def _split_move_pairs(moves: List[Any]) -> tuple[List[str], List[str]]: + src_files: List[str] = [] + dst_files: List[str] = [] + for move in moves: + if not isinstance(move, (list, tuple)) or len(move) != 2: + continue + src_files.append(str(move[0])) + dst_files.append(str(move[1])) + return src_files, dst_files + def remove_path(self, path: str): """Remove a path from watching (e.g., when custom root is removed).""" if not self._running or not self._observer: diff --git a/mjr_am_backend/features/index/watcher_scope.py b/mjr_am_backend/features/index/watcher_scope.py index 502efcfe..a73ea486 100644 --- a/mjr_am_backend/features/index/watcher_scope.py +++ b/mjr_am_backend/features/index/watcher_scope.py @@ -24,49 +24,76 @@ def build_watch_paths(scope: str, custom_root_id: Optional[str]) -> list[dict]: s = normalize_scope(scope) if s == "output": - output_root = str(get_runtime_output_root() or "") - if output_root and os.path.isdir(output_root): - watch_paths.append({"path": output_root, "source": "output", "root_id": None}) + output_entry = _output_watch_entry() + if output_entry is not None: + watch_paths.append(output_entry) return watch_paths if s == "custom": - rid = str(custom_root_id or "").strip() - if not rid: - return watch_paths - try: - from ...custom_roots import resolve_custom_root - root_result = resolve_custom_root(rid) - if root_result.ok and root_result.data: - root_path = str(root_result.data) - if root_path and os.path.isdir(root_path): - watch_paths.append({"path": root_path, "source": "custom", "root_id": rid}) - except Exception: - return watch_paths + custom_entry = _custom_watch_entry(custom_root_id) + if custom_entry is not None: + watch_paths.append(custom_entry) return watch_paths return watch_paths +def _output_watch_entry() -> Optional[dict]: + output_root = str(get_runtime_output_root() or "") + if output_root and os.path.isdir(output_root): + return {"path": output_root, "source": "output", "root_id": None} + return None + + +def _custom_watch_entry(custom_root_id: Optional[str]) -> Optional[dict]: + rid = str(custom_root_id or "").strip() + if not rid: + return None + try: + from ...custom_roots import resolve_custom_root + root_result = resolve_custom_root(rid) + except Exception: + return None + if not root_result.ok or not root_result.data: + return None + root_path = str(root_result.data) + if not (root_path and os.path.isdir(root_path)): + return None + return {"path": root_path, "source": "custom", "root_id": rid} + + async def load_watcher_scope(db) -> dict: scope = "output" custom_root_id = "" - try: - res = await db.aquery( - "SELECT key, value FROM metadata WHERE key IN (?, ?)", - (WATCHER_SCOPE_KEY, WATCHER_CUSTOM_ROOT_ID_KEY), - ) - if res.ok and res.data: - for row in res.data: - key = str(row.get("key") or "") - val = str(row.get("value") or "") - if key == WATCHER_SCOPE_KEY and val: - scope = val - elif key == WATCHER_CUSTOM_ROOT_ID_KEY: - custom_root_id = val - except Exception: - pass + res = await _query_watcher_scope_rows(db) + if res: + scope, custom_root_id = _parse_watcher_scope_rows(res, scope, custom_root_id) normalized_scope = normalize_scope(scope) return { "scope": normalized_scope, "custom_root_id": str(custom_root_id or "").strip() if normalized_scope == "custom" else "", } + + +async def _query_watcher_scope_rows(db): + try: + res = await db.aquery( + "SELECT key, value FROM metadata WHERE key IN (?, ?)", + (WATCHER_SCOPE_KEY, WATCHER_CUSTOM_ROOT_ID_KEY), + ) + except Exception: + return None + if not res.ok or not res.data: + return None + return res.data + + +def _parse_watcher_scope_rows(rows, scope: str, custom_root_id: str) -> tuple[str, str]: + for row in rows: + key = str(row.get("key") or "") + val = str(row.get("value") or "") + if key == WATCHER_SCOPE_KEY and val: + scope = val + elif key == WATCHER_CUSTOM_ROOT_ID_KEY: + custom_root_id = val + return scope, custom_root_id diff --git a/mjr_am_backend/features/metadata/extractors.py b/mjr_am_backend/features/metadata/extractors.py index deccbea8..136acda2 100644 --- a/mjr_am_backend/features/metadata/extractors.py +++ b/mjr_am_backend/features/metadata/extractors.py @@ -3,7 +3,7 @@ Metadata extractors for different file types. Extracts ComfyUI workflow and generation parameters from PNG, WEBP, MP4. """ import os -from typing import Optional, Dict, Any, Tuple +from typing import Optional, Dict, Any, Tuple, Set, List from ...shared import Result, ErrorCode, get_logger from .parsing_utils import ( @@ -18,6 +18,749 @@ logger = get_logger(__name__) # Constants removed (moved to parsing_utils) MAX_TAG_LENGTH = 100 +_WEBP_WORKFLOW_KEYS = ("EXIF:Make", "IFD0:Make", "Keys:Workflow", "comfyui:workflow") +_WEBP_PROMPT_KEYS = ("EXIF:Model", "IFD0:Model", "Keys:Prompt", "comfyui:prompt") +_WEBP_TEXT_KEYS = ( + "EXIF:ImageDescription", + "IFD0:ImageDescription", + "ImageDescription", + "EXIF:UserComment", + "IFD0:UserComment", + "UserComment", + "EXIF:Comment", + "IFD0:Comment", + "EXIF:Subject", + "IFD0:Subject", +) +_VIDEO_WORKFLOW_KEYS = ("QuickTime:Workflow", "Keys:Workflow", "comfyui:workflow") +_VIDEO_PROMPT_KEYS = ("QuickTime:Prompt", "Keys:Prompt", "comfyui:prompt") +_RATING_CANDIDATE_KEYS = ( + "xmp:rating", + "xmp-xmp:rating", + "microsoft:ratingpercent", + "xmp-microsoft:ratingpercent", + "microsoft:shareduserrating", + "xmp-microsoft:shareduserrating", + "rating", + "ratingpercent", + "shareduserrating", +) +_TAG_CANDIDATE_KEYS = ( + "dc:subject", + "xmp-dc:subject", + "xmp:subject", + "iptc:keywords", + "photoshop:keywords", + "lr:hierarchicalsubject", + "microsoft:category", + "xmp-microsoft:category", + "xpkeywords", + "keywords", + "subject", + "category", +) +_DIM_WIDTH_KEYS = ( + "Image:ImageWidth", + "EXIF:ImageWidth", + "EXIF:ExifImageWidth", + "IFD0:ImageWidth", + "Composite:ImageWidth", + "QuickTime:ImageWidth", + "PNG:ImageWidth", + "File:ImageWidth", + "width", +) +_DIM_HEIGHT_KEYS = ( + "Image:ImageHeight", + "EXIF:ImageHeight", + "EXIF:ExifImageHeight", + "IFD0:ImageHeight", + "Composite:ImageHeight", + "QuickTime:ImageHeight", + "PNG:ImageHeight", + "File:ImageHeight", + "height", +) +_DIM_PAIR_KEYS = ("Composite:ImageSize", "Image:ImageSize", "QuickTime:ImageSize", "ImageSize") + + +def _inspect_json_field(container: Any, key_names: Tuple[str, ...]) -> Optional[Any]: + if not isinstance(container, dict): + return None + for key in key_names: + parsed = parse_json_value(container.get(key)) + if parsed is not None: + return parsed + return None + + +def _unwrap_workflow_prompt_container(container: Any) -> Tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]]]: + """ + Some pipelines embed: + {"workflow": {...}, "prompt": "{...json...}"} + in a single metadata tag. + """ + if not isinstance(container, dict): + return (None, None) + + wf = _container_json_candidate(container, "workflow") + pr = _container_json_candidate(container, "prompt") + wf_out = wf if isinstance(wf, dict) and looks_like_comfyui_workflow(wf) else None + pr_out = _prompt_graph_from_container_value(pr) + return wf_out, pr_out + + +def _container_json_candidate(container: Dict[str, Any], key: str) -> Any: + return container.get(key) or container.get(key.capitalize()) + + +def _prompt_graph_from_container_value(value: Any) -> Optional[Dict[str, Any]]: + if isinstance(value, dict) and looks_like_comfyui_prompt_graph(value): + return value + if not isinstance(value, str): + return None + parsed = try_parse_json_text(value) + if isinstance(parsed, dict) and looks_like_comfyui_prompt_graph(parsed): + return parsed + return None + + +def _merge_workflow_prompt_candidate( + candidate: Any, + workflow: Optional[Dict[str, Any]], + prompt: Optional[Dict[str, Any]], +) -> Tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]], bool]: + matched = False + wf_candidate: Optional[Dict[str, Any]] = None + pr_candidate: Optional[Dict[str, Any]] = None + + if looks_like_comfyui_workflow(candidate): + wf_candidate = candidate + if looks_like_comfyui_prompt_graph(candidate): + pr_candidate = candidate + if wf_candidate is None and pr_candidate is None: + wf_candidate, pr_candidate = _unwrap_workflow_prompt_container(candidate) + + if workflow is None and wf_candidate is not None: + workflow = wf_candidate + matched = True + if prompt is None and pr_candidate is not None: + prompt = pr_candidate + matched = True + return workflow, prompt, matched + + +def _coerce_dim(value: Any) -> Optional[int]: + try: + if value is None: + return None + if isinstance(value, str): + txt = value.strip().lower().replace("px", "").strip() + if not txt or "x" in txt: + return None + value = txt + out = int(float(value)) + return out if out > 0 else None + except Exception: + return None + + +def _first_available_dim(exif_data: Dict[str, Any], keys: Tuple[str, ...]) -> Optional[int]: + for key in keys: + value = _coerce_dim(exif_data.get(key)) + if value is not None: + return value + return None + + +def _parse_size_pair(raw: Any) -> Tuple[Optional[int], Optional[int]]: + if raw is None: + return (None, None) + try: + txt = str(raw).strip().lower().replace("px", "").replace("×", "x") + txt = txt.replace(",", " ").replace(";", " ") + if "x" in txt: + left, right = [s.strip() for s in txt.split("x", 1)] + else: + parts = [part for part in txt.split() if part] + if len(parts) < 2: + return (None, None) + left, right = parts[0], parts[1] + return (_coerce_dim(left), _coerce_dim(right)) + except Exception: + return (None, None) + + +def _apply_dimensions_from_exif(metadata: Dict[str, Any], exif_data: Dict[str, Any]) -> None: + width = metadata.get("width") or _first_available_dim(exif_data, _DIM_WIDTH_KEYS) + height = metadata.get("height") or _first_available_dim(exif_data, _DIM_HEIGHT_KEYS) + width, height = _fill_missing_dims_from_pairs(width, height, exif_data) + if width is not None and height is not None: + metadata["width"] = int(width) + metadata["height"] = int(height) + + +def _fill_missing_dims_from_pairs( + width: Optional[int], + height: Optional[int], + exif_data: Dict[str, Any], +) -> tuple[Optional[int], Optional[int]]: + if width is not None and height is not None: + return width, height + for key in _DIM_PAIR_KEYS: + pair_w, pair_h = _parse_size_pair(exif_data.get(key)) + if width is None and pair_w is not None: + width = pair_w + if height is None and pair_h is not None: + height = pair_h + if width is not None and height is not None: + break + return width, height + + +def _extract_ffprobe_format_tags(ffprobe_data: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]: + if not isinstance(ffprobe_data, dict): + return None + fmt = ffprobe_data.get("format") or {} + if not isinstance(fmt, dict): + return None + tags = fmt.get("tags") + return tags if isinstance(tags, dict) else None + + +def _extract_ffprobe_stream_tag_dicts(ffprobe_data: Optional[Dict[str, Any]]) -> list[Dict[str, Any]]: + out: list[Dict[str, Any]] = [] + if not isinstance(ffprobe_data, dict): + return out + streams = ffprobe_data.get("streams") + if not isinstance(streams, list): + return out + for stream in streams: + if not isinstance(stream, dict): + continue + tags = stream.get("tags") + if isinstance(tags, dict): + out.append(tags) + return out + + +def _collect_text_candidates(container: Any) -> list[Tuple[str, str]]: + if not isinstance(container, dict): + return [] + out: list[Tuple[str, str]] = [] + for key, value in container.items(): + if isinstance(value, str): + out.append((str(key), value)) + elif isinstance(value, (list, tuple)): + for item in value: + if isinstance(item, str): + out.append((str(key), item)) + return out + + +def _apply_auto1111_text_candidates( + metadata: Dict[str, Any], + candidates: list[Tuple[str, str]], + *, + prompt_graph: Optional[Dict[str, Any]], + preserve_existing_prompt_text: bool, +) -> None: + if metadata.get("parameters") is not None: + return + for _, text in candidates: + parsed = parse_auto1111_params(text) + if not parsed: + continue + metadata["parameters"] = text + metadata.update(parsed) + if metadata.get("quality") != "full": + _bump_quality(metadata, "partial") + if prompt_graph is None and parsed.get("prompt"): + if not preserve_existing_prompt_text or not metadata.get("prompt"): + metadata["prompt"] = parsed.get("prompt") + if parsed.get("negative_prompt"): + metadata["negative_prompt"] = parsed.get("negative_prompt") + return + + +def _workflow_build_link_lookup(links: Any) -> Dict[Any, Tuple[Any, Any, Any, Any]]: + link_lookup: Dict[Any, Tuple[Any, Any, Any, Any]] = {} + if not isinstance(links, list): + return link_lookup + for link in links: + if isinstance(link, list) and len(link) >= 5: + link_lookup[link[0]] = (link[1], link[2], link[3], link[4]) + return link_lookup + + +def _workflow_get_source_data( + node_map: Dict[Any, Dict[str, Any]], + link_lookup: Dict[Any, Tuple[Any, Any, Any, Any]], + target_node_id: Any, + input_name: str, +) -> Optional[Tuple[Optional[Dict[str, Any]], Any]]: + target_node = node_map.get(target_node_id) + if not isinstance(target_node, dict): + return None + + inp_def = next((item for item in target_node.get("inputs", []) if item.get("name") == input_name), None) + if not isinstance(inp_def, dict): + return None + + link_id = inp_def.get("link") + if not link_id: + return None + + link_info = link_lookup.get(link_id) + if not link_info: + return None + + origin_id = link_info[0] + return (node_map.get(origin_id), link_id) + + +def _workflow_get_node_text(node: Dict[str, Any], context: Optional[str] = None) -> Optional[str]: + widgets = node.get("widgets_values") + if not isinstance(widgets, list) or not widgets: + return None + + str_widgets = [value for value in widgets if isinstance(value, str) and value.strip()] + if not str_widgets: + return None + if len(str_widgets) == 1: + return str_widgets[0] + if context == "negative": + return str_widgets[-1] + return str_widgets[0] + + +def _workflow_find_upstream_text( + node_map: Dict[Any, Dict[str, Any]], + link_lookup: Dict[Any, Tuple[Any, Any, Any, Any]], + start_node: Optional[Dict[str, Any]], + *, + context: Optional[str] = None, + depth: int = 0, + seen: Optional[Set[Any]] = None, +) -> List[str]: + if not isinstance(start_node, dict): + return [] + + local_seen: Set[Any] = seen if seen is not None else set() + found_texts: List[str] = [] + stack: List[Tuple[Dict[str, Any], int]] = [(start_node, depth)] + max_depth = 32 + + while stack: + node, d = stack.pop() + if not _workflow_trace_node_allowed(node, d, max_depth, local_seen): + continue + _workflow_collect_node_text(node, context, found_texts) + _workflow_push_upstream_inputs(node_map, link_lookup, node, context, d, stack) + return found_texts + + +def _workflow_trace_node_allowed(node: Any, depth: int, max_depth: int, local_seen: Set[Any]) -> bool: + if not isinstance(node, dict) or depth > max_depth: + return False + node_id = node.get("id") + if node_id is None: + return True + if node_id in local_seen: + return False + local_seen.add(node_id) + return True + + +def _workflow_collect_node_text(node: Dict[str, Any], context: Optional[str], found_texts: List[str]) -> None: + node_type = str(node.get("type", "")).lower() + if "loader" in node_type or "checkpoint" in node_type or "loadimage" in node_type: + return + text = _workflow_get_node_text(node, context) + if text: + found_texts.append(text) + + +def _workflow_push_upstream_inputs( + node_map: Dict[Any, Dict[str, Any]], + link_lookup: Dict[Any, Tuple[Any, Any, Any, Any]], + node: Dict[str, Any], + context: Optional[str], + depth: int, + stack: List[Tuple[Dict[str, Any], int]], +) -> None: + inputs = node.get("inputs", []) + if not isinstance(inputs, list) or not inputs: + return + for inp in reversed(inputs): + if not isinstance(inp, dict): + continue + if not _workflow_input_context_allowed(inp, context): + continue + link_id = inp.get("link") + if not link_id: + continue + link_info = link_lookup.get(link_id) + if not link_info: + continue + origin_node = node_map.get(link_info[0]) + if isinstance(origin_node, dict): + stack.append((origin_node, depth + 1)) + + +def _workflow_input_context_allowed(inp: Dict[str, Any], context: Optional[str]) -> bool: + input_name = str(inp.get("name", "")).lower() + if context == "positive" and "negative" in input_name: + return False + if context == "negative" and "positive" in input_name: + return False + return True + + +def _workflow_apply_sampler_widget_params(params: Dict[str, Any], widgets: Any) -> None: + if not isinstance(widgets, list): + return + for value in widgets: + _workflow_apply_sampler_widget_value(params, value) + + +def _workflow_apply_sampler_widget_value(params: Dict[str, Any], value: Any) -> None: + if isinstance(value, int): + _workflow_apply_sampler_int(params, value) + return + if isinstance(value, float): + if 1.0 <= value <= 30.0 and "cfg" not in params: + params["cfg"] = value + return + if isinstance(value, str): + _workflow_apply_sampler_name(params, value) + + +def _workflow_apply_sampler_int(params: Dict[str, Any], value: int) -> None: + if value > 10000 and "seed" not in params: + params["seed"] = value + elif 1 <= value <= 200 and "steps" not in params: + params["steps"] = value + + +def _workflow_apply_sampler_name(params: Dict[str, Any], value: str) -> None: + if "sampler" in params: + return + if value.lower() in ["euler", "euler_a", "dpm++", "ddim", "uni_pc"]: + params["sampler"] = value + + +def _workflow_trace_sampler_fields( + node: Dict[str, Any], + node_map: Dict[Any, Dict[str, Any]], + link_lookup: Dict[Any, Tuple[Any, Any, Any, Any]], + params: Dict[str, Any], + unique_prompts: Set[str], + unique_negatives: Set[str], + visited_nodes: Set[Any], +) -> None: + node_id = node.get("id") + if node_id is None: + return + _workflow_apply_sampler_widget_params(params, node.get("widgets_values", [])) + + _workflow_collect_source_texts(node_map, link_lookup, node_id, "positive", "positive", unique_prompts, visited_nodes) + _workflow_collect_source_texts(node_map, link_lookup, node_id, "negative", "negative", unique_negatives, visited_nodes) + _workflow_trace_guider_fields(node_map, link_lookup, node_id, unique_prompts, unique_negatives, visited_nodes) + _workflow_trace_model_field(node_map, link_lookup, node_id, params) + + +def _workflow_trace_guider_fields( + node_map: Dict[Any, Dict[str, Any]], + link_lookup: Dict[Any, Tuple[Any, Any, Any, Any]], + node_id: Any, + unique_prompts: Set[str], + unique_negatives: Set[str], + visited_nodes: Set[Any], +) -> None: + guider_info = _workflow_get_source_data(node_map, link_lookup, node_id, "guider") + if not guider_info or not isinstance(guider_info[0], dict): + return + guider_id = guider_info[0].get("id") + cond_info = _workflow_get_source_data(node_map, link_lookup, guider_id, "conditioning") + if not cond_info: + cond_info = _workflow_get_source_data(node_map, link_lookup, guider_id, "positive") + if cond_info: + _workflow_collect_from_source_info(node_map, link_lookup, cond_info, "positive", unique_prompts, visited_nodes) + neg_guider_info = _workflow_get_source_data(node_map, link_lookup, guider_id, "negative") + if neg_guider_info: + _workflow_collect_from_source_info(node_map, link_lookup, neg_guider_info, "negative", unique_negatives, visited_nodes) + + +def _workflow_trace_model_field( + node_map: Dict[Any, Dict[str, Any]], + link_lookup: Dict[Any, Tuple[Any, Any, Any, Any]], + node_id: Any, + params: Dict[str, Any], +) -> None: + model_info = _workflow_get_source_data(node_map, link_lookup, node_id, "model") + if not model_info or not isinstance(model_info[0], dict): + return + for item in model_info[0].get("widgets_values", []): + if isinstance(item, str) and (".safetensors" in item or ".ckpt" in item): + params["model"] = item + return + + +def _workflow_collect_source_texts( + node_map: Dict[Any, Dict[str, Any]], + link_lookup: Dict[Any, Tuple[Any, Any, Any, Any]], + node_id: Any, + source_name: str, + context: str, + out: Set[str], + visited_nodes: Set[Any], +) -> None: + source_info = _workflow_get_source_data(node_map, link_lookup, node_id, source_name) + if source_info: + _workflow_collect_from_source_info(node_map, link_lookup, source_info, context, out, visited_nodes) + + +def _workflow_collect_from_source_info( + node_map: Dict[Any, Dict[str, Any]], + link_lookup: Dict[Any, Tuple[Any, Any, Any, Any]], + source_info: Tuple[Any, Any], + context: str, + out: Set[str], + visited_nodes: Set[Any], +) -> None: + trace_seen: Set[Any] = set() + texts = _workflow_find_upstream_text( + node_map, + link_lookup, + source_info[0], + context=context, + seen=trace_seen, + ) + out.update(texts) + visited_nodes.update(trace_seen) + + +def _workflow_classify_unconnected_node( + start_node: Dict[str, Any], + node_map: Dict[Any, Dict[str, Any]], + link_lookup: Dict[Any, Tuple[Any, Any, Any, Any]], + *, + depth: int = 0, + visited_trace: Optional[Set[Any]] = None, +) -> str: + local_visited = visited_trace if visited_trace is not None else set() + nid = start_node.get("id") + if nid in local_visited or depth > 6: + return "unknown" + local_visited.add(nid) + + outputs = start_node.get("outputs", []) + if not isinstance(outputs, list): + return "unknown" + for link_id in _workflow_iter_output_link_ids(outputs): + tgt_node = _workflow_target_node_from_link(link_id, node_map, link_lookup) + if not isinstance(tgt_node, dict): + continue + direct_role = _workflow_role_from_target_input_link(tgt_node, link_id) + if direct_role != "unknown": + return direct_role + role = _workflow_classify_unconnected_node( + tgt_node, + node_map, + link_lookup, + depth=depth + 1, + visited_trace=local_visited, + ) + if role != "unknown": + return role + + return "unknown" + + +def _workflow_iter_output_link_ids(outputs: List[Any]) -> List[Any]: + link_ids: List[Any] = [] + for output in outputs: + if not isinstance(output, dict): + continue + links = output.get("links") + if not links: + continue + if isinstance(links, list): + link_ids.extend(links) + else: + link_ids.append(links) + return link_ids + + +def _workflow_target_node_from_link( + link_id: Any, + node_map: Dict[Any, Dict[str, Any]], + link_lookup: Dict[Any, Tuple[Any, Any, Any, Any]], +) -> Optional[Dict[str, Any]]: + link_target = link_lookup.get(link_id) + if not link_target: + return None + tgt_node = node_map.get(link_target[2]) + return tgt_node if isinstance(tgt_node, dict) else None + + +def _workflow_role_from_target_input_link(target_node: Dict[str, Any], link_id: Any) -> str: + tgt_inputs = target_node.get("inputs", []) + if not isinstance(tgt_inputs, list): + return "unknown" + for inp in tgt_inputs: + if not isinstance(inp, dict): + continue + if inp.get("link") != link_id: + continue + input_name = str(inp.get("name", "")).lower() + if "positive" in input_name: + return "positive" + if "negative" in input_name: + return "negative" + return "unknown" + + +def _workflow_collect_unconnected_texts( + nodes: List[Dict[str, Any]], + node_map: Dict[Any, Dict[str, Any]], + link_lookup: Dict[Any, Tuple[Any, Any, Any, Any]], + unique_prompts: Set[str], + unique_negatives: Set[str], + visited_nodes: Set[Any], + sampler_found: bool, +) -> None: + for node in nodes: + uid = node.get("id") + if uid in visited_nodes: + continue + + node_type = str(node.get("type", "")).lower() + title = str(node.get("title", "")).lower() + if not _workflow_is_text_prompt_node(node_type): + continue + + text = _workflow_get_node_text(node) + if not text: + continue + + _workflow_store_unconnected_text( + text, + title, + node, + node_map, + link_lookup, + unique_prompts, + unique_negatives, + sampler_found, + ) + + +def _workflow_store_unconnected_text( + text: str, + title: str, + node: Dict[str, Any], + node_map: Dict[Any, Dict[str, Any]], + link_lookup: Dict[Any, Tuple[Any, Any, Any, Any]], + unique_prompts: Set[str], + unique_negatives: Set[str], + sampler_found: bool, +) -> None: + if "negative" in title: + unique_negatives.add(text) + return + if "positive" in title: + unique_prompts.add(text) + return + role = _workflow_classify_unconnected_node(node, node_map, link_lookup) + if role == "positive": + unique_prompts.add(text) + elif role == "negative": + unique_negatives.add(text) + elif not sampler_found: + pass + + +def _workflow_apply_prompt_fallback( + nodes: List[Dict[str, Any]], + unique_prompts: Set[str], + unique_negatives: Set[str], +) -> None: + if unique_prompts or unique_negatives: + return + + for node in nodes: + widgets = node.get("widgets_values") + if not isinstance(widgets, list) or not widgets: + continue + node_type = str(node.get("type", "")).lower() + title = str(node.get("title", "")).lower() + if not _workflow_is_text_prompt_node(node_type): + continue + value = _workflow_first_nontrivial_text(widgets) + if not value: + continue + if _workflow_is_negative_prompt_node(title, node_type): + unique_negatives.add(value) + else: + unique_prompts.add(value) + + +def _workflow_is_text_prompt_node(node_type: str) -> bool: + return "text" in node_type or "string" in node_type or "prompt" in node_type + + +def _workflow_first_nontrivial_text(widgets: List[Any]) -> Optional[str]: + for value in widgets: + if isinstance(value, str) and len(value) > 2: + return value + return None + + +def _workflow_is_negative_prompt_node(title: str, node_type: str) -> bool: + return "negative" in title or "negative" in node_type + + +def _workflow_finalize_params( + params: Dict[str, Any], + unique_prompts: Set[str], + unique_negatives: Set[str], +) -> None: + if unique_prompts: + params["positive_prompt"] = "\n".join(unique_prompts) + if unique_negatives: + params["negative_prompt"] = "\n".join(unique_negatives) + fake_text_parts = _workflow_prompt_lines(params) + details = _workflow_detail_lines(params) + if details: + fake_text_parts.append(", ".join(details)) + if fake_text_parts: + params["parameters"] = "\n".join(fake_text_parts) + + +def _workflow_prompt_lines(params: Dict[str, Any]) -> List[str]: + out: List[str] = [] + if "positive_prompt" in params: + out.append(params["positive_prompt"]) + if "negative_prompt" in params: + out.append(f"Negative prompt: {params['negative_prompt']}") + return out + + +def _workflow_detail_lines(params: Dict[str, Any]) -> List[str]: + out: List[str] = [] + if "steps" in params: + out.append(f"Steps: {params['steps']}") + if "sampler" in params: + out.append(f"Sampler: {params['sampler']}") + if "cfg" in params: + out.append(f"CFG scale: {params['cfg']}") + if "seed" in params: + out.append(f"Seed: {params['seed']}") + if "model" in params: + out.append(f"Model: {params['model']}") + return out def _reconstruct_params_from_workflow(workflow: Dict[str, Any]) -> Dict[str, Any]: @@ -26,335 +769,49 @@ def _reconstruct_params_from_workflow(workflow: Dict[str, Any]) -> Dict[str, Any Uses link traversal to distinguish Positive vs Negative prompts. """ params: Dict[str, Any] = {} - nodes = workflow.get("nodes", []) - links = workflow.get("links", []) # [[id, src_node, src_slot, tgt_node, tgt_slot, type], ...] - if not nodes: + raw_nodes = workflow.get("nodes", []) + links = workflow.get("links", []) + if not isinstance(raw_nodes, list) or not raw_nodes: return params - # 1. Index Nodes by ID and Links by Target - node_map = {n["id"]: n for n in nodes if "id" in n} - - # link_lookup: link_id -> (origin_node_id, origin_slot_idx, target_node_id, target_slot_idx) - link_lookup = {} - for link in links: - # [id, src_node, src_slot, tgt_node, tgt_slot, type] - if len(link) >= 4: - link_lookup[link[0]] = (link[1], link[2], link[3], link[4]) + nodes = [node for node in raw_nodes if isinstance(node, dict)] + node_map = {node["id"]: node for node in nodes if "id" in node} + link_lookup = _workflow_build_link_lookup(links) - def get_source_data(target_node_id, input_name): - """Find the (origin_node, link_id) connected to a specific named input.""" - t_node = node_map.get(target_node_id) - if not t_node: return None - - inp_def = next((i for i in t_node.get("inputs", []) if i.get("name") == input_name), None) - if not inp_def: return None - - link_id = inp_def.get("link") - if not link_id: return None - - link_info = link_lookup.get(link_id) - if not link_info: return None - - origin_id = link_info[0] - return (node_map.get(origin_id), link_id) - - def get_node_text(node, context=None): - """ - Extract text widget value from a node. - If context='negative' and multiple widgets exist, prefers the second/last one. - If context='positive' and multiple widgets exist, prefers the first one. - """ - widgets = node.get("widgets_values") - if not widgets or not isinstance(widgets, list): - return None - - # Filter for string widgets that look like prompts - s_widgets = [w for w in widgets if isinstance(w, str) and len(w.strip()) > 0] - - if not s_widgets: - return None - - if len(s_widgets) == 1: - return s_widgets[0] - - # Heuristic for multi-widget nodes (like Weaver's PromptToLoras) - if context == "negative": - # Return the last one (usually negative) - return s_widgets[-1] - - # Default/Positive: Return the first one - return s_widgets[0] - - def find_upstream_text(start_node, depth=0, seen=None, context=None): - """Find prompt-ish text upstream with semantic filtering. - - Converted to an iterative traversal to avoid relying on Python recursion - limits on deep or cyclic graphs. - """ - if not start_node: - return [] - - if seen is None: - seen = set() - - # Allow deeper traversal than before without recursion risk. - max_depth = 32 - - found_texts = [] - stack = [(start_node, depth)] - - while stack: - node, d = stack.pop() - if not node or d > max_depth: - continue - - node_id = node.get("id") - if node_id is not None: - if node_id in seen: - continue - seen.add(node_id) - - node_type = str(node.get("type", "")).lower() - - # 1. Widgets in this node - # We assume if a node has text widgets, it contributes to the prompt. - # But we must be careful with Loaders (filenames). - if "loader" not in node_type and "checkpoint" not in node_type and "loadimage" not in node_type: - txt = get_node_text(node, context) - if txt: - found_texts.append(txt) - - # 2. Walk upstream - inputs = node.get("inputs", []) - if not isinstance(inputs, list) or not inputs: - continue - - # Use reversed order so traversal order stays close to the previous - # recursive implementation. - for inp in reversed(inputs): - link_id = inp.get("link") - if not link_id: - continue - - name = str(inp.get("name", "")).lower() - if context == "positive" and "negative" in name: - continue - if context == "negative" and "positive" in name: - continue - - link_info = link_lookup.get(link_id) - if not link_info: - continue - - origin_id = link_info[0] - origin_node = node_map.get(origin_id) - if origin_node: - stack.append((origin_node, d + 1)) - - return found_texts - - # Detect Samplers - unique_prompts = set() - unique_negatives = set() - visited_nodes = set() + unique_prompts: Set[str] = set() + unique_negatives: Set[str] = set() + visited_nodes: Set[Any] = set() sampler_found = False for node in nodes: node_type = str(node.get("type", "")).lower() - if "ksampler" in node_type or "sampler" in node_type: - sampler_found = True - - # Extract Params from Widgets - widgets = node.get("widgets_values", []) - for val in widgets: - if isinstance(val, int): - if val > 10000 and "seed" not in params: params["seed"] = val - elif 1 <= val <= 200 and "steps" not in params: params["steps"] = val - elif isinstance(val, float): - if 1.0 <= val <= 30.0 and "cfg" not in params: params["cfg"] = val - elif isinstance(val, str): - v_low = val.lower() - if v_low in ["euler", "euler_a", "dpm++", "ddim", "uni_pc"] and "sampler" not in params: - params["sampler"] = val - - # Trace Prompts - # 1. Positive - pos_info = get_source_data(node["id"], "positive") - if pos_info: - # Use a local set for this trace to allow the same node to be visited - # in the negative trace (e.g. dual-role nodes like PromptToLoras) - trace_seen: set[Any] = set() - pos_texts = find_upstream_text(pos_info[0], context="positive", seen=trace_seen) - unique_prompts.update(pos_texts) - visited_nodes.update(trace_seen) - - # 2. Negative - neg_info = get_source_data(node["id"], "negative") - if neg_info: - trace_seen = set() - neg_texts = find_upstream_text(neg_info[0], context="negative", seen=trace_seen) - unique_negatives.update(neg_texts) - visited_nodes.update(trace_seen) - - # 3. Flux/Advanced Sampler (via Guider) - guider_info = get_source_data(node["id"], "guider") - if guider_info: - guider_node = guider_info[0] - if guider_node: - # Guider usually has 'conditioning' or 'positive' input - cond_info = get_source_data(guider_node["id"], "conditioning") or get_source_data(guider_node["id"], "positive") - if cond_info: - trace_seen = set() - pos_texts = find_upstream_text(cond_info[0], context="positive", seen=trace_seen) - unique_prompts.update(pos_texts) - visited_nodes.update(trace_seen) - - # Some guiders (CFGGuider) also have negative - neg_guider_info = get_source_data(guider_node["id"], "negative") - if neg_guider_info: - trace_seen = set() - neg_texts = find_upstream_text(neg_guider_info[0], context="negative", seen=trace_seen) - unique_negatives.update(neg_texts) - visited_nodes.update(trace_seen) - - # 4. Model - model_info = get_source_data(node["id"], "model") - if model_info: - model_node = model_info[0] - # Often connects to CheckpointLoaderSimple - widgets_m = model_node.get("widgets_values", []) - for w in widgets_m: - if isinstance(w, str) and (".safetensors" in w or ".ckpt" in w): - params["model"] = w - break - - # HEURISTIC FALLBACK / POST-PASS - # If we found disjoint text nodes (not connected to KSampler directly, e.g. WanVideo Bridge), - # we should try to include them by inferring their role from their DOWNSTREAM connections. - - # helper to check if a node connects to a positive/negative input anywhere - # Now recursive (depth-limited) to catch Resize/VAE chains in I2I/Video workflows - def classify_unconnected_node(start_node, depth=0, visited_trace=None): - if visited_trace is None: visited_trace = set() - - nid = start_node.get("id") - if nid in visited_trace or depth > 6: - return "unknown" - visited_trace.add(nid) - - outputs = start_node.get("outputs", []) - for out in outputs: - links_arr = out.get("links") - if not links_arr: continue - if not isinstance(links_arr, list): links_arr = [links_arr] - - for lid in links_arr: - # Check where this link goes - link_target = link_lookup.get(lid) - if not link_target: continue - - tgt_id = link_target[2] - tgt_node = node_map.get(tgt_id) - if not tgt_node: continue - - # 1. Direct Input Check - tgt_inputs = tgt_node.get("inputs", []) - for inp in tgt_inputs: - if inp.get("link") == lid: - iname = str(inp.get("name", "")).lower() - if "positive" in iname: return "positive" - if "negative" in iname: return "negative" - - # 2. Recursive downstream - role = classify_unconnected_node(tgt_node, depth + 1, visited_trace) - if role != "unknown": - return role - - return "unknown" - - for node in nodes: - uid = node.get("id") - if uid in visited_nodes: continue - - node_type = str(node.get("type", "")).lower() - title = str(node.get("title", "")).lower() - - # Check if it looks like a text node - if "text" in node_type or "string" in node_type or "prompt" in node_type: - # Check widgets first - txt = get_node_text(node) - if not txt: continue - - # Check Title hint - if "negative" in title: - unique_negatives.add(txt) - continue - if "positive" in title: - unique_prompts.add(txt) - continue - - # Check Downstream connections - role = classify_unconnected_node(node) - if role == "positive": - unique_prompts.add(txt) - elif role == "negative": - unique_negatives.add(txt) - elif not sampler_found: - # If NO sampler was found at all, we fall back to dumping everything into positive/negative based on heuristic - # But here we are in post-pass. - pass - - # HEURISTIC FALLBACK: If no sampler tracing worked (e.g. non-standard sampler names or broken links), - # revert to the "title contains negative" scan, BUT ONLY for nodes not already found. - if not unique_prompts and not unique_negatives: - for node in nodes: - widgets = node.get("widgets_values") - if not widgets or not isinstance(widgets, list): continue - - node_type = str(node.get("type", "")).lower() - title = str(node.get("title", "")).lower() - - if "text" in node_type or "string" in node_type or "prompt" in node_type: - for val in widgets: - if isinstance(val, str) and len(val) > 2: - if "negative" in title or "negative" in node_type: - unique_negatives.add(val) - else: - unique_prompts.add(val) - break - - # Construct output - # NOTE: Use "positive_prompt" not "prompt" to avoid overwriting the prompt graph dict - # in metadata. The prompt graph (dict) is used by geninfo parser; positive_prompt (str) - # is just the extracted text for display purposes. - if unique_prompts: - params["positive_prompt"] = "\n".join(unique_prompts) - - if unique_negatives: - params["negative_prompt"] = "\n".join(unique_negatives) - - # Build fake parameters text - fake_text_parts = [] - if "positive_prompt" in params: - fake_text_parts.append(params["positive_prompt"]) - if "negative_prompt" in params: - fake_text_parts.append(f"Negative prompt: {params['negative_prompt']}") - - details = [] - if "steps" in params: details.append(f"Steps: {params['steps']}") - if "sampler" in params: details.append(f"Sampler: {params['sampler']}") - if "cfg" in params: details.append(f"CFG scale: {params['cfg']}") - if "seed" in params: details.append(f"Seed: {params['seed']}") - if "model" in params: details.append(f"Model: {params['model']}") - - if details: - fake_text_parts.append(", ".join(details)) - - if fake_text_parts: - params["parameters"] = "\n".join(fake_text_parts) + if "ksampler" not in node_type and "sampler" not in node_type: + continue + sampler_found = True + _workflow_trace_sampler_fields( + node, + node_map, + link_lookup, + params, + unique_prompts, + unique_negatives, + visited_nodes, + ) + _workflow_collect_unconnected_texts( + nodes, + node_map, + link_lookup, + unique_prompts, + unique_negatives, + visited_nodes, + sampler_found, + ) + _workflow_apply_prompt_fallback(nodes, unique_prompts, unique_negatives) + _workflow_finalize_params(params, unique_prompts, unique_negatives) return params + def _coerce_rating_to_stars(value: Any) -> Optional[int]: """ Normalize rating values to 0..5 stars. @@ -364,6 +821,16 @@ def _coerce_rating_to_stars(value: Any) -> Optional[int]: - 0..100-ish "percent" (Windows SharedUserRating / RatingPercent) - string numbers """ + v = _coerce_rating_number(value) + if v is None: + return None + if v <= 5.0: + stars = int(round(v)) + return max(0, min(5, stars)) + return _coerce_percent_rating_to_stars(v) + + +def _coerce_rating_number(value: Any) -> Optional[float]: if value is None: return None try: @@ -375,26 +842,22 @@ def _coerce_rating_to_stars(value: Any) -> Optional[int]: return None value = float(value.replace(",", ".")) if isinstance(value, (int, float)): - v = float(value) - else: - return None + return float(value) except Exception: return None + return None - if v <= 5.0: - stars = int(round(v)) - return max(0, min(5, stars)) - # Treat as percent-like - if v <= 0: +def _coerce_percent_rating_to_stars(value: float) -> int: + if value <= 0: return 0 - if v >= 88: + if value >= 88: return 5 - if v >= 63: + if value >= 63: return 4 - if v >= 38: + if value >= 38: return 3 - if v >= 13: + if value >= 13: return 2 return 1 @@ -426,104 +889,85 @@ def _split_tags(text: str) -> list[str]: def _extract_rating_tags(exif_data: Optional[Dict]) -> tuple[Optional[int], list[str]]: if not exif_data or not isinstance(exif_data, dict): return (None, []) + norm = _build_exif_norm_map(exif_data) + rating = _extract_rating_from_norm(norm) + tags = _extract_tags_from_norm(norm) + return (rating, _dedupe_tag_list(tags)) - # ExifTool key normalization: - # - We run ExifTool with `-G1 -s`, which typically yields keys like `XMP-xmp:Rating` - # instead of the simpler `XMP:Rating`. - # - To avoid missing OS metadata (Windows Explorer stars/tags), normalize keys into - # a few predictable aliases and match against those. - # - # Examples: - # - `XMP-xmp:Rating` -> `xmp-xmp:rating`, `xmp:rating`, `rating` - # - `XMP-microsoft:RatingPercent` -> `xmp-microsoft:ratingpercent`, `microsoft:ratingpercent`, `ratingpercent` - def _build_norm_map(d: Dict[str, Any]) -> Dict[str, Any]: - out: Dict[str, Any] = {} - for k, v in d.items(): - try: - key = str(k) - except Exception: - continue - if not key: - continue - kl = key.strip().lower() - if not kl: - continue - out.setdefault(kl, v) - if ":" in kl: - group, tag = kl.split(":", 1) - if tag: - out.setdefault(tag, v) - # Reduce group like `xmp-xmp` -> `xmp` - group_last = group.split("-")[-1] if group else "" - if group_last and group_last != group: - out.setdefault(f"{group_last}:{tag}", v) - return out +def _build_exif_norm_map(data: Dict[str, Any]) -> Dict[str, Any]: + out: Dict[str, Any] = {} + for key, value in data.items(): + kl = _normalized_exif_key(key) + if not kl: + continue + out.setdefault(kl, value) + _add_exif_aliases(out, kl, value) + return out - norm = _build_norm_map(exif_data) - rating_candidates = ( - # Common XMP rating keys (as seen with ExifTool -G1) - "xmp:rating", - "xmp-xmp:rating", - # Windows Explorer often maps stars into XMP-microsoft rating percent - "microsoft:ratingpercent", - "xmp-microsoft:ratingpercent", - # Windows Property System / legacy keys - "microsoft:shareduserrating", - "xmp-microsoft:shareduserrating", - # Generic tag-only fallbacks - "rating", - "ratingpercent", - "shareduserrating", - ) - rating: Optional[int] = None - for key in rating_candidates: - if key in norm: - rating = _coerce_rating_to_stars(norm.get(key)) - if rating is not None: - break +def _normalized_exif_key(key: Any) -> str: + try: + text = str(key) + except Exception: + return "" + if not text: + return "" + return text.strip().lower() - tag_candidates = ( - # Common XMP subjects/keywords (ExifTool -G1) - "dc:subject", - "xmp-dc:subject", - "xmp:subject", - "iptc:keywords", - "photoshop:keywords", - "lr:hierarchicalsubject", - # Windows Explorer category/keywords - "microsoft:category", - "xmp-microsoft:category", - "xpkeywords", - # Generic tag-only fallbacks - "keywords", - "subject", - "category", - ) - tags: list[str] = [] - for key in tag_candidates: + +def _add_exif_aliases(out: Dict[str, Any], key_lower: str, value: Any) -> None: + if ":" not in key_lower: + return + group, tag = key_lower.split(":", 1) + if not tag: + return + out.setdefault(tag, value) + group_last = group.split("-")[-1] if group else "" + if group_last and group_last != group: + out.setdefault(f"{group_last}:{tag}", value) + + +def _extract_rating_from_norm(norm: Dict[str, Any]) -> Optional[int]: + for key in _RATING_CANDIDATE_KEYS: if key not in norm: continue - val = norm.get(key) - if val is None: + rating = _coerce_rating_to_stars(norm.get(key)) + if rating is not None: + return rating + return None + + +def _extract_tags_from_norm(norm: Dict[str, Any]) -> list[str]: + tags: list[str] = [] + for key in _TAG_CANDIDATE_KEYS: + if key not in norm: continue - if isinstance(val, list): - for item in val: - if item is None: - continue - tags.extend(_split_tags(str(item))) - else: - tags.extend(_split_tags(str(val))) - # De-dupe while keeping order - seen = set() + tags.extend(_split_tag_value(norm.get(key))) + return tags + + +def _split_tag_value(value: Any) -> list[str]: + if value is None: + return [] + if isinstance(value, list): + out: list[str] = [] + for item in value: + if item is not None: + out.extend(_split_tags(str(item))) + return out + return _split_tags(str(value)) + + +def _dedupe_tag_list(tags: list[str]) -> list[str]: + seen: Set[str] = set() deduped: list[str] = [] - for t in tags: - if t in seen: + for tag in tags: + if tag in seen: continue - seen.add(t) - deduped.append(t) - return (rating, deduped) + seen.add(tag) + deduped.append(tag) + return deduped def _extract_date_created(exif_data: Optional[Dict[str, Any]]) -> Optional[str]: @@ -532,7 +976,7 @@ def _extract_date_created(exif_data: Optional[Dict[str, Any]]) -> Optional[str]: """ if not exif_data: return None - + # Priority list for "Date Taken" / "Content Created" candidates = [ "ExifIFD:DateTimeOriginal", @@ -546,14 +990,14 @@ def _extract_date_created(exif_data: Optional[Dict[str, Any]]) -> Optional[str]: "XMP-photoshop:DateCreated", "Composite:DateTimeCreated" ] - + for key in candidates: val = exif_data.get(key) if val and isinstance(val, str): # Basic validation/cleaning could happen here # ExifTool usually returns "YYYY:MM:DD HH:MM:SS" or similar return val - + return None @@ -580,16 +1024,11 @@ def _bump_quality(metadata: Dict[str, Any], quality: str) -> None: metadata["quality"] = quality -def _apply_common_exif_fields( - metadata: Dict[str, Any], +def _resolve_workflow_prompt_fields( exif_data: Dict[str, Any], - *, - workflow: Optional[Dict[str, Any]] = None, - prompt: Optional[Dict[str, Any]] = None, -) -> None: - """ - Apply workflow/prompt + rating/tags into an existing metadata dict. - """ + workflow: Optional[Dict[str, Any]], + prompt: Optional[Dict[str, Any]], +) -> Tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]]]: wf = workflow pr = prompt if wf is None or pr is None: @@ -598,24 +1037,35 @@ def _apply_common_exif_fields( wf = scanned_workflow if pr is None: pr = scanned_prompt + return wf, pr - if looks_like_comfyui_workflow(wf): - metadata["workflow"] = wf + +def _apply_workflow_prompt_quality( + metadata: Dict[str, Any], + workflow: Optional[Dict[str, Any]], + prompt: Optional[Dict[str, Any]], +) -> None: + if looks_like_comfyui_workflow(workflow): + metadata["workflow"] = workflow _bump_quality(metadata, "full") - - if looks_like_comfyui_prompt_graph(pr): - metadata["prompt"] = pr + if looks_like_comfyui_prompt_graph(prompt): + metadata["prompt"] = prompt if str(metadata.get("quality") or "none") != "full": _bump_quality(metadata, "partial") - # FALLBACK: If we have a workflow but NO parameters text, reconstruct it - if metadata.get("workflow") and not metadata.get("parameters"): - reconstructed = _reconstruct_params_from_workflow(metadata["workflow"]) - if reconstructed: - metadata.update(reconstructed) - if metadata.get("quality") != "full": - _bump_quality(metadata, "partial") +def _apply_reconstructed_workflow_params(metadata: Dict[str, Any]) -> None: + if not metadata.get("workflow") or metadata.get("parameters"): + return + reconstructed = _reconstruct_params_from_workflow(metadata["workflow"]) + if not reconstructed: + return + metadata.update(reconstructed) + if metadata.get("quality") != "full": + _bump_quality(metadata, "partial") + + +def _apply_rating_tags_and_generation_time(metadata: Dict[str, Any], exif_data: Dict[str, Any]) -> None: rating, tags = _extract_rating_tags(exif_data) if rating is not None: metadata["rating"] = rating @@ -626,93 +1076,23 @@ def _apply_common_exif_fields( if date_created: metadata["generation_time"] = date_created - # Normalize visual dimensions from common ExifTool keys when present. - # This covers PNG/WEBP/JPG and fallback readers that expose ImageWidth/ImageHeight. + +def _apply_common_exif_fields( + metadata: Dict[str, Any], + exif_data: Dict[str, Any], + *, + workflow: Optional[Dict[str, Any]] = None, + prompt: Optional[Dict[str, Any]] = None, +) -> None: + """ + Apply workflow/prompt + rating/tags into an existing metadata dict. + """ + wf, pr = _resolve_workflow_prompt_fields(exif_data, workflow, prompt) + _apply_workflow_prompt_quality(metadata, wf, pr) + _apply_reconstructed_workflow_params(metadata) + _apply_rating_tags_and_generation_time(metadata, exif_data) if metadata.get("width") is None or metadata.get("height") is None: - width_keys = ( - "Image:ImageWidth", - "EXIF:ImageWidth", - "EXIF:ExifImageWidth", - "IFD0:ImageWidth", - "Composite:ImageWidth", - "QuickTime:ImageWidth", - "PNG:ImageWidth", - "File:ImageWidth", - "width", - ) - height_keys = ( - "Image:ImageHeight", - "EXIF:ImageHeight", - "EXIF:ExifImageHeight", - "IFD0:ImageHeight", - "Composite:ImageHeight", - "QuickTime:ImageHeight", - "PNG:ImageHeight", - "File:ImageHeight", - "height", - ) - - def _coerce_dim(value: Any) -> Optional[int]: - try: - if value is None: - return None - if isinstance(value, str): - txt = value.strip().lower().replace("px", "").strip() - if not txt: - return None - if "x" in txt: - # Avoid parsing "1920x1080" here (handled by pair parsing elsewhere). - return None - value = txt - out = int(float(value)) - return out if out > 0 else None - except Exception: - return None - - w = metadata.get("width") - h = metadata.get("height") - if w is None: - for k in width_keys: - w = _coerce_dim(exif_data.get(k)) - if w is not None: - break - if h is None: - for k in height_keys: - h = _coerce_dim(exif_data.get(k)) - if h is not None: - break - - # Parse paired size string as a final fallback (e.g. "1920 1080", "1920x1080"). - if w is None or h is None: - pair_keys = ("Composite:ImageSize", "Image:ImageSize", "QuickTime:ImageSize", "ImageSize") - for k in pair_keys: - raw = exif_data.get(k) - if not raw: - continue - try: - txt = str(raw).strip().lower().replace("px", "").replace("×", "x") - txt = txt.replace(",", " ").replace(";", " ") - if "x" in txt: - a, b = [s.strip() for s in txt.split("x", 1)] - else: - parts = [p for p in txt.split() if p] - if len(parts) < 2: - continue - a, b = parts[0], parts[1] - ww = _coerce_dim(a) - hh = _coerce_dim(b) - if ww and hh: - if w is None: - w = ww - if h is None: - h = hh - break - except Exception: - continue - - if w is not None and h is not None: - metadata["width"] = int(w) - metadata["height"] = int(h) + _apply_dimensions_from_exif(metadata, exif_data) def _build_a1111_geninfo(parsed: Dict[str, Any]) -> Optional[Dict[str, Any]]: @@ -722,6 +1102,15 @@ def _build_a1111_geninfo(parsed: Dict[str, Any]) -> Optional[Dict[str, Any]]: out: Dict[str, Any] = {"engine": {"parser_version": "geninfo-params-v1", "source": "parameters"}} + _apply_a1111_prompt_fields(out, parsed) + _apply_a1111_numeric_fields(out, parsed) + _apply_a1111_size_field(out, parsed) + _apply_a1111_checkpoint_field(out, parsed) + + return out if len(out) > 1 else None + + +def _apply_a1111_prompt_fields(out: Dict[str, Any], parsed: Dict[str, Any]) -> None: pos = parsed.get("prompt") neg = parsed.get("negative_prompt") if isinstance(pos, str) and pos.strip(): @@ -729,50 +1118,50 @@ def _build_a1111_geninfo(parsed: Dict[str, Any]) -> Optional[Dict[str, Any]]: if isinstance(neg, str) and neg.strip(): out["negative"] = {"value": neg.strip(), "confidence": "high", "source": "parameters"} - try: - steps = parsed.get("steps") - if steps is not None: - out["steps"] = {"value": int(steps), "confidence": "high", "source": "parameters"} - except Exception: - pass - try: - cfg = parsed.get("cfg") - if cfg is not None: - out["cfg"] = {"value": float(cfg), "confidence": "high", "source": "parameters"} - except Exception: - pass - try: - seed = parsed.get("seed") - if seed is not None: - out["seed"] = {"value": int(seed), "confidence": "high", "source": "parameters"} - except Exception: - pass - try: - width = parsed.get("width") - height = parsed.get("height") - if width is not None and height is not None: - out["size"] = { - "width": int(width), - "height": int(height), - "confidence": "high", - "source": "parameters", - } - except Exception: - pass +def _apply_a1111_numeric_fields(out: Dict[str, Any], parsed: Dict[str, Any]) -> None: + _apply_a1111_numeric_field(out, "steps", parsed.get("steps"), int) + _apply_a1111_numeric_field(out, "cfg", parsed.get("cfg"), float) + _apply_a1111_numeric_field(out, "seed", parsed.get("seed"), int) + +def _apply_a1111_numeric_field(out: Dict[str, Any], key: str, value: Any, caster: Any) -> None: + if value is None: + return + try: + out[key] = {"value": caster(value), "confidence": "high", "source": "parameters"} + except Exception: + return + + +def _apply_a1111_size_field(out: Dict[str, Any], parsed: Dict[str, Any]) -> None: + width = parsed.get("width") + height = parsed.get("height") + if width is None or height is None: + return + try: + out["size"] = { + "width": int(width), + "height": int(height), + "confidence": "high", + "source": "parameters", + } + except Exception: + return + + +def _apply_a1111_checkpoint_field(out: Dict[str, Any], parsed: Dict[str, Any]) -> None: model = parsed.get("model") - if isinstance(model, str) and model.strip(): - checkpoint = model.strip().replace("\\", "/").split("/")[-1] - lower = checkpoint.lower() - for ext in (".safetensors", ".ckpt", ".pt", ".pth", ".bin", ".gguf", ".json"): - if lower.endswith(ext): - checkpoint = checkpoint[: -len(ext)] - break - if checkpoint: - out["checkpoint"] = {"name": checkpoint, "confidence": "high", "source": "parameters"} - - return out if len(out) > 1 else None + if not isinstance(model, str) or not model.strip(): + return + checkpoint = model.strip().replace("\\", "/").split("/")[-1] + lower = checkpoint.lower() + for ext in (".safetensors", ".ckpt", ".pt", ".pth", ".bin", ".gguf", ".json"): + if lower.endswith(ext): + checkpoint = checkpoint[: -len(ext)] + break + if checkpoint: + out["checkpoint"] = {"name": checkpoint, "confidence": "high", "source": "parameters"} def extract_png_metadata(file_path: str, exif_data: Optional[Dict] = None) -> Result[Dict[str, Any]]: @@ -806,6 +1195,117 @@ def extract_png_metadata(file_path: str, exif_data: Optional[Dict] = None) -> Re logger.warning(f"PNG metadata extraction error: {e}") return Result.Err(ErrorCode.PARSE_ERROR, str(e), quality="degraded") + +def _merge_scanned_workflow_prompt( + workflow: Optional[Dict[str, Any]], + prompt: Optional[Dict[str, Any]], + container: Optional[Dict[str, Any]], +) -> Tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]]]: + if not isinstance(container, dict): + return workflow, prompt + scanned_workflow, scanned_prompt = _extract_json_fields(container) + if workflow is None and looks_like_comfyui_workflow(scanned_workflow): + workflow = scanned_workflow + if prompt is None and looks_like_comfyui_prompt_graph(scanned_prompt): + prompt = scanned_prompt + return workflow, prompt + + +def _scan_webp_text_fields( + metadata: Dict[str, Any], + exif_data: Dict[str, Any], + workflow: Optional[Dict[str, Any]], + prompt: Optional[Dict[str, Any]], +) -> Tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]]]: + for key in _WEBP_TEXT_KEYS: + candidate = exif_data.get(key) + if not isinstance(candidate, str) or not candidate: + continue + workflow, prompt, matched = _try_merge_webp_json_candidate(candidate, workflow, prompt) + if matched: + continue + _apply_webp_auto1111_candidate(metadata, candidate, prompt) + return workflow, prompt + + +def _try_merge_webp_json_candidate( + candidate_text: str, + workflow: Optional[Dict[str, Any]], + prompt: Optional[Dict[str, Any]], +) -> Tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]], bool]: + parsed_json = try_parse_json_text(candidate_text) + if parsed_json is None: + return workflow, prompt, False + return _merge_workflow_prompt_candidate(parsed_json, workflow, prompt) + + +def _apply_webp_auto1111_candidate( + metadata: Dict[str, Any], + candidate_text: str, + prompt: Optional[Dict[str, Any]], +) -> None: + parsed_a1111 = parse_auto1111_params(candidate_text) + if not parsed_a1111: + return + metadata["parameters"] = candidate_text + metadata.update(parsed_a1111) + if metadata.get("quality") != "full": + _bump_quality(metadata, "partial") + if prompt is None and not metadata.get("prompt") and parsed_a1111.get("prompt"): + metadata["prompt"] = parsed_a1111["prompt"] + + +def _apply_video_ffprobe_fields(metadata: Dict[str, Any], ffprobe_data: Optional[Dict[str, Any]]) -> None: + if not isinstance(ffprobe_data, dict): + return + video_stream = ffprobe_data.get("video_stream", {}) + format_info = ffprobe_data.get("format", {}) + metadata["resolution"] = ( + video_stream.get("width"), + video_stream.get("height"), + ) + metadata["fps"] = video_stream.get("r_frame_rate") + metadata["duration"] = format_info.get("duration") + + +def _scan_video_workflow_prompt_from_sources( + exif_data: Dict[str, Any], + ffprobe_data: Optional[Dict[str, Any]], +) -> Tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]], Dict[str, Any], List[Dict[str, Any]]]: + workflow: Optional[Dict[str, Any]] = None + prompt: Optional[Dict[str, Any]] = None + potential_workflow = _inspect_json_field(exif_data, _VIDEO_WORKFLOW_KEYS) + potential_prompt = _inspect_json_field(exif_data, _VIDEO_PROMPT_KEYS) + workflow, prompt, _ = _merge_workflow_prompt_candidate(potential_workflow, workflow, prompt) + workflow, prompt, _ = _merge_workflow_prompt_candidate(potential_prompt, workflow, prompt) + if workflow is None or prompt is None: + workflow, prompt = _merge_scanned_workflow_prompt(workflow, prompt, exif_data) + + format_tags = _extract_ffprobe_format_tags(ffprobe_data) + if workflow is None or prompt is None: + workflow, prompt = _merge_scanned_workflow_prompt(workflow, prompt, format_tags) + + stream_tag_dicts = _extract_ffprobe_stream_tag_dicts(ffprobe_data) + if workflow is None or prompt is None: + for tags in stream_tag_dicts: + workflow, prompt = _merge_scanned_workflow_prompt(workflow, prompt, tags) + if workflow is not None and prompt is not None: + break + return workflow, prompt, format_tags, stream_tag_dicts + + +def _collect_video_text_candidates( + exif_data: Dict[str, Any], + format_tags: Dict[str, Any], + stream_tag_dicts: List[Dict[str, Any]], +) -> List[str]: + text_candidates = _collect_text_candidates(exif_data) + if format_tags: + text_candidates.extend(_collect_text_candidates(format_tags)) + for tags in stream_tag_dicts: + text_candidates.extend(_collect_text_candidates(tags)) + return text_candidates + def extract_webp_metadata(file_path: str, exif_data: Optional[Dict] = None) -> Result[Dict[str, Any]]: """ Extract ComfyUI metadata from WEBP file. @@ -825,76 +1325,19 @@ def extract_webp_metadata(file_path: str, exif_data: Optional[Dict] = None) -> R return Result.Ok(metadata, quality="none") try: - # 1) Direct check in standard ComfyUI locations for WebP - def _inspect_json_field(key_names, container): - for key in key_names: - value = container.get(key) if isinstance(container, dict) else None - parsed = parse_json_value(value) - if parsed is not None: - return parsed - return None - workflow = None prompt = None - potential_workflow = _inspect_json_field(["EXIF:Make", "IFD0:Make", "Keys:Workflow", "comfyui:workflow"], exif_data) - potential_prompt = _inspect_json_field(["EXIF:Model", "IFD0:Model", "Keys:Prompt", "comfyui:prompt"], exif_data) + potential_workflow = _inspect_json_field(exif_data, _WEBP_WORKFLOW_KEYS) + potential_prompt = _inspect_json_field(exif_data, _WEBP_PROMPT_KEYS) - # Cross-check because sometimes Prompt is in Make, or Workflow is in ImageDescription - if potential_workflow: - if looks_like_comfyui_workflow(potential_workflow): - workflow = potential_workflow - elif looks_like_comfyui_prompt_graph(potential_workflow) and not prompt: - prompt = potential_workflow + workflow, prompt, _ = _merge_workflow_prompt_candidate(potential_workflow, workflow, prompt) + workflow, prompt, _ = _merge_workflow_prompt_candidate(potential_prompt, workflow, prompt) - if potential_prompt: - if looks_like_comfyui_prompt_graph(potential_prompt) and not prompt: - prompt = potential_prompt - elif looks_like_comfyui_workflow(potential_prompt) and not workflow: - workflow = potential_prompt + if workflow is None or prompt is None: + workflow, prompt = _merge_scanned_workflow_prompt(workflow, prompt, exif_data) - # 2) Fallback: scan all tags (optimized scan) - if not workflow or not prompt: - scanned_workflow, scanned_prompt = _extract_json_fields(exif_data) - if not workflow: - workflow = scanned_workflow - if not prompt: - prompt = scanned_prompt - - # 3) Look for Auto1111-style text OR prefixed JSON in description fields - # This fixes files where workflow is stored as "Workflow: {...}" in ImageDescription - text_keys = [ - "EXIF:ImageDescription", "IFD0:ImageDescription", "ImageDescription", - "EXIF:UserComment", "IFD0:UserComment", "UserComment", - "EXIF:Comment", "IFD0:Comment", - "EXIF:Subject", "IFD0:Subject", - ] - - for key in text_keys: - candidate = exif_data.get(key) - if not candidate or not isinstance(candidate, str): - continue - - # A) Try to parse as JSON first (handles "Workflow: {...}") - parsed_json = try_parse_json_text(candidate) - if parsed_json: - if workflow is None and looks_like_comfyui_workflow(parsed_json): - workflow = parsed_json - continue - if prompt is None and looks_like_comfyui_prompt_graph(parsed_json): - prompt = parsed_json - continue - - # B) Try to parse as Auto1111 parameters - parsed_a1111 = parse_auto1111_params(candidate) - if parsed_a1111: - metadata["parameters"] = candidate - metadata.update(parsed_a1111) - if metadata["quality"] != "full": - _bump_quality(metadata, "partial") - # Don't overwrite existing prompt graph with simple text prompt - if not prompt and not metadata.get("prompt") and parsed_a1111.get("prompt"): - metadata["prompt"] = parsed_a1111["prompt"] + workflow, prompt = _scan_webp_text_fields(metadata, exif_data, workflow, prompt) _apply_common_exif_fields(metadata, exif_data, workflow=workflow, prompt=prompt) @@ -936,188 +1379,17 @@ def extract_video_metadata(file_path: str, exif_data: Optional[Dict] = None, ffp } try: - # Extract video properties from ffprobe - if ffprobe_data: - video_stream = ffprobe_data.get("video_stream", {}) - format_info = ffprobe_data.get("format", {}) - - metadata["resolution"] = ( - video_stream.get("width"), - video_stream.get("height") - ) - metadata["fps"] = video_stream.get("r_frame_rate") - metadata["duration"] = format_info.get("duration") - - def _unwrap_workflow_prompt_container(container: Any) -> Tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]]]: - """ - Some pipelines embed a wrapper object like: - {"workflow": {...}, "prompt": "{...json...}"} - inside a single tag (often `ItemList:Comment`). - - Only return values if they match ComfyUI shapes (prevents false positives). - """ - if not isinstance(container, dict): - return (None, None) - - wf = container.get("workflow") or container.get("Workflow") or None - pr = container.get("prompt") or container.get("Prompt") or None - - wf_out: Optional[Dict[str, Any]] = None - pr_out: Optional[Dict[str, Any]] = None - - if isinstance(wf, dict) and looks_like_comfyui_workflow(wf): - wf_out = wf - - if isinstance(pr, dict) and looks_like_comfyui_prompt_graph(pr): - pr_out = pr - elif isinstance(pr, str): - # Prompt can be a JSON string literal. - parsed = try_parse_json_text(pr) - if isinstance(parsed, dict) and looks_like_comfyui_prompt_graph(parsed): - pr_out = parsed - - if wf_out or pr_out: - return (wf_out, pr_out) - return (None, None) - - def _inspect_json_field(key_names, container): - for key in key_names: - value = container.get(key) if isinstance(container, dict) else None - parsed = parse_json_value(value) - - # Check for wrapped container (Legacy behavior restoration) - if isinstance(parsed, dict): - wf_wrapped, pr_wrapped = _unwrap_workflow_prompt_container(parsed) - if wf_wrapped or pr_wrapped: - return parsed - - if parsed is not None: - return parsed - return None - - def _collect_text_candidates(container): - if not isinstance(container, dict): - return [] - out = [] - for k, v in container.items(): - if isinstance(v, str): - out.append((k, v)) - elif isinstance(v, (list, tuple)): - for item in v: - if isinstance(item, str): - out.append((k, item)) - return out - - # NOTE: Some encoders store JSON in ItemList:Comment that is NOT a ComfyUI workflow. - # We intentionally do not treat ItemList:Comment as a workflow source unless it - # clearly contains a workflow payload (see fallback scan below). - workflow = None - prompt = None - - potential_workflow = _inspect_json_field( - ["QuickTime:Workflow", "Keys:Workflow", "comfyui:workflow"], - exif_data + _apply_video_ffprobe_fields(metadata, ffprobe_data) + workflow, prompt, format_tags, stream_tag_dicts = _scan_video_workflow_prompt_from_sources(exif_data, ffprobe_data) + text_candidates = _collect_video_text_candidates(exif_data, format_tags, stream_tag_dicts) + _apply_auto1111_text_candidates( + metadata, + text_candidates, + prompt_graph=prompt, + preserve_existing_prompt_text=False, ) - potential_prompt = _inspect_json_field( - ["QuickTime:Prompt", "Keys:Prompt", "comfyui:prompt"], - exif_data - ) - - # Cross-check to handle swapped tags (e.g. prompt in workflow tag) - if potential_workflow: - if looks_like_comfyui_workflow(potential_workflow): - workflow = potential_workflow - elif looks_like_comfyui_prompt_graph(potential_workflow) and not prompt: - prompt = potential_workflow - else: - # Check for legacy wrapper format in tags like "QuickTime:Workflow" - wf_w, pr_w = _unwrap_workflow_prompt_container(potential_workflow) - if wf_w and not workflow: workflow = wf_w - if pr_w and not prompt: prompt = pr_w - - if potential_prompt: - if looks_like_comfyui_prompt_graph(potential_prompt) and not prompt: - prompt = potential_prompt - elif looks_like_comfyui_workflow(potential_prompt) and not workflow: - workflow = potential_prompt - else: - wf_w, pr_w = _unwrap_workflow_prompt_container(potential_prompt) - if wf_w and not workflow: workflow = wf_w - if pr_w and not prompt: prompt = pr_w - - # Fallback scan across all ExifTool tags, using shape-based heuristics. - if workflow is None or prompt is None: - scanned_workflow, scanned_prompt = _extract_json_fields(exif_data) - if workflow is None and looks_like_comfyui_workflow(scanned_workflow): - workflow = scanned_workflow - if prompt is None and looks_like_comfyui_prompt_graph(scanned_prompt): - prompt = scanned_prompt - - # Fallback: check ffprobe container tags (some encoders store metadata here). - format_tags = None - if isinstance(ffprobe_data, dict): - fmt = ffprobe_data.get("format") or {} - if isinstance(fmt, dict): - tags = fmt.get("tags") - if isinstance(tags, dict): - format_tags = tags - - if format_tags and (workflow is None or prompt is None): - scanned_workflow, scanned_prompt = _extract_json_fields(format_tags) - if workflow is None and looks_like_comfyui_workflow(scanned_workflow): - workflow = scanned_workflow - if prompt is None and looks_like_comfyui_prompt_graph(scanned_prompt): - prompt = scanned_prompt - - # VHS/other pipelines may store tags on the stream rather than container. - stream_tag_dicts = [] - try: - if isinstance(ffprobe_data, dict) and isinstance(ffprobe_data.get("streams"), list): - for s in ffprobe_data.get("streams") or []: - if not isinstance(s, dict): - continue - tags = s.get("tags") - if isinstance(tags, dict): - stream_tag_dicts.append(tags) - except Exception: - stream_tag_dicts = [] - - if stream_tag_dicts and (workflow is None or prompt is None): - for tags in stream_tag_dicts: - scanned_workflow, scanned_prompt = _extract_json_fields(tags) - if workflow is None and looks_like_comfyui_workflow(scanned_workflow): - workflow = scanned_workflow - if prompt is None and looks_like_comfyui_prompt_graph(scanned_prompt): - prompt = scanned_prompt - if workflow is not None and prompt is not None: - break - - # Auto1111-style "parameters" text can also appear in video tags (comment/description). - if metadata.get("parameters") is None: - candidates = _collect_text_candidates(exif_data) - if format_tags: - candidates.extend(_collect_text_candidates(format_tags)) - for tags in stream_tag_dicts or []: - candidates.extend(_collect_text_candidates(tags)) - for _, text in candidates: - parsed = parse_auto1111_params(text) - if not parsed: - continue - metadata["parameters"] = text - # Keep prompt as a string-based fallback only if we don't have a prompt graph. - if prompt is None and parsed.get("prompt"): - metadata["prompt"] = parsed.get("prompt") - if parsed.get("negative_prompt"): - metadata["negative_prompt"] = parsed.get("negative_prompt") - if metadata["quality"] != "full": - _bump_quality(metadata, "partial") - break - - # NOTE: No sidecar fallback here (e.g. .json/.txt). We only trust embedded tags - # (ExifTool/FFprobe) for video generation metadata to avoid hidden filesystem coupling. _apply_common_exif_fields(metadata, exif_data, workflow=workflow, prompt=prompt) - return Result.Ok(metadata, quality=metadata["quality"]) except Exception as e: @@ -1129,85 +1401,109 @@ def _extract_json_fields(exif_data: Dict[str, Any]) -> Tuple[Optional[Dict[str, Optimized scan of EXIF/ffprobe metadata for workflow/prompt JSON fields. Prioritizes known keys before scanning strictly. """ - workflow = None - prompt = None - - # 1. Prioritize known keys where JSON is typically found - # This avoids expensive regex/zlib on thousands of irrelevant EXIF tags - PRIORITY_KEYS = [ - "UserComment", "Comment", "Description", "ImageDescription", - "Parameters", "Workflow", "Prompt", "ExifOffset", "Make", "Model" - ] - - # Helper to check if a key matches our priority list (case-insensitive partial match) - def is_priority_key(k: str) -> bool: - k_lower = k.lower() - return any(pk.lower() in k_lower for pk in PRIORITY_KEYS) - - # Sort keys to process priority ones first - sorted_items = sorted( - exif_data.items(), - key=lambda item: 0 if is_priority_key(str(item[0])) else 1 - ) - - def _unwrap_workflow_prompt_container(container: Any) -> Tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]]]: - if not isinstance(container, dict): - return (None, None) - - wf = container.get("workflow") or container.get("Workflow") - pr = container.get("prompt") or container.get("Prompt") - - wf_out, pr_out = None, None - - if isinstance(wf, dict) and looks_like_comfyui_workflow(wf): - wf_out = wf - if isinstance(pr, dict) and looks_like_comfyui_prompt_graph(pr): - pr_out = pr - elif isinstance(pr, str): - parsed = try_parse_json_text(pr) - if isinstance(parsed, dict) and looks_like_comfyui_prompt_graph(parsed): - pr_out = parsed - return (wf_out, pr_out) - - for key, value in sorted_items: - # Stop early if we found both + workflow: Optional[Dict[str, Any]] = None + prompt: Optional[Dict[str, Any]] = None + for key, value in _priority_sorted_exif_items(exif_data): if workflow and prompt: break - - # Skip obviously short strings to save CPU if isinstance(value, str) and len(value) < 10: continue - - normalized = key.lower() if isinstance(key, str) else "" parsed = parse_json_value(value) - if not parsed: continue - - # Check for container pattern - if workflow is None or prompt is None: - wf_candidate, pr_candidate = _unwrap_workflow_prompt_container(parsed) - if workflow is None and wf_candidate: - workflow = wf_candidate - if prompt is None and pr_candidate: - prompt = pr_candidate - if workflow and prompt: - break - - # Check direct match - if workflow is None and looks_like_comfyui_workflow(parsed): - workflow = parsed - if prompt is None and looks_like_comfyui_prompt_graph(parsed): - prompt = parsed - - # Check prefix usage (workflow: ...) - if isinstance(value, str): - text_lower = value.strip().lower() - if workflow is None and (text_lower.startswith("workflow:") or "workflow" in normalized): - if looks_like_comfyui_workflow(parsed): - workflow = parsed - if prompt is None and (text_lower.startswith("prompt:") or "prompt" in normalized): - if looks_like_comfyui_prompt_graph(parsed): - prompt = parsed - + workflow, prompt = _merge_json_candidate(workflow, prompt, parsed, key, value) return workflow, prompt + + +def _priority_sorted_exif_items(exif_data: Dict[str, Any]) -> List[Tuple[Any, Any]]: + priority_keys = ( + "UserComment", + "Comment", + "Description", + "ImageDescription", + "Parameters", + "Workflow", + "Prompt", + "ExifOffset", + "Make", + "Model", + ) + return sorted( + exif_data.items(), + key=lambda item: 0 if _is_priority_key(str(item[0]), priority_keys) else 1, + ) + + +def _is_priority_key(key: str, priority_keys: Tuple[str, ...]) -> bool: + key_lower = key.lower() + return any(token.lower() in key_lower for token in priority_keys) + + +def _merge_json_candidate( + workflow: Optional[Dict[str, Any]], + prompt: Optional[Dict[str, Any]], + parsed: Any, + key: Any, + value: Any, +) -> Tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]]]: + workflow, prompt = _merge_container_candidate(workflow, prompt, parsed) + workflow = _merge_direct_workflow_candidate(workflow, parsed) + prompt = _merge_direct_prompt_candidate(prompt, parsed) + return _merge_prefixed_json_candidate(workflow, prompt, parsed, key, value) + + +def _merge_container_candidate( + workflow: Optional[Dict[str, Any]], + prompt: Optional[Dict[str, Any]], + parsed: Any, +) -> Tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]]]: + if workflow is not None and prompt is not None: + return workflow, prompt + wf_candidate, pr_candidate = _unwrap_workflow_prompt_container(parsed) + if workflow is None and wf_candidate: + workflow = wf_candidate + if prompt is None and pr_candidate: + prompt = pr_candidate + return workflow, prompt + + +def _merge_direct_workflow_candidate(workflow: Optional[Dict[str, Any]], parsed: Any) -> Optional[Dict[str, Any]]: + if workflow is None and looks_like_comfyui_workflow(parsed): + return parsed + return workflow + + +def _merge_direct_prompt_candidate(prompt: Optional[Dict[str, Any]], parsed: Any) -> Optional[Dict[str, Any]]: + if prompt is None and looks_like_comfyui_prompt_graph(parsed): + return parsed + return prompt + + +def _merge_prefixed_json_candidate( + workflow: Optional[Dict[str, Any]], + prompt: Optional[Dict[str, Any]], + parsed: Any, + key: Any, + value: Any, +) -> Tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]]]: + if not isinstance(value, str): + return workflow, prompt + normalized = _normalized_json_key(key) + text_lower = value.strip().lower() + if workflow is None and _looks_like_workflow_prefixed(text_lower, normalized) and looks_like_comfyui_workflow(parsed): + workflow = parsed + if prompt is None and _looks_like_prompt_prefixed(text_lower, normalized) and looks_like_comfyui_prompt_graph(parsed): + prompt = parsed + return workflow, prompt + + +def _normalized_json_key(key: Any) -> str: + return key.lower() if isinstance(key, str) else "" + + +def _looks_like_workflow_prefixed(text_lower: str, normalized: str) -> bool: + return text_lower.startswith("workflow:") or "workflow" in normalized + + +def _looks_like_prompt_prefixed(text_lower: str, normalized: str) -> bool: + return text_lower.startswith("prompt:") or "prompt" in normalized diff --git a/mjr_am_backend/features/metadata/fallback_readers.py b/mjr_am_backend/features/metadata/fallback_readers.py index b0b323ba..063e9fce 100644 --- a/mjr_am_backend/features/metadata/fallback_readers.py +++ b/mjr_am_backend/features/metadata/fallback_readers.py @@ -12,6 +12,56 @@ from typing import Any, Dict, Optional from PIL import Image, ExifTags +def _apply_image_info_exif_like(out: Dict[str, Any], info: Dict[str, Any]) -> None: + for key, value in info.items(): + if value is None: + continue + key_s = str(key) + key_l = key_s.lower() + if key_l == "parameters": + out["PNG:Parameters"] = value + out["Parameters"] = value + continue + if key_l == "workflow": + out["PNG:Workflow"] = value + out["Keys:Workflow"] = value + out["comfyui:workflow"] = value + continue + if key_l == "prompt": + out["PNG:Prompt"] = value + out["Keys:Prompt"] = value + out["comfyui:prompt"] = value + continue + if key_l in ("comment", "description"): + out["EXIF:ImageDescription"] = value + out["ImageDescription"] = value + continue + out[f"Pillow:{key_s}"] = value + + +def _safe_image_exif_dict(img: Any) -> Dict[Any, Any]: + try: + exif = img.getexif() + except Exception: + exif = None + return dict(exif) if exif else {} + + +def _apply_pillow_exif_fields(out: Dict[str, Any], exif_map: Dict[Any, Any]) -> None: + for tag_id, value in exif_map.items(): + try: + tag_name = ExifTags.TAGS.get(tag_id, str(tag_id)) + except Exception: + tag_name = str(tag_id) + out[f"EXIF:{tag_name}"] = value + if tag_name == "ImageDescription": + out["IFD0:ImageDescription"] = value + if tag_name == "Make": + out["IFD0:Make"] = value + if tag_name == "Model": + out["IFD0:Model"] = value + + def read_image_exif_like(path: str) -> Dict[str, Any]: """ Build an ExifTool-like dict from Pillow data for image files. @@ -26,45 +76,8 @@ def read_image_exif_like(path: str) -> Dict[str, Any]: out["Image:ImageHeight"] = int(img.height) info = dict(getattr(img, "info", {}) or {}) - for key, value in info.items(): - if value is None: - continue - key_s = str(key) - key_l = key_s.lower() - if key_l == "parameters": - out["PNG:Parameters"] = value - out["Parameters"] = value - elif key_l == "workflow": - out["PNG:Workflow"] = value - out["Keys:Workflow"] = value - out["comfyui:workflow"] = value - elif key_l == "prompt": - out["PNG:Prompt"] = value - out["Keys:Prompt"] = value - out["comfyui:prompt"] = value - elif key_l in ("comment", "description"): - out["EXIF:ImageDescription"] = value - out["ImageDescription"] = value - else: - out[f"Pillow:{key_s}"] = value - - try: - exif = img.getexif() - except Exception: - exif = None - if exif: - for tag_id, value in dict(exif).items(): - try: - tag_name = ExifTags.TAGS.get(tag_id, str(tag_id)) - except Exception: - tag_name = str(tag_id) - out[f"EXIF:{tag_name}"] = value - if tag_name == "ImageDescription": - out["IFD0:ImageDescription"] = value - if tag_name == "Make": - out["IFD0:Make"] = value - if tag_name == "Model": - out["IFD0:Model"] = value + _apply_image_info_exif_like(out, info) + _apply_pillow_exif_fields(out, _safe_image_exif_dict(img)) except Exception: return {} return out @@ -76,80 +89,135 @@ def read_media_probe_like(path: str) -> Dict[str, Any]: Returns an empty dict when hachoir is unavailable or parsing fails. """ - try: - from hachoir.parser import createParser - from hachoir.metadata import extractMetadata - except Exception: + parser, extract_metadata = _load_hachoir_parser(path) + if parser is None or extract_metadata is None: return {} - - try: - parser = createParser(path) - if parser is None: - return {} - except Exception: + lines = _extract_hachoir_lines(parser, extract_metadata) + if lines is None: return {} format_info: Dict[str, Any] = {"tags": {}} video_stream: Dict[str, Any] = {} audio_stream: Dict[str, Any] = {} + for raw_line in lines: + _apply_probe_line(str(raw_line), format_info, video_stream, audio_stream) + + return _build_media_probe_payload(format_info, video_stream, audio_stream) + + +def _load_hachoir_parser(path: str): + try: + from hachoir.parser import createParser + from hachoir.metadata import extractMetadata + except Exception: + return None, None + try: + parser = createParser(path) + if parser is None: + return None, None + except Exception: + return None, None + return parser, extractMetadata + + +def _extract_hachoir_lines(parser: Any, extract_metadata: Any) -> Optional[list[Any]]: try: with parser: - meta = extractMetadata(parser) + meta = extract_metadata(parser) if meta is None: - return {} - lines = list(meta.exportPlaintext() or []) + return None + return list(meta.exportPlaintext() or []) except Exception: - return {} + return None - for raw_line in lines: - line = str(raw_line).strip() - if ":" not in line: - continue - key, value = line.split(":", 1) - key_l = key.strip().lower() - val = value.strip() - if key_l == "duration": - format_info["duration"] = val - elif key_l in ("image width", "width"): - w = _to_int(val) - if w is not None: - video_stream["width"] = w - elif key_l in ("image height", "height"): - h = _to_int(val) - if h is not None: - video_stream["height"] = h - elif key_l in ("frame rate", "framerate", "fps"): - video_stream["r_frame_rate"] = val - elif key_l in ("audio codec", "codec"): - audio_stream["codec_name"] = val - elif key_l in ("sample rate", "samplerate"): - sr = _to_int(val) - audio_stream["sample_rate"] = sr if sr is not None else val - elif key_l in ("channel", "channels"): - ch = _to_int(val) - audio_stream["channels"] = ch if ch is not None else val - elif key_l in ("bit rate", "bitrate"): - br = _to_int(val) - if br is not None: - format_info["bit_rate"] = str(br) - audio_stream["bit_rate"] = str(br) - else: - if val: - format_info["tags"][key.strip()] = val +def _apply_probe_line( + raw_line: str, + format_info: Dict[str, Any], + video_stream: Dict[str, Any], + audio_stream: Dict[str, Any], +) -> None: + line = str(raw_line).strip() + if ":" not in line: + return + key, value = line.split(":", 1) + key_l = key.strip().lower() + val = value.strip() + if _apply_known_probe_line(key_l, val, format_info, video_stream, audio_stream): + return + if val: + format_info["tags"][key.strip()] = val - streams = [] + +def _apply_known_probe_line( + key_l: str, + val: str, + format_info: Dict[str, Any], + video_stream: Dict[str, Any], + audio_stream: Dict[str, Any], +) -> bool: + if key_l == "duration": + format_info["duration"] = val + return True + if key_l in ("image width", "width"): + _set_int_if_present(video_stream, "width", val) + return True + if key_l in ("image height", "height"): + _set_int_if_present(video_stream, "height", val) + return True + if key_l in ("frame rate", "framerate", "fps"): + video_stream["r_frame_rate"] = val + return True + if key_l in ("audio codec", "codec"): + audio_stream["codec_name"] = val + return True + if key_l in ("sample rate", "samplerate"): + _set_int_or_raw(audio_stream, "sample_rate", val) + return True + if key_l in ("channel", "channels"): + _set_int_or_raw(audio_stream, "channels", val) + return True + if key_l in ("bit rate", "bitrate"): + _apply_probe_bitrate(val, format_info, audio_stream) + return True + return False + + +def _apply_probe_bitrate(val: str, format_info: Dict[str, Any], audio_stream: Dict[str, Any]) -> None: + bitrate = _to_int(val) + if bitrate is None: + return + bitrate_str = str(bitrate) + format_info["bit_rate"] = bitrate_str + audio_stream["bit_rate"] = bitrate_str + + +def _set_int_if_present(target: Dict[str, Any], key: str, value: Any) -> None: + parsed = _to_int(value) + if parsed is not None: + target[key] = parsed + + +def _set_int_or_raw(target: Dict[str, Any], key: str, value: Any) -> None: + parsed = _to_int(value) + target[key] = parsed if parsed is not None else value + + +def _build_media_probe_payload( + format_info: Dict[str, Any], + video_stream: Dict[str, Any], + audio_stream: Dict[str, Any], +) -> Dict[str, Any]: + streams: list[Dict[str, Any]] = [] if video_stream: video_stream["codec_type"] = "video" streams.append(video_stream) if audio_stream: audio_stream["codec_type"] = "audio" streams.append(audio_stream) - if not format_info.get("duration") and not streams and not format_info.get("tags"): return {} - return { "format": format_info, "streams": streams, diff --git a/mjr_am_backend/features/metadata/parsing_utils.py b/mjr_am_backend/features/metadata/parsing_utils.py index d0c4225e..f654bcc2 100644 --- a/mjr_am_backend/features/metadata/parsing_utils.py +++ b/mjr_am_backend/features/metadata/parsing_utils.py @@ -12,7 +12,7 @@ logger = logging.getLogger(__name__) MAX_METADATA_JSON_SIZE = 10 * 1024 * 1024 # 10MB # 50MB Decompressed limit (Critical-003) -MAX_DECOMPRESSED_SIZE = 50 * 1024 * 1024 +MAX_DECOMPRESSED_SIZE = 50 * 1024 * 1024 MIN_BASE64_CANDIDATE_LEN = 80 _BASE64_RE = re.compile(r"^[A-Za-z0-9+/=\s]+$") @@ -26,10 +26,10 @@ def _safe_zlib_decompress(data: bytes, max_size: int = MAX_DECOMPRESSED_SIZE) -> try: decompressor = zlib.decompressobj() result = bytearray() - + chunk_size = 81920 # 80KB chunks offset = 0 - + while offset < len(data): chunk = decompressor.decompress(data[offset:offset + chunk_size]) if chunk: @@ -37,14 +37,14 @@ def _safe_zlib_decompress(data: bytes, max_size: int = MAX_DECOMPRESSED_SIZE) -> if len(result) > max_size: return None offset += chunk_size - + chunk = decompressor.flush() if chunk: result.extend(chunk) - + if len(result) > max_size: return None - + return bytes(result) except Exception: return None @@ -58,64 +58,80 @@ def try_parse_json_text(text: str) -> Optional[Dict[str, Any]]: if not raw: return None - # Handle common prefixes - lower_raw = raw.lower() - if lower_raw.startswith("workflow:"): - raw = raw[9:].strip() - elif lower_raw.startswith("prompt:"): - raw = raw[7:].strip() - elif lower_raw.startswith("makeprompt:"): - raw = raw[11:].strip() + raw = _strip_known_json_prefix(raw) if len(raw) > MAX_METADATA_JSON_SIZE: return None - def _loads_maybe(s: str) -> Optional[Dict[str, Any]]: - try: - parsed = json.loads(s) - except Exception: - return None - if isinstance(parsed, dict): - return parsed - if isinstance(parsed, str): - try: - nested = json.loads(parsed) - except Exception: - return None - return nested if isinstance(nested, dict) else None - return None - - direct = _loads_maybe(raw) + direct = _loads_maybe_dict(raw) if direct is not None: return direct - # Base64 check - if len(raw) < MIN_BASE64_CANDIDATE_LEN or len(raw) > (MAX_METADATA_JSON_SIZE * 2): - return None - if not _BASE64_RE.match(raw): + decoded = _decode_base64_candidate(raw) + if decoded is None: return None - try: - decoded = base64.b64decode(raw, validate=False) - except Exception: - return None - - # Zlib check - if decoded.startswith(b"x\x9c") or decoded.startswith(b"x\xda"): - decompressed = _safe_zlib_decompress(decoded) - if decompressed is not None: - decoded = decompressed - # Fallback to usage as-is if decompression fails or is skipped (unlikely valid though) - - try: - decoded_text = decoded.decode("utf-8", errors="replace").strip() - except Exception: + decoded = _maybe_decompress_zlib(decoded) + decoded_text = _decode_utf8_text(decoded) + if decoded_text is None: return None if not decoded_text or len(decoded_text) > MAX_METADATA_JSON_SIZE: return None - return _loads_maybe(decoded_text) + return _loads_maybe_dict(decoded_text) + + +def _strip_known_json_prefix(raw: str) -> str: + lower_raw = raw.lower() + if lower_raw.startswith("workflow:"): + return raw[9:].strip() + if lower_raw.startswith("prompt:"): + return raw[7:].strip() + if lower_raw.startswith("makeprompt:"): + return raw[11:].strip() + return raw + + +def _loads_maybe_dict(text: str) -> Optional[Dict[str, Any]]: + try: + parsed = json.loads(text) + except Exception: + return None + if isinstance(parsed, dict): + return parsed + if isinstance(parsed, str): + try: + nested = json.loads(parsed) + except Exception: + return None + return nested if isinstance(nested, dict) else None + return None + + +def _decode_base64_candidate(raw: str) -> Optional[bytes]: + if len(raw) < MIN_BASE64_CANDIDATE_LEN or len(raw) > (MAX_METADATA_JSON_SIZE * 2): + return None + if not _BASE64_RE.match(raw): + return None + try: + return base64.b64decode(raw, validate=False) + except Exception: + return None + + +def _maybe_decompress_zlib(decoded: bytes) -> bytes: + if not (decoded.startswith(b"x\x9c") or decoded.startswith(b"x\xda")): + return decoded + decompressed = _safe_zlib_decompress(decoded) + return decompressed if decompressed is not None else decoded + + +def _decode_utf8_text(decoded: bytes) -> Optional[str]: + try: + return decoded.decode("utf-8", errors="replace").strip() + except Exception: + return None def parse_json_value(value: Any) -> Optional[Dict[str, Any]]: @@ -162,21 +178,21 @@ def looks_like_comfyui_workflow(value: Optional[Dict[str, Any]]) -> bool: return False sample = nodes[:5] - valid_nodes = 0 - for node in sample: - if not isinstance(node, dict): - continue - if "type" in node and "id" in node: - valid_nodes += 1 - continue - if "id" in node and ("title" in node or "outputs" in node or "inputs" in node): - valid_nodes += 1 - continue - + valid_nodes = sum(1 for node in sample if _looks_like_workflow_node(node)) min_required = max(1, len(sample) // 2) return valid_nodes >= min_required +def _looks_like_workflow_node(node: Any) -> bool: + if not isinstance(node, dict): + return False + if "type" in node and "id" in node: + return True + if "id" not in node: + return False + return any(key in node for key in ("title", "outputs", "inputs")) + + def looks_like_comfyui_prompt_graph(value: Optional[Dict[str, Any]]) -> bool: """Heuristic check for ComfyUI prompt graph (runtime prompt dict).""" if not isinstance(value, dict) or not value: @@ -187,20 +203,28 @@ def looks_like_comfyui_prompt_graph(value: Optional[Dict[str, Any]]) -> bool: return False keys = list(value.keys())[:8] + digit_keys, valid_nodes = _prompt_graph_key_and_node_counts(value, keys) + threshold = max(2, len(keys) // 2) + return digit_keys >= threshold and valid_nodes >= threshold + + +def _prompt_graph_key_and_node_counts(graph: Dict[str, Any], keys: list[Any]) -> tuple[int, int]: digit_keys = 0 valid_nodes = 0 - for k in keys: - if looks_like_prompt_node_id(k): + for key in keys: + if looks_like_prompt_node_id(key): digit_keys += 1 - node = value.get(k) - if not isinstance(node, dict): - continue - ct = node.get("class_type") or node.get("type") - ins = node.get("inputs") - if isinstance(ct, str) and isinstance(ins, dict): + if _prompt_graph_node_looks_valid(graph.get(key)): valid_nodes += 1 + return digit_keys, valid_nodes - return digit_keys >= max(2, len(keys) // 2) and valid_nodes >= max(2, len(keys) // 2) + +def _prompt_graph_node_looks_valid(node: Any) -> bool: + if not isinstance(node, dict): + return False + ct = node.get("class_type") or node.get("type") + ins = node.get("inputs") + return isinstance(ct, str) and isinstance(ins, dict) def parse_auto1111_params(params_text: str) -> Optional[Dict[str, Any]]: @@ -214,92 +238,111 @@ def parse_auto1111_params(params_text: str) -> Optional[Dict[str, Any]]: return None result: Dict[str, Any] = {} - - remaining = "" - neg_marker = "Negative prompt:" - neg_idx = text.find(neg_marker) - if neg_idx != -1: - result["prompt"] = text[:neg_idx].strip() - after = text[neg_idx + len(neg_marker):].lstrip() - - # Find the start of the key-value parameters block (e.g. "Steps: 20, ...") - # We look for a line starting with "Steps:" or a standard key. - # Using regex to find the boundary more reliably than just "\n". - # Common starting keys: Steps, Size, Model hash, Seed, CFG scale. - # We look for "\n" followed by one of those keys. - param_start_match = re.search(r"\n(?:Steps|Size|Model|Seed|CFG|Sampler|Denoising|Ens|Version):", after) - - if param_start_match: - split_idx = param_start_match.start() - result["negative_prompt"] = after[:split_idx].strip() - remaining = after[split_idx:].strip() - else: - # Fallback: if no parameters found, assume remaining is all negative prompt - # (unless it's empty) - result["negative_prompt"] = after.strip() - remaining = "" - else: - # No negative prompt marker. - # Split positive prompt from params directly. - param_start_match = re.search(r"\n(?:Steps|Size|Model|Seed|CFG|Sampler|Denoising|Ens|Version):", text) - if param_start_match: - split_idx = param_start_match.start() - result["prompt"] = text[:split_idx].strip() - remaining = text[split_idx:].strip() - else: - steps_idx = text.find("\nSteps:") - if steps_idx == -1: - if text.startswith("Steps:"): - result["prompt"] = "" - remaining = text - else: - result["prompt"] = text - remaining = "" - else: - result["prompt"] = text[:steps_idx].strip() - remaining = text[steps_idx:].lstrip() - + remaining = _split_auto1111_prompt_and_remaining(text, result) if remaining: - matches = list(_AUTO1111_KV_RE.finditer(remaining)) - for i, match in enumerate(matches): - key = match.group(1).strip().lower().replace(" ", "_") - value_start = match.end() - value_end = matches[i + 1].start() if (i + 1) < len(matches) else len(remaining) - value = remaining[value_start:value_end].strip().strip(",").strip() - - if not key: - continue - - if key == "steps": - try: - result["steps"] = int(value) - except (ValueError, TypeError): - pass - elif key == "sampler": - result["sampler"] = value - elif key in ("cfg_scale", "cfg"): - try: - result["cfg"] = float(value) - except (ValueError, TypeError): - pass - elif key == "seed": - try: - result["seed"] = int(value) - except (ValueError, TypeError): - pass - elif key in ("size", "hires_resize"): - if "x" in value: - w, h = value.split("x", 1) - try: - result["width"] = int(w.strip()) - result["height"] = int(h.strip()) - except (ValueError, TypeError): - pass - elif key == "model": - result["model"] = value + _parse_auto1111_kv_block(remaining, result) return result or None except Exception as e: logger.debug(f"Failed to parse Auto1111 params: {e}") return None + + +def _split_auto1111_prompt_and_remaining(text: str, result: Dict[str, Any]) -> str: + neg_marker = "Negative prompt:" + neg_idx = text.find(neg_marker) + if neg_idx != -1: + result["prompt"] = text[:neg_idx].strip() + after = text[neg_idx + len(neg_marker) :].lstrip() + split_idx = _find_auto1111_params_start(after) + if split_idx is not None: + result["negative_prompt"] = after[:split_idx].strip() + return after[split_idx:].strip() + result["negative_prompt"] = after.strip() + return "" + split_idx = _find_auto1111_params_start(text) + if split_idx is not None: + result["prompt"] = text[:split_idx].strip() + return text[split_idx:].strip() + return _fallback_split_prompt_and_remaining(text, result) + + +def _find_auto1111_params_start(text: str) -> Optional[int]: + match = re.search(r"\n(?:Steps|Size|Model|Seed|CFG|Sampler|Denoising|Ens|Version):", text) + if not match: + return None + return match.start() + + +def _fallback_split_prompt_and_remaining(text: str, result: Dict[str, Any]) -> str: + steps_idx = text.find("\nSteps:") + if steps_idx == -1: + if text.startswith("Steps:"): + result["prompt"] = "" + return text + result["prompt"] = text + return "" + result["prompt"] = text[:steps_idx].strip() + return text[steps_idx:].lstrip() + + +def _parse_auto1111_kv_block(remaining: str, result: Dict[str, Any]) -> None: + matches = list(_AUTO1111_KV_RE.finditer(remaining)) + for i, match in enumerate(matches): + key = match.group(1).strip().lower().replace(" ", "_") + if not key: + continue + value = _extract_auto1111_kv_value(remaining, matches, i, match) + _apply_auto1111_kv(result, key, value) + + +def _extract_auto1111_kv_value(remaining: str, matches: list[Any], idx: int, match: Any) -> str: + value_start = match.end() + value_end = matches[idx + 1].start() if (idx + 1) < len(matches) else len(remaining) + return remaining[value_start:value_end].strip().strip(",").strip() + + +def _apply_auto1111_kv(result: Dict[str, Any], key: str, value: str) -> None: + if key == "steps": + _set_int_kv(result, "steps", value) + return + if key == "sampler": + result["sampler"] = value + return + if key in ("cfg_scale", "cfg"): + _set_float_kv(result, "cfg", value) + return + if key == "seed": + _set_int_kv(result, "seed", value) + return + if key in ("size", "hires_resize"): + _set_size_kv(result, value) + return + if key == "model": + result["model"] = value + + +def _set_int_kv(result: Dict[str, Any], key: str, value: str) -> None: + try: + result[key] = int(value) + except (ValueError, TypeError): + return + + +def _set_float_kv(result: Dict[str, Any], key: str, value: str) -> None: + try: + result[key] = float(value) + except (ValueError, TypeError): + return + + +def _set_size_kv(result: Dict[str, Any], value: str) -> None: + if "x" not in value: + return + w, h = value.split("x", 1) + try: + result["width"] = int(w.strip()) + result["height"] = int(h.strip()) + except (ValueError, TypeError): + return diff --git a/mjr_am_backend/features/metadata/service.py b/mjr_am_backend/features/metadata/service.py index 515f4207..cdd391eb 100644 --- a/mjr_am_backend/features/metadata/service.py +++ b/mjr_am_backend/features/metadata/service.py @@ -64,23 +64,7 @@ def _build_geninfo_from_parameters(meta: Dict[str, Any]) -> Optional[Dict[str, A if not isinstance(meta, dict): return None - parsed_meta = dict(meta) - params_text = parsed_meta.get("parameters") - if isinstance(params_text, str) and params_text.strip(): - parsed_params = parse_auto1111_params(params_text) - if parsed_params: - if logger.isEnabledFor(logging.DEBUG): - logger.debug( - "Parsed PNG:Parameters for geninfo fallback (%s keys)", - ", ".join(sorted(parsed_params.keys())), - ) - for key, value in parsed_params.items(): - if value is None: - continue - parsed_meta[key] = value - elif logger.isEnabledFor(logging.DEBUG): - logger.debug("Auto1111 parameter parser returned no fields for provided PNG:Parameters text.") - + parsed_meta = _merge_parsed_params(meta) pos = parsed_meta.get("prompt") neg = parsed_meta.get("negative_prompt") steps = parsed_meta.get("steps") @@ -92,51 +76,91 @@ def _build_geninfo_from_parameters(meta: Dict[str, Any]) -> Optional[Dict[str, A h = parsed_meta.get("height") model = parsed_meta.get("model") - has_any = any(v is not None and v != "" for v in (pos, neg, steps, sampler, cfg, seed, w, h, model)) + has_any = _has_any_parameter_signal(pos, neg, steps, sampler, cfg, seed, w, h, model) if not has_any: # Return an empty dict instead of None to ensure geninfo is always a dict return {} out: Dict[str, Any] = {"engine": {"parser_version": "geninfo-params-v1", "source": "parameters"}} + _apply_prompt_fields(out, pos, neg) + _apply_sampler_fields(out, sampler, scheduler) + _apply_numeric_fields(out, steps, cfg, seed) + _apply_size_field(out, w, h) + _apply_checkpoint_fields(out, model) + + return out if len(out.keys()) > 1 else None + + +def _merge_parsed_params(meta: Dict[str, Any]) -> Dict[str, Any]: + parsed_meta = dict(meta) + params_text = parsed_meta.get("parameters") + if not isinstance(params_text, str) or not params_text.strip(): + return parsed_meta + parsed_params = parse_auto1111_params(params_text) + if parsed_params: + if logger.isEnabledFor(logging.DEBUG): + logger.debug( + "Parsed PNG:Parameters for geninfo fallback (%s keys)", + ", ".join(sorted(parsed_params.keys())), + ) + for key, value in parsed_params.items(): + if value is not None: + parsed_meta[key] = value + return parsed_meta + if logger.isEnabledFor(logging.DEBUG): + logger.debug("Auto1111 parameter parser returned no fields for provided PNG:Parameters text.") + return parsed_meta + + +def _has_any_parameter_signal(*values: Any) -> bool: + return any(value is not None and value != "" for value in values) + + +def _apply_prompt_fields(out: Dict[str, Any], pos: Any, neg: Any) -> None: if isinstance(pos, str) and pos.strip(): out["positive"] = {"value": pos.strip(), "confidence": "high", "source": "parameters"} if isinstance(neg, str) and neg.strip(): out["negative"] = {"value": neg.strip(), "confidence": "high", "source": "parameters"} + +def _apply_sampler_fields(out: Dict[str, Any], sampler: Any, scheduler: Any) -> None: if sampler is not None: out["sampler"] = {"name": str(sampler), "confidence": "high", "source": "parameters"} if scheduler is not None: out["scheduler"] = {"name": str(scheduler), "confidence": "high", "source": "parameters"} - if steps is not None: - try: - out["steps"] = {"value": int(steps), "confidence": "high", "source": "parameters"} - except Exception: - pass - if cfg is not None: - try: - out["cfg"] = {"value": float(cfg), "confidence": "high", "source": "parameters"} - except Exception: - pass - if seed is not None: - try: - out["seed"] = {"value": int(seed), "confidence": "high", "source": "parameters"} - except Exception: - pass - if w is not None and h is not None: - try: - out["size"] = {"width": int(w), "height": int(h), "confidence": "high", "source": "parameters"} - except Exception: - pass +def _apply_numeric_fields(out: Dict[str, Any], steps: Any, cfg: Any, seed: Any) -> None: + _set_numeric_field(out, "steps", steps, int) + _set_numeric_field(out, "cfg", cfg, float) + _set_numeric_field(out, "seed", seed, int) + +def _set_numeric_field(out: Dict[str, Any], key: str, value: Any, caster: Callable[[Any], Any]) -> None: + if value is None: + return + try: + out[key] = {"value": caster(value), "confidence": "high", "source": "parameters"} + except Exception: + return + + +def _apply_size_field(out: Dict[str, Any], width: Any, height: Any) -> None: + if width is None or height is None: + return + try: + out["size"] = {"width": int(width), "height": int(height), "confidence": "high", "source": "parameters"} + except Exception: + return + + +def _apply_checkpoint_fields(out: Dict[str, Any], model: Any) -> None: ckpt = _clean_model_name(model) - if ckpt: - out["checkpoint"] = {"name": ckpt, "confidence": "high", "source": "parameters"} - out["models"] = {"checkpoint": out["checkpoint"]} - - return out if len(out.keys()) > 1 else None + if not ckpt: + return + out["checkpoint"] = {"name": ckpt, "confidence": "high", "source": "parameters"} + out["models"] = {"checkpoint": out["checkpoint"]} def _looks_like_media_pipeline(prompt_graph: Any) -> bool: @@ -148,30 +172,46 @@ def _looks_like_media_pipeline(prompt_graph: Any) -> bool: if not isinstance(prompt_graph, dict) or not prompt_graph: return False try: - types = [] - for node in prompt_graph.values(): - if not isinstance(node, dict): - continue - ct = str(node.get("class_type") or node.get("type") or "").lower() - if ct: - types.append(ct) + types = _collect_prompt_graph_types(prompt_graph) if not types: return False - # If any sampler-like node exists, it's not media-only. - if any(("ksampler" in t and "select" not in t) or ("samplercustom" in t) for t in types): - return False - if any("sampler" in t and "select" not in t for t in types): + if _has_generation_sampler(types): return False - has_load = any("loadvideo" in t or "vhs_loadvideo" in t for t in types) - has_combine = any("videocombine" in t or "video_combine" in t or "vhs_videocombine" in t for t in types) - has_save = any(t.startswith("save") or "savevideo" in t or "savegif" in t or "saveanimatedwebp" in t for t in types) + has_load, has_combine, has_save = _classify_media_nodes(types) return has_load and (has_combine or has_save) except Exception: return False +def _collect_prompt_graph_types(prompt_graph: Dict[str, Any]) -> list[str]: + types: list[str] = [] + for node in prompt_graph.values(): + if not isinstance(node, dict): + continue + ct = str(node.get("class_type") or node.get("type") or "").lower() + if ct: + types.append(ct) + return types + + +def _has_generation_sampler(types: list[str]) -> bool: + if any(("ksampler" in token and "select" not in token) or ("samplercustom" in token) for token in types): + return True + return any("sampler" in token and "select" not in token for token in types) + + +def _classify_media_nodes(types: list[str]) -> tuple[bool, bool, bool]: + has_load = any("loadvideo" in token or "vhs_loadvideo" in token for token in types) + has_combine = any("videocombine" in token or "video_combine" in token or "vhs_videocombine" in token for token in types) + has_save = any( + token.startswith("save") or "savevideo" in token or "savegif" in token or "saveanimatedwebp" in token + for token in types + ) + return has_load, has_combine, has_save + + def _should_parse_geninfo(meta: Dict[str, Any]) -> bool: """ Fast gate to avoid expensive geninfo parsing when generation signals are absent. @@ -192,6 +232,81 @@ def _should_parse_geninfo(meta: Dict[str, Any]) -> bool: return False return False + +def _group_existing_paths(paths: list[str]) -> tuple[list[str], list[str], list[str], list[str]]: + images: list[str] = [] + videos: list[str] = [] + audios: list[str] = [] + others: list[str] = [] + for path in paths: + if not os.path.exists(path): + continue + kind = classify_file(path) + if kind == "image": + images.append(path) + elif kind == "video": + videos.append(path) + elif kind == "audio": + audios.append(path) + else: + others.append(path) + return images, videos, audios, others + + +def _build_batch_probe_targets(paths: list[str], probe_mode: str) -> tuple[list[str], list[str]]: + exif_targets: list[str] = [] + ffprobe_targets: list[str] = [] + seen_exif: set[str] = set() + seen_ffprobe: set[str] = set() + for path in paths: + backends = pick_probe_backend(path, settings_override=probe_mode) + if "exiftool" in backends and path not in seen_exif: + seen_exif.add(path) + exif_targets.append(path) + if "ffprobe" in backends and path not in seen_ffprobe: + seen_ffprobe.add(path) + ffprobe_targets.append(path) + return exif_targets, ffprobe_targets + + +def _batch_tool_data( + result_map: Dict[str, Result[Dict[str, Any]]], + path: str, +) -> Optional[Dict[str, Any]]: + item = result_map.get(path) + return item.data if item and item.ok else None + + +def _expand_resolution_scalars(payload: Dict[str, Any]) -> None: + resolution = payload.get("resolution") + if not resolution: + return + if payload.get("width") is not None and payload.get("height") is not None: + return + dims = _coerce_resolution_pair(resolution) + if dims is None: + return + w, h = dims + if payload.get("width") is None: + payload["width"] = _coerce_dimension_value(w) + if payload.get("height") is None: + payload["height"] = _coerce_dimension_value(h) + + +def _coerce_resolution_pair(resolution: Any) -> Optional[tuple[Any, Any]]: + try: + w, h = resolution + return w, h + except Exception: + return None + + +def _coerce_dimension_value(value: Any) -> Any: + try: + return int(value) if value is not None else None + except Exception: + return value + class MetadataService: """ Metadata extraction service. @@ -433,39 +548,16 @@ class MetadataService: return Result.Ok({"workflow": None, "prompt": None, "quality": "none"}, quality="none") ext = os.path.splitext(file_path)[1].lower() - - exif_start = time.perf_counter() - exif_result = await self._exif_read(file_path) - exif_duration = time.perf_counter() - exif_start - exif_data = exif_result.data if exif_result.ok else None image_fallback_enabled, _ = await self._resolve_fallback_prefs() - if not exif_result.ok: - self._log_metadata_issue( - logging.DEBUG, - "ExifTool failed while reading workflow-only metadata", - file_path, - scan_id=scan_id, - tool="exiftool", - error=exif_result.error, - duration_seconds=exif_duration, - ) - if kind == "image" and image_fallback_enabled: - exif_data = await asyncio.to_thread(read_image_exif_like, file_path) - else: - exif_data = {} + exif_data = await self._read_workflow_only_exif( + file_path=file_path, + kind=kind, + image_fallback_enabled=image_fallback_enabled, + scan_id=scan_id, + ) if not exif_data: return Result.Ok({"workflow": None, "prompt": None, "quality": "none"}, quality="none") - - if kind == "image": - if ext == ".png": - res = extract_png_metadata(file_path, exif_data) - elif ext == ".webp": - res = extract_webp_metadata(file_path, exif_data) - else: - res = Result.Ok({"workflow": None, "prompt": None, "quality": "none"}, quality="none") - else: - # No ffprobe here (can be slow); extractor supports workflow/prompt from ExifTool tags. - res = extract_video_metadata(file_path, exif_data, ffprobe_data=None) + res = self._extract_workflow_only_payload(kind, ext, file_path, exif_data) if not res.ok: return res @@ -478,6 +570,50 @@ class MetadataService: } return Result.Ok(payload, quality=payload.get("quality", "none")) + async def _read_workflow_only_exif( + self, + *, + file_path: str, + kind: str, + image_fallback_enabled: bool, + scan_id: Optional[str], + ) -> Dict[str, Any]: + exif_start = time.perf_counter() + exif_result = await self._exif_read(file_path) + exif_duration = time.perf_counter() - exif_start + exif_data = exif_result.data if exif_result.ok else None + if exif_result.ok: + return exif_data or {} + self._log_metadata_issue( + logging.DEBUG, + "ExifTool failed while reading workflow-only metadata", + file_path, + scan_id=scan_id, + tool="exiftool", + error=exif_result.error, + duration_seconds=exif_duration, + ) + if kind == "image" and image_fallback_enabled: + fallback = await asyncio.to_thread(read_image_exif_like, file_path) + return fallback or {} + return {} + + @staticmethod + def _extract_workflow_only_payload( + kind: str, + ext: str, + file_path: str, + exif_data: Dict[str, Any], + ) -> Result[Dict[str, Any]]: + if kind == "image": + if ext == ".png": + return extract_png_metadata(file_path, exif_data) + if ext == ".webp": + return extract_webp_metadata(file_path, exif_data) + return Result.Ok({"workflow": None, "prompt": None, "quality": "none"}, quality="none") + # No ffprobe here (can be slow); extractor supports workflow/prompt from ExifTool tags. + return extract_video_metadata(file_path, exif_data, ffprobe_data=None) + async def _extract_image_metadata( self, file_path: str, @@ -486,40 +622,8 @@ class MetadataService: ) -> Result[Dict[str, Any]]: """Extract metadata from image file.""" ext = os.path.splitext(file_path)[1].lower() - - exif_data = None - image_fallback_enabled, _ = await self._resolve_fallback_prefs() - if allow_exif: - exif_start = time.perf_counter() - exif_result = await self._exif_read(file_path) - exif_duration = time.perf_counter() - exif_start - exif_data = exif_result.data if exif_result.ok else None - if not exif_result.ok: - self._log_metadata_issue( - logging.WARNING, - "ExifTool failed while reading image metadata", - file_path, - scan_id=scan_id, - tool="exiftool", - error=exif_result.error, - duration_seconds=exif_duration - ) - if not exif_data and image_fallback_enabled: - exif_data = await asyncio.to_thread(read_image_exif_like, file_path) - - # Extract based on format - if ext == ".png": - metadata_result = extract_png_metadata(file_path, exif_data) - elif ext == ".webp": - metadata_result = extract_webp_metadata(file_path, exif_data) - else: - # Generic image - metadata_result = Result.Ok({ - "workflow": None, - "prompt": None, - "quality": "partial" if exif_data else "none" - }) - + exif_data = await self._resolve_image_exif_data(file_path, scan_id=scan_id, allow_exif=allow_exif) + metadata_result = self._extract_image_by_extension(file_path, ext, exif_data) if not metadata_result.ok: self._log_metadata_issue( logging.WARNING, @@ -528,16 +632,11 @@ class MetadataService: scan_id=scan_id, tool="extractor", error=metadata_result.error, - duration_seconds=exif_duration + duration_seconds=0.0, ) return metadata_result - # Combine results - combined = { - "file_info": self._get_file_info(file_path), - "exif": exif_data, - **(metadata_result.data or {}) - } + combined = self._build_image_metadata_payload(file_path, exif_data, metadata_result) self._normalize_visual_dimensions(combined, exif_data=exif_data, ffprobe_data=None) # Optional: parse deterministic generation info from ComfyUI prompt graph (backend-first). @@ -547,6 +646,69 @@ class MetadataService: quality = metadata_result.meta.get("quality", "none") return Result.Ok(combined, quality=quality) + async def _resolve_image_exif_data( + self, + file_path: str, + *, + scan_id: Optional[str], + allow_exif: bool, + ) -> Optional[Dict[str, Any]]: + exif_data: Optional[Dict[str, Any]] = None + image_fallback_enabled, _ = await self._resolve_fallback_prefs() + if allow_exif: + exif_data = await self._read_image_exif_if_allowed(file_path, scan_id=scan_id) + if not exif_data and image_fallback_enabled: + exif_data = await asyncio.to_thread(read_image_exif_like, file_path) + return exif_data + + async def _read_image_exif_if_allowed( + self, + file_path: str, + *, + scan_id: Optional[str], + ) -> Optional[Dict[str, Any]]: + exif_start = time.perf_counter() + exif_result = await self._exif_read(file_path) + exif_duration = time.perf_counter() - exif_start + if exif_result.ok: + return exif_result.data if isinstance(exif_result.data, dict) else None + self._log_metadata_issue( + logging.WARNING, + "ExifTool failed while reading image metadata", + file_path, + scan_id=scan_id, + tool="exiftool", + error=exif_result.error, + duration_seconds=exif_duration, + ) + return None + + @staticmethod + def _extract_image_by_extension(file_path: str, ext: str, exif_data: Optional[Dict[str, Any]]) -> Result[Dict[str, Any]]: + if ext == ".png": + return extract_png_metadata(file_path, exif_data) + if ext == ".webp": + return extract_webp_metadata(file_path, exif_data) + return Result.Ok( + { + "workflow": None, + "prompt": None, + "quality": "partial" if exif_data else "none", + } + ) + + def _build_image_metadata_payload( + self, + file_path: str, + exif_data: Optional[Dict[str, Any]], + metadata_result: Result[Dict[str, Any]], + ) -> Dict[str, Any]: + return { + "file_info": self._get_file_info(file_path), + "exif": exif_data, + **(metadata_result.data or {}), + } + async def _extract_video_metadata( self, file_path: str, @@ -555,47 +717,9 @@ class MetadataService: allow_ffprobe: bool = True, ) -> Result[Dict[str, Any]]: """Extract metadata from video file.""" - exif_data: Dict[str, Any] | None = {} - exif_duration = 0.0 _, media_fallback_enabled = await self._resolve_fallback_prefs() - if allow_exif: - exif_start = time.perf_counter() - exif_result = await self._exif_read(file_path) - exif_duration = time.perf_counter() - exif_start - exif_data = exif_result.data if exif_result.ok else None - if not exif_result.ok: - self._log_metadata_issue( - logging.WARNING, - "ExifTool failed while reading video metadata", - file_path, - scan_id=scan_id, - tool="exiftool", - error=exif_result.error, - duration_seconds=exif_duration - ) - exif_data = None - else: - exif_data = {} - - ffprobe_data: Dict[str, Any] | None = {} - ffprobe_duration = 0.0 - if allow_ffprobe: - ffprobe_start = time.perf_counter() - ffprobe_result = await self._ffprobe_read(file_path) - ffprobe_duration = time.perf_counter() - ffprobe_start - if not ffprobe_result.ok: - self._log_metadata_issue( - logging.WARNING, - "FFprobe failed while reading video metadata", - file_path, - scan_id=scan_id, - tool="ffprobe", - error=ffprobe_result.error, - duration_seconds=ffprobe_duration - ) - ffprobe_data = None - else: - ffprobe_data = ffprobe_result.data or {} + exif_data, exif_duration = await self._read_video_exif_if_allowed(file_path, scan_id, allow_exif) + ffprobe_data, ffprobe_duration = await self._read_video_ffprobe_if_allowed(file_path, scan_id, allow_ffprobe) if not ffprobe_data and media_fallback_enabled: ffprobe_data = await asyncio.to_thread(read_media_probe_like, file_path) @@ -614,15 +738,7 @@ class MetadataService: ) return metadata_result - # Unpack resolution tuple to width/height for DB compatibility - data = metadata_result.data or {} - if data.get("resolution"): - try: - w, h = data["resolution"] - data["width"] = w - data["height"] = h - except (ValueError, TypeError): - pass + self._expand_video_resolution_fields(metadata_result) # Combine results combined = { @@ -650,63 +766,11 @@ class MetadataService: ) -> Result[Dict[str, Any]]: """Extract metadata from audio file.""" if not allow_exif and not allow_ffprobe: - _, media_fallback_enabled = await self._resolve_fallback_prefs() - fallback_ffprobe = ( - await asyncio.to_thread(read_media_probe_like, file_path) - if media_fallback_enabled - else None - ) - metadata_result = extract_audio_metadata(file_path, exif_data=None, ffprobe_data=fallback_ffprobe or None) - if not metadata_result.ok: - return metadata_result - combined = { - "file_info": self._get_file_info(file_path), - "exif": None, - "ffprobe": fallback_ffprobe or None, - **(metadata_result.data or {}), - } - quality = metadata_result.meta.get("quality", "none") if metadata_result.ok else "none" - return Result.Ok(combined, quality=quality) + return await self._extract_audio_metadata_fallback_only(file_path) - exif_data = None - ffprobe_data = None - exif_duration = 0.0 - ffprobe_duration = 0.0 - - if allow_exif: - exif_start = time.perf_counter() - exif_result = await self._exif_read(file_path) - exif_duration = time.perf_counter() - exif_start - exif_data = exif_result.data if exif_result.ok else None - if not exif_result.ok: - self._log_metadata_issue( - logging.WARNING, - "ExifTool failed while reading audio metadata", - file_path, - scan_id=scan_id, - tool="exiftool", - error=exif_result.error, - duration_seconds=exif_duration, - ) - - if allow_ffprobe: - ffprobe_start = time.perf_counter() - ffprobe_result = await self._ffprobe_read(file_path) - ffprobe_duration = time.perf_counter() - ffprobe_start - ffprobe_data = ffprobe_result.data if ffprobe_result.ok else None - if not ffprobe_result.ok: - self._log_metadata_issue( - logging.WARNING, - "FFprobe failed while reading audio metadata", - file_path, - scan_id=scan_id, - tool="ffprobe", - error=ffprobe_result.error, - duration_seconds=ffprobe_duration, - ) - _, media_fallback_enabled = await self._resolve_fallback_prefs() - if not ffprobe_data and media_fallback_enabled: - ffprobe_data = await asyncio.to_thread(read_media_probe_like, file_path) + exif_data, exif_duration = await self._read_audio_exif_if_allowed(file_path, scan_id, allow_exif) + ffprobe_data, ffprobe_duration = await self._read_audio_ffprobe_if_allowed(file_path, scan_id, allow_ffprobe) + ffprobe_data = await self._maybe_audio_ffprobe_fallback(file_path, ffprobe_data) metadata_result = extract_audio_metadata(file_path, exif_data=exif_data, ffprobe_data=ffprobe_data) if not metadata_result.ok: @@ -721,6 +785,144 @@ class MetadataService: ) return metadata_result + return await self._finalize_audio_metadata_ok(file_path, exif_data, ffprobe_data, metadata_result) + + async def _extract_audio_metadata_fallback_only(self, file_path: str) -> Result[Dict[str, Any]]: + _, media_fallback_enabled = await self._resolve_fallback_prefs() + fallback_ffprobe = await asyncio.to_thread(read_media_probe_like, file_path) if media_fallback_enabled else None + metadata_result = extract_audio_metadata(file_path, exif_data=None, ffprobe_data=fallback_ffprobe or None) + if not metadata_result.ok: + return metadata_result + combined = { + "file_info": self._get_file_info(file_path), + "exif": None, + "ffprobe": fallback_ffprobe or None, + **(metadata_result.data or {}), + } + quality = metadata_result.meta.get("quality", "none") + return Result.Ok(combined, quality=quality) + + async def _read_video_exif_if_allowed( + self, file_path: str, scan_id: Optional[str], allow_exif: bool + ) -> tuple[Dict[str, Any] | None, float]: + if not allow_exif: + return {}, 0.0 + exif_start = time.perf_counter() + exif_result = await self._exif_read(file_path) + exif_duration = time.perf_counter() - exif_start + if exif_result.ok: + return exif_result.data or {}, exif_duration + self._log_metadata_issue( + logging.WARNING, + "ExifTool failed while reading video metadata", + file_path, + scan_id=scan_id, + tool="exiftool", + error=exif_result.error, + duration_seconds=exif_duration, + ) + return None, exif_duration + + async def _read_video_ffprobe_if_allowed( + self, file_path: str, scan_id: Optional[str], allow_ffprobe: bool + ) -> tuple[Dict[str, Any] | None, float]: + if not allow_ffprobe: + return {}, 0.0 + ffprobe_start = time.perf_counter() + ffprobe_result = await self._ffprobe_read(file_path) + ffprobe_duration = time.perf_counter() - ffprobe_start + if ffprobe_result.ok: + return ffprobe_result.data or {}, ffprobe_duration + self._log_metadata_issue( + logging.WARNING, + "FFprobe failed while reading video metadata", + file_path, + scan_id=scan_id, + tool="ffprobe", + error=ffprobe_result.error, + duration_seconds=ffprobe_duration, + ) + return None, ffprobe_duration + + @staticmethod + def _expand_video_resolution_fields(metadata_result: Result[Dict[str, Any]]) -> None: + data = metadata_result.data or {} + if not data.get("resolution"): + return + try: + w, h = data["resolution"] + data["width"] = w + data["height"] = h + except (ValueError, TypeError): + return + + async def _read_audio_exif_if_allowed( + self, + file_path: str, + scan_id: Optional[str], + allow_exif: bool, + ) -> tuple[Optional[Dict[str, Any]], float]: + if not allow_exif: + return None, 0.0 + exif_start = time.perf_counter() + exif_result = await self._exif_read(file_path) + exif_duration = time.perf_counter() - exif_start + exif_data = exif_result.data if exif_result.ok else None + if not exif_result.ok: + self._log_metadata_issue( + logging.WARNING, + "ExifTool failed while reading audio metadata", + file_path, + scan_id=scan_id, + tool="exiftool", + error=exif_result.error, + duration_seconds=exif_duration, + ) + return exif_data, exif_duration + + async def _read_audio_ffprobe_if_allowed( + self, + file_path: str, + scan_id: Optional[str], + allow_ffprobe: bool, + ) -> tuple[Optional[Dict[str, Any]], float]: + if not allow_ffprobe: + return None, 0.0 + ffprobe_start = time.perf_counter() + ffprobe_result = await self._ffprobe_read(file_path) + ffprobe_duration = time.perf_counter() - ffprobe_start + ffprobe_data = ffprobe_result.data if ffprobe_result.ok else None + if not ffprobe_result.ok: + self._log_metadata_issue( + logging.WARNING, + "FFprobe failed while reading audio metadata", + file_path, + scan_id=scan_id, + tool="ffprobe", + error=ffprobe_result.error, + duration_seconds=ffprobe_duration, + ) + return ffprobe_data, ffprobe_duration + + async def _maybe_audio_ffprobe_fallback( + self, + file_path: str, + ffprobe_data: Optional[Dict[str, Any]], + ) -> Optional[Dict[str, Any]]: + if ffprobe_data: + return ffprobe_data + _, media_fallback_enabled = await self._resolve_fallback_prefs() + if not media_fallback_enabled: + return ffprobe_data + return await asyncio.to_thread(read_media_probe_like, file_path) + + async def _finalize_audio_metadata_ok( + self, + file_path: str, + exif_data: Optional[Dict[str, Any]], + ffprobe_data: Optional[Dict[str, Any]], + metadata_result: Result[Dict[str, Any]], + ) -> Result[Dict[str, Any]]: combined = { "file_info": self._get_file_info(file_path), "exif": exif_data, @@ -763,6 +965,103 @@ class MetadataService: probe_mode_override=probe_mode_override, ) + async def _finalize_batch_ok( + self, + combined: Dict[str, Any], + metadata_result: Result[Dict[str, Any]], + ) -> Result[Dict[str, Any]]: + if _should_parse_geninfo(combined): + await self._enrich_with_geninfo_async(combined) + elif "geninfo" not in combined: + combined["geninfo"] = {} + quality = metadata_result.meta.get("quality", "none") + return Result.Ok(combined, quality=quality) + + async def _process_image_batch_item( + self, + path: str, + exif_data: Optional[Dict[str, Any]], + image_fallback_enabled: bool, + ) -> Result[Dict[str, Any]]: + resolved_exif = exif_data + if not resolved_exif and image_fallback_enabled: + resolved_exif = read_image_exif_like(path) + + ext = os.path.splitext(path)[1].lower() + if ext == ".png": + metadata_result = extract_png_metadata(path, resolved_exif) + elif ext == ".webp": + metadata_result = extract_webp_metadata(path, resolved_exif) + else: + metadata_result = Result.Ok( + { + "workflow": None, + "prompt": None, + "quality": "partial" if resolved_exif else "none", + } + ) + + if not metadata_result.ok: + return metadata_result + + combined = { + "file_info": self._get_file_info(path), + "exif": resolved_exif, + **(metadata_result.data or {}), + } + self._normalize_visual_dimensions(combined, exif_data=resolved_exif, ffprobe_data=None) + return await self._finalize_batch_ok(combined, metadata_result) + + async def _process_video_batch_item( + self, + path: str, + exif_data: Optional[Dict[str, Any]], + ffprobe_data: Optional[Dict[str, Any]], + media_fallback_enabled: bool, + ) -> Result[Dict[str, Any]]: + resolved_ffprobe = ffprobe_data + if not resolved_ffprobe and media_fallback_enabled: + resolved_ffprobe = read_media_probe_like(path) + + metadata_result = extract_video_metadata(path, exif_data, resolved_ffprobe) + if not metadata_result.ok: + return metadata_result + + payload = dict(metadata_result.data or {}) + _expand_resolution_scalars(payload) + + combined = { + "file_info": self._get_file_info(path), + "exif": exif_data, + "ffprobe": resolved_ffprobe, + **payload, + } + self._normalize_visual_dimensions(combined, exif_data=exif_data, ffprobe_data=resolved_ffprobe) + return await self._finalize_batch_ok(combined, metadata_result) + + async def _process_audio_batch_item( + self, + path: str, + exif_data: Optional[Dict[str, Any]], + ffprobe_data: Optional[Dict[str, Any]], + media_fallback_enabled: bool, + ) -> Result[Dict[str, Any]]: + resolved_ffprobe = ffprobe_data + if not resolved_ffprobe and media_fallback_enabled: + resolved_ffprobe = read_media_probe_like(path) + + metadata_result = extract_audio_metadata(path, exif_data=exif_data, ffprobe_data=resolved_ffprobe) + if not metadata_result.ok: + return metadata_result + + combined = { + "file_info": self._get_file_info(path), + "exif": exif_data, + "ffprobe": resolved_ffprobe, + **(metadata_result.data or {}), + } + return await self._finalize_batch_ok(combined, metadata_result) + async def _get_metadata_batch_impl( self, file_paths: list[str], @@ -772,47 +1071,12 @@ class MetadataService: if not file_paths: return {} - from ...shared import classify_file - - def _group_existing_paths(paths: list[str]): - images: list[str] = [] - videos: list[str] = [] - audios: list[str] = [] - others: list[str] = [] - for path in paths: - if not os.path.exists(path): - continue - kind = classify_file(path) - if kind == "image": - images.append(path) - elif kind == "video": - videos.append(path) - elif kind == "audio": - audios.append(path) - else: - others.append(path) - return images, videos, audios, others - images, videos, audios, others = _group_existing_paths(file_paths) - results: Dict[str, Result[Dict[str, Any]]] = {} probe_mode = await self._resolve_probe_mode(probe_mode_override) image_fallback_enabled, media_fallback_enabled = await self._resolve_fallback_prefs() - # Schedule tool targets (dedupe to avoid duplicate subprocess work). - exif_targets: list[str] = [] - ffprobe_targets: list[str] = [] - seen_exif: set[str] = set() - seen_ffprobe: set[str] = set() - - for path in [*images, *videos, *audios]: - backends = pick_probe_backend(path, settings_override=probe_mode) - if "exiftool" in backends and path not in seen_exif: - seen_exif.add(path) - exif_targets.append(path) - if "ffprobe" in backends and path not in seen_ffprobe: - seen_ffprobe.add(path) - ffprobe_targets.append(path) + exif_targets, ffprobe_targets = _build_batch_probe_targets([*images, *videos, *audios], probe_mode) exif_results: Dict[str, Result[Dict[str, Any]]] = ( await self.exiftool.aread_batch(exif_targets) if exif_targets else {} @@ -821,115 +1085,75 @@ class MetadataService: await self._ffprobe_read_batch(ffprobe_targets) if ffprobe_targets else {} ) - def _exif_for(path: str) -> Optional[Dict[str, Any]]: - ex_res = exif_results.get(path) - return ex_res.data if ex_res and ex_res.ok else None + await self._fill_image_batch_results(results, images, exif_results, image_fallback_enabled) + await self._fill_video_batch_results(results, videos, exif_results, ffprobe_results, media_fallback_enabled) + await self._fill_audio_batch_results(results, audios, exif_results, ffprobe_results, media_fallback_enabled) + self._fill_other_batch_results(results, others) - def _ffprobe_for(path: str) -> Optional[Dict[str, Any]]: - ff_res = ffprobe_results.get(path) - return ff_res.data if ff_res and ff_res.ok else None + return results - async def _finalize_ok(path: str, combined: Dict[str, Any], metadata_result: Result[Dict[str, Any]]): - if _should_parse_geninfo(combined): - await self._enrich_with_geninfo_async(combined) - elif "geninfo" not in combined: - combined["geninfo"] = {} - quality = metadata_result.meta.get("quality", "none") - results[path] = Result.Ok(combined, quality=quality) - - # Images - for path in images: + async def _fill_image_batch_results( + self, + results: Dict[str, Result[Dict[str, Any]]], + paths: list[str], + exif_results: Dict[str, Result[Dict[str, Any]]], + image_fallback_enabled: bool, + ) -> None: + for path in paths: if path in results: continue - ext = os.path.splitext(path)[1].lower() - exif_data = _exif_for(path) - if not exif_data and image_fallback_enabled: - exif_data = read_image_exif_like(path) + results[path] = await self._process_image_batch_item( + path, + _batch_tool_data(exif_results, path), + image_fallback_enabled, + ) - if ext == ".png": - metadata_result = extract_png_metadata(path, exif_data) - elif ext == ".webp": - metadata_result = extract_webp_metadata(path, exif_data) - else: - metadata_result = Result.Ok({ - "workflow": None, - "prompt": None, - "quality": "partial" if exif_data else "none", - }) - - if metadata_result.ok: - combined = { - "file_info": self._get_file_info(path), - "exif": exif_data, - **(metadata_result.data or {}), - } - self._normalize_visual_dimensions(combined, exif_data=exif_data, ffprobe_data=None) - await _finalize_ok(path, combined, metadata_result) - else: - results[path] = metadata_result - - # Videos - for path in videos: + async def _fill_video_batch_results( + self, + results: Dict[str, Result[Dict[str, Any]]], + paths: list[str], + exif_results: Dict[str, Result[Dict[str, Any]]], + ffprobe_results: Dict[str, Result[Dict[str, Any]]], + media_fallback_enabled: bool, + ) -> None: + for path in paths: if path in results: continue - exif_data = _exif_for(path) - ffprobe_data = _ffprobe_for(path) - if not ffprobe_data and media_fallback_enabled: - ffprobe_data = read_media_probe_like(path) + results[path] = await self._process_video_batch_item( + path, + _batch_tool_data(exif_results, path), + _batch_tool_data(ffprobe_results, path), + media_fallback_enabled, + ) - metadata_result = extract_video_metadata(path, exif_data, ffprobe_data) - if metadata_result.ok: - payload = dict(metadata_result.data or {}) - # Keep DB compatibility: expose width/height scalars even when extractor - # returns only a resolution tuple. - if payload.get("resolution") and (payload.get("width") is None or payload.get("height") is None): - try: - w, h = payload.get("resolution") - if payload.get("width") is None: - payload["width"] = int(w) if w is not None else None - if payload.get("height") is None: - payload["height"] = int(h) if h is not None else None - except Exception: - pass - combined = { - "file_info": self._get_file_info(path), - "exif": exif_data, - "ffprobe": ffprobe_data, - **payload, - } - self._normalize_visual_dimensions(combined, exif_data=exif_data, ffprobe_data=ffprobe_data) - await _finalize_ok(path, combined, metadata_result) - else: - results[path] = metadata_result - - # Audio - for path in audios: + async def _fill_audio_batch_results( + self, + results: Dict[str, Result[Dict[str, Any]]], + paths: list[str], + exif_results: Dict[str, Result[Dict[str, Any]]], + ffprobe_results: Dict[str, Result[Dict[str, Any]]], + media_fallback_enabled: bool, + ) -> None: + for path in paths: if path in results: continue - exif_data = _exif_for(path) - ffprobe_data = _ffprobe_for(path) - if not ffprobe_data and media_fallback_enabled: - ffprobe_data = read_media_probe_like(path) - metadata_result = extract_audio_metadata(path, exif_data=exif_data, ffprobe_data=ffprobe_data) - if metadata_result.ok: - combined = { - "file_info": self._get_file_info(path), - "exif": exif_data, - "ffprobe": ffprobe_data, - **(metadata_result.data or {}), - } - await _finalize_ok(path, combined, metadata_result) - else: - results[path] = metadata_result + results[path] = await self._process_audio_batch_item( + path, + _batch_tool_data(exif_results, path), + _batch_tool_data(ffprobe_results, path), + media_fallback_enabled, + ) - # Other - for path in others: + def _fill_other_batch_results( + self, + results: Dict[str, Result[Dict[str, Any]]], + paths: list[str], + ) -> None: + for path in paths: if path in results: continue results[path] = Result.Ok({"file_info": self._get_file_info(path), "quality": "none"}, quality="none") - return results - def extract_rating_tags_only(self, file_path: str, scan_id: Optional[str] = None) -> Result[Dict[str, Any]]: """ Lightweight extraction for rating/tags only (no workflow/geninfo parsing). @@ -1039,41 +1263,97 @@ class MetadataService: w = _coerce(payload.get("width")) h = _coerce(payload.get("height")) - - if (w is None or h is None) and payload.get("resolution"): - try: - rw, rh = payload.get("resolution") - if w is None: - w = _coerce(rw) - if h is None: - h = _coerce(rh) - except Exception: - pass - - if (w is None or h is None) and isinstance(ffprobe_data, dict): - try: - video_stream = ffprobe_data.get("video_stream") or {} - if w is None: - w = _coerce(video_stream.get("width")) - if h is None: - h = _coerce(video_stream.get("height")) - except Exception: - pass - - if (w is None or h is None) and isinstance(exif_data, dict): - width_keys = ("Image:ImageWidth", "EXIF:ImageWidth", "EXIF:ExifImageWidth", "IFD0:ImageWidth", "Composite:ImageWidth", "QuickTime:ImageWidth", "PNG:ImageWidth", "File:ImageWidth", "width") - height_keys = ("Image:ImageHeight", "EXIF:ImageHeight", "EXIF:ExifImageHeight", "IFD0:ImageHeight", "Composite:ImageHeight", "QuickTime:ImageHeight", "PNG:ImageHeight", "File:ImageHeight", "height") - if w is None: - for k in width_keys: - w = _coerce(exif_data.get(k)) - if w is not None: - break - if h is None: - for k in height_keys: - h = _coerce(exif_data.get(k)) - if h is not None: - break + w, h = MetadataService._fill_dims_from_resolution(payload, w, h, _coerce) + w, h = MetadataService._fill_dims_from_ffprobe(ffprobe_data, w, h, _coerce) + w, h = MetadataService._fill_dims_from_exif(exif_data, w, h, _coerce) if w is not None and h is not None: payload["width"] = int(w) payload["height"] = int(h) + + @staticmethod + def _fill_dims_from_resolution( + payload: Dict[str, Any], + w: Optional[int], + h: Optional[int], + coerce: Callable[[Any], Optional[int]], + ) -> tuple[Optional[int], Optional[int]]: + if (w is not None and h is not None) or not payload.get("resolution"): + return w, h + try: + rw, rh = payload.get("resolution") + if w is None: + w = coerce(rw) + if h is None: + h = coerce(rh) + except Exception: + return w, h + return w, h + + @staticmethod + def _fill_dims_from_ffprobe( + ffprobe_data: Optional[Dict[str, Any]], + w: Optional[int], + h: Optional[int], + coerce: Callable[[Any], Optional[int]], + ) -> tuple[Optional[int], Optional[int]]: + if w is not None and h is not None: + return w, h + if not isinstance(ffprobe_data, dict): + return w, h + try: + video_stream = ffprobe_data.get("video_stream") or {} + if w is None: + w = coerce(video_stream.get("width")) + if h is None: + h = coerce(video_stream.get("height")) + except Exception: + return w, h + return w, h + + @staticmethod + def _fill_dims_from_exif( + exif_data: Optional[Dict[str, Any]], + w: Optional[int], + h: Optional[int], + coerce: Callable[[Any], Optional[int]], + ) -> tuple[Optional[int], Optional[int]]: + if w is not None and h is not None: + return w, h + if not isinstance(exif_data, dict): + return w, h + width_keys = ( + "Image:ImageWidth", + "EXIF:ImageWidth", + "EXIF:ExifImageWidth", + "IFD0:ImageWidth", + "Composite:ImageWidth", + "QuickTime:ImageWidth", + "PNG:ImageWidth", + "File:ImageWidth", + "width", + ) + height_keys = ( + "Image:ImageHeight", + "EXIF:ImageHeight", + "EXIF:ExifImageHeight", + "IFD0:ImageHeight", + "Composite:ImageHeight", + "QuickTime:ImageHeight", + "PNG:ImageHeight", + "File:ImageHeight", + "height", + ) + if w is None: + w = MetadataService._pick_first_coerced(exif_data, width_keys, coerce) + if h is None: + h = MetadataService._pick_first_coerced(exif_data, height_keys, coerce) + return w, h + + @staticmethod + def _pick_first_coerced(source: Dict[str, Any], keys: tuple[str, ...], coerce: Callable[[Any], Optional[int]]) -> Optional[int]: + for key in keys: + value = coerce(source.get(key)) + if value is not None: + return value + return None diff --git a/mjr_am_backend/features/tags/sync.py b/mjr_am_backend/features/tags/sync.py index 11fd3588..71f44954 100644 --- a/mjr_am_backend/features/tags/sync.py +++ b/mjr_am_backend/features/tags/sync.py @@ -59,17 +59,7 @@ def _normalize_tags(tags: List[str]) -> List[str]: return out -def write_exif_rating_tags(exiftool: ExifTool, file_path: str, rating: int, tags: List[str]) -> Result[bool]: - """ - Try writing rating/tags into the file metadata via ExifTool. - """ - try: - if not exiftool or not exiftool.is_available(): - return Result.Err(ErrorCode.TOOL_MISSING, "ExifTool not available") - except (AttributeError, TypeError, ValueError) as exc: - logger.debug("ExifTool availability check failed: %s", exc) - return Result.Err(ErrorCode.TOOL_MISSING, "ExifTool not available") - +def _validate_exiftool_file_path(file_path: str) -> Result[Path]: try: p = Path(str(file_path)) except (OSError, TypeError, ValueError) as exc: @@ -77,57 +67,75 @@ def write_exif_rating_tags(exiftool: ExifTool, file_path: str, rating: int, tags return Result.Err(ErrorCode.INVALID_INPUT, "Invalid file path") if not p.exists() or not p.is_file(): return Result.Err(ErrorCode.NOT_FOUND, f"File not found: {file_path}") + return Result.Ok(p) + + +def _build_exiftool_rating_tags_payload(stars: int, tags_norm: List[str]) -> Dict[str, Any]: + joined = "; ".join(tags_norm) + payload: Dict[str, Any] = { + "XMP:Rating": stars, + "xmp:rating": stars, + "rating": stars, + "ratingpercent": _windows_rating_percent(stars), + # Windows Property System (best-effort; writable depends on container/handler) + "Microsoft:SharedUserRating": _windows_rating_percent(stars), + "Microsoft:Category": joined, + "XMP:Subject": tags_norm, + "IPTC:Keywords": tags_norm, + "XPKeywords": joined, + "Keywords": joined, + "Subject": joined, + } + if tags_norm == []: + payload["XMP:Subject"] = [] + payload["IPTC:Keywords"] = [] + payload["Microsoft:Category"] = "" + payload["XPKeywords"] = "" + payload["Keywords"] = "" + payload["Subject"] = "" + return payload + + +def _exiftool_available(exiftool: ExifTool) -> bool: + try: + return bool(exiftool and exiftool.is_available()) + except (AttributeError, TypeError, ValueError) as exc: + logger.debug("ExifTool availability check failed: %s", exc) + return False + + +def _write_exiftool_payload(exiftool: ExifTool, path: Path, payload: Dict[str, Any]) -> Result[bool]: + try: + res = exiftool.write(str(path), payload, preserve_workflow=True) + except Exception as exc: + return Result.Err(ErrorCode.EXIFTOOL_ERROR, f"ExifTool write failed: {exc}") + if not res.ok: + return Result.Err(res.code or ErrorCode.EXIFTOOL_ERROR, res.error or "ExifTool write failed") + return Result.Ok(True) + + +def write_exif_rating_tags(exiftool: ExifTool, file_path: str, rating: int, tags: List[str]) -> Result[bool]: + """ + Try writing rating/tags into the file metadata via ExifTool. + """ + if not _exiftool_available(exiftool): + return Result.Err(ErrorCode.TOOL_MISSING, "ExifTool not available") + + path_res = _validate_exiftool_file_path(file_path) + if not path_res.ok: + return Result.Err(path_res.code or ErrorCode.INVALID_INPUT, path_res.error or "Invalid file path") + p = path_res.data stars = max(0, min(5, int(rating or 0))) tags_norm = _normalize_tags(tags) - joined = "; ".join(tags_norm) - try: - original_mtime = None - try: - original_mtime = os.path.getmtime(str(p)) - except OSError: - original_mtime = None - - # Multi-namespace writes to maximize Windows Explorer compatibility. - # Arrays are written as repeated values by the ExifTool adapter. - payload: Dict[str, Any] = { - "XMP:Rating": stars, - "xmp:rating": stars, - "rating": stars, - "ratingpercent": _windows_rating_percent(stars), - # Windows Property System (best-effort; writable depends on container/handler) - "Microsoft:SharedUserRating": _windows_rating_percent(stars), - "Microsoft:Category": joined, - "XMP:Subject": tags_norm, - "IPTC:Keywords": tags_norm, - "XPKeywords": joined, - "Keywords": joined, - "Subject": joined, - } - - # Clear tags when explicitly empty. - if tags_norm == []: - payload["XMP:Subject"] = [] - payload["IPTC:Keywords"] = [] - payload["Microsoft:Category"] = "" - payload["XPKeywords"] = "" - payload["Keywords"] = "" - payload["Subject"] = "" - - res = exiftool.write(str(p), payload, preserve_workflow=True) - if not res.ok: - return Result.Err(res.code or ErrorCode.EXIFTOOL_ERROR, res.error or "ExifTool write failed") - - if original_mtime is not None: - try: - os.utime(str(p), (original_mtime, original_mtime)) - except OSError as exc: - logger.debug("Failed to restore mtime after ExifTool write: %s", exc) - - return Result.Ok(True) - except Exception as exc: - return Result.Err(ErrorCode.EXIFTOOL_ERROR, f"ExifTool write failed: {exc}") + original_mtime = _get_file_mtime(p) + payload = _build_exiftool_rating_tags_payload(stars, tags_norm) + write_result = _write_exiftool_payload(exiftool, p, payload) + if not write_result.ok: + return write_result + _restore_file_mtime(p, original_mtime) + return Result.Ok(True) @dataclass(frozen=True) @@ -174,23 +182,7 @@ def _resolve_shell_indices(folder) -> tuple[int, int]: if _WIN_RATING_IDX is not None and _WIN_TAGS_IDX is not None: return _WIN_RATING_IDX, _WIN_TAGS_IDX - # Common Windows Explorer column indices; will be replaced by discovery below. - rating_idx = WIN_SHELL_DEFAULT_RATING_COL_IDX - tags_idx = WIN_SHELL_DEFAULT_TAGS_COL_IDX - try: - for i in range(0, WIN_SHELL_COL_SCAN_MAX): - name = _normalize_label(folder.GetDetailsOf(None, i)) - if not name: - continue - if _WIN_RATING_IDX is None and any(k in name for k in _RATING_KEYS): - _WIN_RATING_IDX = i - if _WIN_TAGS_IDX is None and any(k in name for k in _TAGS_KEYS): - _WIN_TAGS_IDX = i - if _WIN_RATING_IDX is not None and _WIN_TAGS_IDX is not None: - break - except Exception as exc: - logger.debug("Windows shell column discovery failed (fallback indices will be used): %s", exc) - + rating_idx, tags_idx = _discover_shell_indices(folder) if _WIN_RATING_IDX is None: _WIN_RATING_IDX = rating_idx if _WIN_TAGS_IDX is None: @@ -198,6 +190,102 @@ def _resolve_shell_indices(folder) -> tuple[int, int]: return _WIN_RATING_IDX, _WIN_TAGS_IDX +def _discover_shell_indices(folder) -> tuple[int, int]: + rating_idx = WIN_SHELL_DEFAULT_RATING_COL_IDX + tags_idx = WIN_SHELL_DEFAULT_TAGS_COL_IDX + try: + for i in range(0, WIN_SHELL_COL_SCAN_MAX): + name = _shell_column_label(folder, i) + if not name: + continue + if _should_pick_rating_index(name): + rating_idx = i + if _should_pick_tags_index(name): + tags_idx = i + if _indices_discovery_complete(rating_idx, tags_idx): + break + except Exception as exc: + logger.debug("Windows shell column discovery failed (fallback indices will be used): %s", exc) + return rating_idx, tags_idx + + +def _shell_column_label(folder, index: int) -> str: + return _normalize_label(folder.GetDetailsOf(None, index)) + + +def _should_pick_rating_index(name: str) -> bool: + return _WIN_RATING_IDX is None and any(key in name for key in _RATING_KEYS) + + +def _should_pick_tags_index(name: str) -> bool: + return _WIN_TAGS_IDX is None and any(key in name for key in _TAGS_KEYS) + + +def _indices_discovery_complete(rating_idx: int, tags_idx: int) -> bool: + rating_done = _WIN_RATING_IDX is not None or rating_idx != WIN_SHELL_DEFAULT_RATING_COL_IDX + tags_done = _WIN_TAGS_IDX is not None or tags_idx != WIN_SHELL_DEFAULT_TAGS_COL_IDX + return rating_done and tags_done + + +def _validate_sync_file_path(file_path: str) -> Result[Path]: + try: + path = Path(str(file_path)) + except (OSError, TypeError, ValueError) as exc: + logger.debug("Invalid file path for Windows shell rating/tags sync: %s", exc) + return Result.Err(ErrorCode.INVALID_INPUT, "Invalid file path") + if not path.exists() or not path.is_file(): + return Result.Err(ErrorCode.NOT_FOUND, f"File not found: {file_path}") + return Result.Ok(path) + + +def _get_file_mtime(path: Path) -> Optional[float]: + try: + return os.path.getmtime(str(path)) + except OSError: + return None + + +def _restore_file_mtime(path: Path, mtime: Optional[float]) -> None: + if mtime is None: + return + try: + os.utime(str(path), (mtime, mtime)) + except OSError as exc: + logger.debug("Failed to restore mtime after Windows shell write: %s", exc) + + +def _write_shell_rating_and_tags(folder, item, rating_idx: int, tags_idx: int, rating_val: int, tags_norm: List[str]) -> None: + try: + folder.GetDetailsOf(item, rating_idx) + folder.SetDetailsOf(item, rating_idx, str(rating_val)) + except Exception as exc: + logger.debug("Windows shell rating write skipped: %s", exc) + + try: + folder.GetDetailsOf(item, tags_idx) + folder.SetDetailsOf(item, tags_idx, "; ".join(tags_norm) if tags_norm else "") + except Exception as exc: + logger.debug("Windows shell tags write skipped: %s", exc) + + +def _co_initialize_pythoncom() -> Any: + try: + import pythoncom # type: ignore + + pythoncom.CoInitialize() + return pythoncom + except Exception: + return None + + +def _shell_write_for_path(win32com, path: Path, rating_val: int, tags_norm: List[str]) -> None: + shell = win32com.Dispatch("Shell.Application") + folder = shell.Namespace(str(path.parent)) + item = folder.ParseName(path.name) + rating_idx, tags_idx = _resolve_shell_indices(folder) + _write_shell_rating_and_tags(folder, item, rating_idx, tags_idx, rating_val, tags_norm) + + def write_windows_rating_tags(file_path: str, rating: int, tags: List[str]) -> Result[bool]: """ Windows-only fallback using Shell.Application SetDetailsOf. @@ -212,51 +300,20 @@ def write_windows_rating_tags(file_path: str, rating: int, tags: List[str]) -> R if not win32com: return Result.Err(ErrorCode.TOOL_MISSING, "pywin32 not installed (win32com unavailable)") - try: - p = Path(str(file_path)) - except (OSError, TypeError, ValueError) as exc: - logger.debug("Invalid file path for Windows shell rating/tags sync: %s", exc) - return Result.Err(ErrorCode.INVALID_INPUT, "Invalid file path") - if not p.exists() or not p.is_file(): - return Result.Err(ErrorCode.NOT_FOUND, f"File not found: {file_path}") + path_res = _validate_sync_file_path(file_path) + if not path_res.ok: + return Result.Err(path_res.code or ErrorCode.INVALID_INPUT, path_res.error or "Invalid file path") + p = path_res.data stars = max(0, min(5, int(rating or 0))) rating_val = _windows_rating_percent(stars) tags_norm = _normalize_tags(tags) - - original_mtime = None - try: - original_mtime = os.path.getmtime(str(p)) - except OSError: - original_mtime = None + original_mtime = _get_file_mtime(p) try: - pythoncom_mod: Any = None + pythoncom_mod = _co_initialize_pythoncom() try: - import pythoncom # type: ignore - - pythoncom_mod = pythoncom - pythoncom_mod.CoInitialize() - except Exception: - pythoncom_mod = None - - try: - shell = win32com.Dispatch("Shell.Application") - folder = shell.Namespace(str(p.parent)) - item = folder.ParseName(p.name) - rating_idx, tags_idx = _resolve_shell_indices(folder) - - try: - folder.GetDetailsOf(item, rating_idx) - folder.SetDetailsOf(item, rating_idx, str(rating_val)) - except Exception as exc: - logger.debug("Windows shell rating write skipped: %s", exc) - - try: - folder.GetDetailsOf(item, tags_idx) - folder.SetDetailsOf(item, tags_idx, "; ".join(tags_norm) if tags_norm else "") - except Exception as exc: - logger.debug("Windows shell tags write skipped: %s", exc) + _shell_write_for_path(win32com, p, rating_val, tags_norm) finally: if pythoncom_mod is not None: try: @@ -264,12 +321,7 @@ def write_windows_rating_tags(file_path: str, rating: int, tags: List[str]) -> R except Exception as exc: logger.debug("pythoncom.CoUninitialize failed: %s", exc) - if original_mtime is not None: - try: - os.utime(str(p), (original_mtime, original_mtime)) - except OSError as exc: - logger.debug("Failed to restore mtime after Windows shell write: %s", exc) - + _restore_file_mtime(p, original_mtime) return Result.Ok(True) except Exception as exc: return Result.Err(ErrorCode.UPDATE_FAILED, f"Windows metadata sync failed: {exc}") diff --git a/mjr_am_backend/features/viewer/info.py b/mjr_am_backend/features/viewer/info.py index 9886a535..2400f16b 100644 --- a/mjr_am_backend/features/viewer/info.py +++ b/mjr_am_backend/features/viewer/info.py @@ -11,7 +11,7 @@ def _parse_fps(value: Any) -> Optional[float]: """Safely parse FPS from various formats (number, string, fraction).""" if value is None: return None - + try: if isinstance(value, (int, float)): v = float(value) @@ -63,6 +63,91 @@ def _pick_ffprobe_audio_stream(metadata_raw: Any) -> Dict[str, Any]: return a if isinstance(a, dict) else {} +def _safe_cast(value: Any, caster: Any) -> Any: + try: + if value is None: + return None + return caster(value) + except Exception: + return None + + +def _asset_field(asset: Dict[str, Any], key: str, caster: Any) -> Any: + try: + value = asset.get(key) + except Exception: + return None + return _safe_cast(value, caster) + + +def _extract_video_fields(info: Dict[str, Any], metadata_raw: Any) -> None: + vs = _pick_ffprobe_video_stream(metadata_raw) + fps_raw = _extract_fps_raw(metadata_raw, vs) + fps = _parse_fps(fps_raw) + info["fps"] = fps + info["fps_raw"] = str(fps_raw) if fps_raw is not None else None + info["frame_count"] = _extract_frame_count(vs, info, fps) + + +def _extract_fps_raw(metadata_raw: Any, video_stream: Dict[str, Any]) -> Any: + if isinstance(metadata_raw, dict): + fps_raw = metadata_raw.get("fps") + if fps_raw is not None: + return fps_raw + return video_stream.get("avg_frame_rate") or video_stream.get("r_frame_rate") + + +def _extract_frame_count(video_stream: Dict[str, Any], info: Dict[str, Any], fps: Optional[float]) -> Optional[int]: + frame_count = _extract_stream_frame_count(video_stream) + if frame_count is not None: + return frame_count + return _estimate_frame_count(info.get("duration_s"), fps) + + +def _extract_stream_frame_count(video_stream: Dict[str, Any]) -> Optional[int]: + for key in ("nb_frames", "nb_read_frames"): + try: + raw_val = video_stream.get(key) + if raw_val is None: + continue + val = int(raw_val) + if val > 0: + return val + except (ValueError, TypeError): + continue + return None + + +def _estimate_frame_count(duration_s: Any, fps: Optional[float]) -> Optional[int]: + if not fps or not isinstance(duration_s, (int, float)) or duration_s <= 0: + return None + try: + return max(1, int(round(duration_s * fps))) + except Exception: + return None + + +def _extract_audio_fields(info: Dict[str, Any], metadata_raw: Any) -> None: + astream = _pick_ffprobe_audio_stream(metadata_raw) + info["audio_codec"] = str(astream.get("codec_name")) if astream.get("codec_name") else None + info["sample_rate"] = _safe_cast(astream.get("sample_rate"), int) + info["channels"] = _safe_cast(astream.get("channels"), int) + info["bitrate"] = _safe_cast(astream.get("bit_rate"), int) + + +def _apply_live_stats(info: Dict[str, Any], resolved_path: Optional[Path]) -> None: + if resolved_path is None: + return + try: + st = resolved_path.stat() + info["size_bytes"] = st.st_size + info["mtime"] = int(st.st_mtime) + except OSError as exc: + logger.warning(f"Failed to access file stats for {resolved_path}: {exc}") + except Exception as exc: + logger.debug(f"Unexpected error reading stats {resolved_path}: {exc}") + + def build_viewer_media_info(asset: Dict[str, Any], resolved_path: Optional[Path] = None, refresh: bool = False) -> Dict[str, Any]: """ Build a compact media info payload for the viewer UI. @@ -70,108 +155,32 @@ def build_viewer_media_info(asset: Dict[str, Any], resolved_path: Optional[Path] """ if not isinstance(asset, dict): return {} - + info: Dict[str, Any] = {} - - # Safe extraction helper - def _safe_get(key: str, type_cast: type = str) -> Any: - try: - val = asset.get(key) - if val is None: - return None - return type_cast(val) - except Exception: - return None # Basic Info info["asset_id"] = asset.get("id") - kind = str(asset.get("kind") or "").lower() info["kind"] = kind or None - ext = str(asset.get("ext") or "").lower() info["ext"] = ext or None - info["width"] = _safe_get("width", int) - info["height"] = _safe_get("height", int) - info["duration_s"] = _safe_get("duration", float) + info["width"] = _asset_field(asset, "width", int) + info["height"] = _asset_field(asset, "height", int) + info["duration_s"] = _asset_field(asset, "duration", float) metadata_raw = asset.get("metadata_raw") - # Video extended info if kind == "video": - vs = _pick_ffprobe_video_stream(metadata_raw) - - fps_raw = None - if isinstance(metadata_raw, dict): - fps_raw = metadata_raw.get("fps") - if fps_raw is None: - fps_raw = vs.get("avg_frame_rate") or vs.get("r_frame_rate") - - fps = _parse_fps(fps_raw) - info["fps"] = fps - info["fps_raw"] = str(fps_raw) if fps_raw is not None else None - - frame_count = None - # Valid frame count sources - for key in ("nb_frames", "nb_read_frames"): - try: - raw_val = vs.get(key) - if raw_val is None: - continue - val = int(raw_val) - if val > 0: - frame_count = val - break - except (ValueError, TypeError): - continue - - # Fallback: Duration * FPS - if frame_count is None: - dur = info.get("duration_s") - if fps and isinstance(dur, (int, float)) and dur > 0: - try: - frame_count = max(1, int(round(dur * fps))) - except Exception: - pass - info["frame_count"] = frame_count + _extract_video_fields(info, metadata_raw) elif kind == "audio": - astream = _pick_ffprobe_audio_stream(metadata_raw) - try: - info["audio_codec"] = str(astream.get("codec_name")) if astream.get("codec_name") else None - except Exception: - info["audio_codec"] = None - try: - sample_rate = astream.get("sample_rate") - info["sample_rate"] = int(sample_rate) if sample_rate is not None else None - except Exception: - info["sample_rate"] = None - try: - channels = astream.get("channels") - info["channels"] = int(channels) if channels is not None else None - except Exception: - info["channels"] = None - try: - bit_rate = astream.get("bit_rate") - info["bitrate"] = int(bit_rate) if bit_rate is not None else None - except Exception: - info["bitrate"] = None + _extract_audio_fields(info, metadata_raw) # File stats - prefer DB values initially - info["size_bytes"] = _safe_get("size", int) - info["mtime"] = _safe_get("mtime", int) + info["size_bytes"] = _asset_field(asset, "size", int) + info["mtime"] = _asset_field(asset, "mtime", int) # Override with live filesystem stats if path is provided - if resolved_path is not None: - try: - st = resolved_path.stat() - info["size_bytes"] = st.st_size - info["mtime"] = int(st.st_mtime) - except OSError as e: - # Log failure but allow partial result - logger.warning(f"Failed to access file stats for {resolved_path}: {e}") - except Exception as e: - logger.debug(f"Unexpected error reading stats {resolved_path}: {e}") + _apply_live_stats(info, resolved_path) return info - diff --git a/mjr_am_backend/features/watcher_settings.py b/mjr_am_backend/features/watcher_settings.py index d32a058a..0169f8fe 100644 --- a/mjr_am_backend/features/watcher_settings.py +++ b/mjr_am_backend/features/watcher_settings.py @@ -1,4 +1,4 @@ -""" +""" Runtime watcher tuning knobs (debounce + dedupe) that can be updated without restarting. """ from __future__ import annotations @@ -42,4 +42,3 @@ def update_watcher_settings(*, debounce_ms: Optional[int] = None, dedupe_ttl_ms: if dedupe_ttl_ms is not None: _state["dedupe_ttl_ms"] = _clamp(int(dedupe_ttl_ms), 100, 30_000) return get_watcher_settings() - diff --git a/mjr_am_backend/observability.py b/mjr_am_backend/observability.py index f25515f9..eabdca65 100644 --- a/mjr_am_backend/observability.py +++ b/mjr_am_backend/observability.py @@ -134,48 +134,29 @@ def _should_emit_log(key: str, *, window_ms: float) -> bool: def _should_log(request: web.Request, *, status: Optional[int], duration_ms: float) -> bool: - # Only log Majoor endpoints by default to avoid spamming ComfyUI logs. - try: - path = request.path or "" - except Exception: - path = "" - - if not path.startswith("/mjr/am/"): + path = _request_path_or_empty(request) + if not _is_majoor_route(path): return False - # Client-side control (e.g. ComfyUI settings toggle). - try: - client_flag = (request.headers.get("X-MJR-OBS") or "").strip().lower() - except Exception: - client_flag = "" - if client_flag in ("0", "false", "off", "disable", "disabled", "no"): + client_flag, client_explicit_on = _client_observability_flags(request) + if _is_obs_client_opt_out(client_flag): return False - client_explicit_on = client_flag in ("1", "true", "on", "enable", "enabled", "yes") - - # Allow opt-in verbose logging if _env_flag("MJR_OBS_LOG_ALL", default=False): return True - # Health polling can be extremely chatty; never log successful polls by default. - if path in ("/mjr/am/health", "/mjr/am/health/counters"): - if status is not None and status >= 400: - return True - return False + if _is_health_poll_path(path): + return _is_error_status(status) - # Always log errors (unless client explicitly opted out above). - if status is not None and status >= 400: + if _is_error_status(status): return True - # By default, do NOT log successes (otherwise polling endpoints spam). if not _env_flag("MJR_OBS_LOG_SUCCESS", default=False): return False - # If the client did not explicitly enable observability, keep successful requests quiet. if not client_explicit_on: return False - # Optional: log slow successful requests (threshold configurable). if not _env_flag("MJR_OBS_LOG_SLOW", default=False): return False @@ -183,6 +164,45 @@ def _should_log(request: web.Request, *, status: Optional[int], duration_ms: flo return duration_ms >= slow_ms +def _request_path_or_empty(request: web.Request) -> str: + try: + return request.path or "" + except Exception: + return "" + + +def _is_majoor_route(path: str) -> bool: + return path.startswith("/mjr/am/") + + +def _is_health_poll_path(path: str) -> bool: + return path in ("/mjr/am/health", "/mjr/am/health/counters") + + +def _is_error_status(status: Optional[int]) -> bool: + return status is not None and status >= 400 + + +def _client_observability_flags(request: web.Request) -> tuple[str, bool]: + client_flag = _request_header_lower(request, "X-MJR-OBS") + return client_flag, _is_obs_client_opt_in(client_flag) + + +def _request_header_lower(request: web.Request, header_name: str) -> str: + try: + return (request.headers.get(header_name) or "").strip().lower() + except Exception: + return "" + + +def _is_obs_client_opt_out(value: str) -> bool: + return value in ("0", "false", "off", "disable", "disabled", "no") + + +def _is_obs_client_opt_in(value: str) -> bool: + return value in ("1", "true", "on", "enable", "enabled", "yes") + + @web.middleware async def request_context_middleware(request: web.Request, handler): """Add request-id correlation and lightweight request logging context.""" @@ -191,83 +211,125 @@ async def request_context_middleware(request: web.Request, handler): rid = _get_request_id(request) request["mjr_request_id"] = rid - token = None - try: - if request_id_var is not None: - token = request_id_var.set(rid) - except Exception: - token = None + token = _set_request_id_context(rid) start = time.perf_counter() status: Optional[int] = None error: Optional[str] = None error_type: Optional[str] = None try: response = await handler(request) - try: - status = int(getattr(response, "status", 200) or 200) - except Exception: - status = 200 - response.headers["X-Request-ID"] = rid + status = _response_status_code(response) + _attach_request_id_header(response, rid) return response except web.HTTPException as exc: - try: - status = int(getattr(exc, "status", 500) or 500) - except Exception: - status = 500 - error_type = exc.__class__.__name__ - try: - error = str(exc) - except Exception: - error = "HTTPException" - try: - exc.headers["X-Request-ID"] = rid - except (AttributeError, TypeError, KeyError) as exc2: - logger.debug("Unable to attach X-Request-ID to HTTPException: %s", exc2) + status = _http_exception_status_code(exc) + error_type, error = _exception_error_fields(exc, fallback="HTTPException") + _attach_request_id_header_to_http_exception(exc, rid) raise except Exception as exc: status = 500 - error_type = exc.__class__.__name__ - try: - error = str(exc) - except Exception: - error = "Unhandled error" + error_type, error = _exception_error_fields(exc, fallback="Unhandled error") raise finally: duration_ms = (time.perf_counter() - start) * 1000.0 request["mjr_duration_ms"] = duration_ms + _reset_request_id_context(token) + _emit_request_log(request, status=status, duration_ms=duration_ms, error=error, error_type=error_type) + + +def _response_status_code(response: Any) -> int: + try: + return int(getattr(response, "status", 200) or 200) + except Exception: + return 200 + + +def _http_exception_status_code(exc: web.HTTPException) -> int: + try: + return int(getattr(exc, "status", 500) or 500) + except Exception: + return 500 + + +def _exception_error_fields(exc: Exception, *, fallback: str) -> tuple[str, str]: + error_type = exc.__class__.__name__ + try: + return error_type, str(exc) + except Exception: + return error_type, fallback + + +def _attach_request_id_header(response: Any, rid: str) -> None: + try: + response.headers["X-Request-ID"] = rid + except Exception: + return + + +def _attach_request_id_header_to_http_exception(exc: web.HTTPException, rid: str) -> None: + try: + exc.headers["X-Request-ID"] = rid + except (AttributeError, TypeError, KeyError) as exc2: + logger.debug("Unable to attach X-Request-ID to HTTPException: %s", exc2) + + +def _set_request_id_context(rid: str): + try: + if request_id_var is not None: + return request_id_var.set(rid) + except Exception: + return None + return None + + +def _reset_request_id_context(token: Any) -> None: + try: + if token is not None and request_id_var is not None: + request_id_var.reset(token) + except Exception: + return + + +def _emit_request_log( + request: web.Request, + *, + status: Optional[int], + duration_ms: float, + error: Optional[str], + error_type: Optional[str], +) -> None: + try: + fields = build_request_log_fields(request, response_status=status) + if error: + fields["error"] = error + fields["error_type"] = error_type + if not _should_log(request, status=status, duration_ms=duration_ms): + return + key = _request_log_key(request, status) + window_ms = env_float("MJR_OBS_RATELIMIT_MS", _DEFAULT_LOG_RATELIMIT_MS) + if _should_emit_log(key, window_ms=window_ms): + _emit_request_log_by_status(status, fields) + except Exception as exc: try: - if token is not None and request_id_var is not None: - request_id_var.reset(token) + logger.debug("Observability log emit failed: %s", exc) except Exception: - pass - try: - fields = build_request_log_fields(request, response_status=status) - if error: - fields["error"] = error - fields["error_type"] = error_type + return - if _should_log(request, status=status, duration_ms=duration_ms): - # Rate-limit repetitive logs (especially error loops / polling). - window_ms = env_float("MJR_OBS_RATELIMIT_MS", _DEFAULT_LOG_RATELIMIT_MS) - try: - key = f"{request.method}:{request.path}:{status}" - except Exception: - key = f"unknown:{status}" - should_emit = _should_emit_log(key, window_ms=window_ms) - if should_emit: - if status is not None and status >= 500: - logger.error("Request handled", extra=fields) - elif status is not None and status >= 400: - logger.warning("Request handled", extra=fields) - else: - logger.info("Request handled", extra=fields) - except Exception as exc: - # Never let logging failures interfere with the request. - try: - logger.debug("Observability log emit failed: %s", exc) - except Exception: - _ = None +def _request_log_key(request: web.Request, status: Optional[int]) -> str: + try: + return f"{request.method}:{request.path}:{status}" + except Exception: + return f"unknown:{status}" + + +def _emit_request_log_by_status(status: Optional[int], fields: Dict[str, Any]) -> None: + if status is not None and status >= 500: + logger.error("Request handled", extra=fields) + elif status is not None and status >= 400: + logger.warning("Request handled", extra=fields) + else: + logger.info("Request handled", extra=fields) def ensure_observability(app: web.Application) -> None: diff --git a/mjr_am_backend/probe_router.py b/mjr_am_backend/probe_router.py index 7958f2ff..0dc3204c 100644 --- a/mjr_am_backend/probe_router.py +++ b/mjr_am_backend/probe_router.py @@ -1,4 +1,4 @@ -""" +""" ProbeRouter: Decides which metadata extraction tools to use based on settings. """ from typing import List @@ -31,67 +31,67 @@ def pick_probe_backend( Returns: List of backends to use, in order: ["exiftool"], ["ffprobe"], or ["exiftool", "ffprobe"] """ - mode = settings_override or MEDIA_PROBE_BACKEND - - # Validate mode - if mode not in ("auto", "exiftool", "ffprobe", "both"): - logger.warning(f"Invalid media_probe_backend '{mode}', falling back to 'auto'") - mode = "auto" - - # Determine file type - ext = Path(filepath).suffix.lower() - is_video = ext in VIDEO_EXTS - is_audio = ext in AUDIO_EXTENSIONS - - # Simple modes + mode = _normalize_probe_mode(settings_override or MEDIA_PROBE_BACKEND) + is_media = _is_audio_or_video(filepath) if mode == "exiftool": - if has_exiftool(): - return ["exiftool"] - logger.warning("ExifTool requested but not available, trying ffprobe") - if has_ffprobe(): - return ["ffprobe"] - return [] - + return _single_mode_fallback(primary="exiftool", secondary="ffprobe") if mode == "ffprobe": - if has_ffprobe(): - return ["ffprobe"] - logger.warning("FFprobe requested but not available, trying exiftool") - if has_exiftool(): - return ["exiftool"] - return [] - + return _single_mode_fallback(primary="ffprobe", secondary="exiftool") if mode == "both": - tools = [] - if has_exiftool(): - tools.append("exiftool") - if has_ffprobe(): - tools.append("ffprobe") - return tools + return _available_tools_pair() + return _auto_mode_tools(is_media=is_media, want_generation_tags=want_generation_tags) - # Auto mode (recommended) - if mode == "auto": - if is_video or is_audio: - # For videos: use both if available - # ExifTool for generation tags, FFprobe for tech metadata - tools = [] - if has_exiftool(): - tools.append("exiftool") - if has_ffprobe(): - tools.append("ffprobe") - return tools if tools else [] - else: - # For images: ExifTool is sufficient - if has_exiftool(): - return ["exiftool"] - # Fallback to ffprobe if available (won't get generation tags though) - if has_ffprobe(): - logger.debug("Using ffprobe for image (ExifTool not available)") - return ["ffprobe"] - return [] +def _normalize_probe_mode(mode: str) -> str: + if mode in ("auto", "exiftool", "ffprobe", "both"): + return mode + logger.warning("Invalid media_probe_backend '%s', falling back to 'auto'", mode) + return "auto" + + +def _is_audio_or_video(filepath: str | Path) -> bool: + ext = Path(filepath).suffix.lower() + return ext in VIDEO_EXTS or ext in AUDIO_EXTENSIONS + + +def _single_mode_fallback(primary: str, secondary: str) -> List[str]: + if _tool_available(primary): + return [primary] + logger.warning("%s requested but not available, trying %s", primary.capitalize(), secondary) + if _tool_available(secondary): + return [secondary] return [] +def _available_tools_pair() -> List[str]: + tools: List[str] = [] + if has_exiftool(): + tools.append("exiftool") + if has_ffprobe(): + tools.append("ffprobe") + return tools + + +def _auto_mode_tools(*, is_media: bool, want_generation_tags: bool) -> List[str]: + _ = want_generation_tags + if is_media: + return _available_tools_pair() + if has_exiftool(): + return ["exiftool"] + if has_ffprobe(): + logger.debug("Using ffprobe for image (ExifTool not available)") + return ["ffprobe"] + return [] + + +def _tool_available(name: str) -> bool: + if name == "exiftool": + return has_exiftool() + if name == "ffprobe": + return has_ffprobe() + return False + + def get_probe_strategy_description(filepath: str | Path) -> str: """ Get a human-readable description of the probe strategy for a file. @@ -103,4 +103,3 @@ def get_probe_strategy_description(filepath: str | Path) -> str: if len(backends) == 1: return f"Using {backends[0]} only" return f"Using {' + '.join(backends)}" - diff --git a/mjr_am_backend/routes/__init__.py b/mjr_am_backend/routes/__init__.py index 5977609a..bc2093e2 100644 --- a/mjr_am_backend/routes/__init__.py +++ b/mjr_am_backend/routes/__init__.py @@ -1,4 +1,4 @@ -""" +""" Modular route system for ComfyUI Majoor Assets Manager. Importing this package is side-effect free; route registration is explicit. """ diff --git a/mjr_am_backend/routes/core/paths.py b/mjr_am_backend/routes/core/paths.py index 6a2a3aa7..0682e98f 100644 --- a/mjr_am_backend/routes/core/paths.py +++ b/mjr_am_backend/routes/core/paths.py @@ -1,4 +1,4 @@ -""" +""" Path validation and security utilities. """ import os @@ -80,39 +80,48 @@ def _normalize_path(value: str) -> Optional[Path]: def _is_path_allowed(candidate: Optional[Path], *, must_exist: bool = False) -> bool: if not candidate: return False - - # Prefer strict resolution when possible (prevents symlink/junction escapes). - cand_norm = candidate - try: - if must_exist: - cand_norm = cand_norm.resolve(strict=True) - elif cand_norm.exists(): - cand_norm = cand_norm.resolve(strict=True) - else: - cand_norm = cand_norm.resolve(strict=False) - except (OSError, RuntimeError, ValueError): + cand_norm = _resolve_candidate_path(candidate, must_exist=must_exist) + if cand_norm is None: return False for root in _get_allowed_directories(): - try: - root_resolved = root.resolve(strict=True) if root.exists() else root.resolve(strict=False) - except OSError: + root_resolved = _resolve_root_path(root) + if root_resolved is None: continue - - try: - if cand_norm == root_resolved or cand_norm.is_relative_to(root_resolved): - return True - except AttributeError: - # Fallback when is_relative_to is unavailable - try: - if os.path.commonpath([str(cand_norm), str(root_resolved)]) == str(root_resolved): - return True - except ValueError: - continue + if _path_relative_to(cand_norm, root_resolved): + return True return False +def _resolve_candidate_path(candidate: Path, *, must_exist: bool) -> Optional[Path]: + try: + if must_exist: + return candidate.resolve(strict=True) + if candidate.exists(): + return candidate.resolve(strict=True) + return candidate.resolve(strict=False) + except (OSError, RuntimeError, ValueError): + return None + + +def _resolve_root_path(root: Path) -> Optional[Path]: + try: + return root.resolve(strict=True) if root.exists() else root.resolve(strict=False) + except OSError: + return None + + +def _path_relative_to(candidate: Path, root: Path) -> bool: + try: + return candidate == root or candidate.is_relative_to(root) + except AttributeError: + try: + return os.path.commonpath([str(candidate), str(root)]) == str(root) + except ValueError: + return False + + def _is_path_allowed_custom(candidate: Optional[Path]) -> bool: """ Allow paths within any registered custom root. @@ -268,4 +277,3 @@ def _is_allowed_view_media_file(path: Path) -> bool: return ct.startswith("image/") or ct.startswith("video/") or ct.startswith("audio/") except Exception: return False - diff --git a/mjr_am_backend/routes/core/request_json.py b/mjr_am_backend/routes/core/request_json.py index f31deaba..5727ce20 100644 --- a/mjr_am_backend/routes/core/request_json.py +++ b/mjr_am_backend/routes/core/request_json.py @@ -1,4 +1,4 @@ -""" +""" Safe JSON request parsing with size limits. Guarantees: @@ -43,17 +43,29 @@ async def _read_json(request: web.Request, *, max_bytes: Optional[int] = None) - limit = int(max_bytes) if max_bytes is not None else _max_json_bytes() limit = max(MIN_JSON_BYTES, limit) - # Fast path: reject obviously too large bodies by Content-Length. + length_error = _content_length_error(request, limit) + if length_error is not None: + return length_error + body = await _read_request_body_limited(request, limit) + if not body.ok: + return body + return _decode_and_parse_json_dict(body.data) + + +def _content_length_error(request: web.Request, limit: int) -> Optional[Result[dict]]: try: cl = request.headers.get("Content-Length") - if cl: - size = int(cl) - if size > limit: - return Result.Err(ErrorCode.INVALID_INPUT, f"JSON body too large ({size} > {limit})", limit=limit, size=size) + if not cl: + return None + size = int(cl) + if size <= limit: + return None + return Result.Err(ErrorCode.INVALID_INPUT, f"JSON body too large ({size} > {limit})", limit=limit, size=size) except Exception: - pass + return None - # Stream the body to enforce the limit even when Content-Length is missing/incorrect. + +async def _read_request_body_limited(request: web.Request, limit: int) -> Result[bytes]: buf = bytearray() try: async for chunk in request.content.iter_chunked(REQUEST_STREAM_CHUNK_BYTES): @@ -64,18 +76,18 @@ async def _read_json(request: web.Request, *, max_bytes: Optional[int] = None) - return Result.Err(ErrorCode.INVALID_INPUT, f"JSON body too large (> {limit})", limit=limit, size=len(buf)) except Exception as exc: return Result.Err(ErrorCode.INVALID_JSON, f"Failed to read request body: {exc}") + return Result.Ok(bytes(buf)) + +def _decode_and_parse_json_dict(body: bytes) -> Result[dict]: try: - text = bytes(buf).decode("utf-8", errors="strict") + text = body.decode("utf-8", errors="strict") except Exception as exc: return Result.Err(ErrorCode.INVALID_JSON, f"Invalid UTF-8 JSON body: {exc}") - try: parsed: Any = json.loads(text) if text else {} except Exception as exc: return Result.Err(ErrorCode.INVALID_JSON, f"Invalid JSON body: {exc}") - if not isinstance(parsed, dict): return Result.Err(ErrorCode.INVALID_JSON, "JSON body must be an object") return Result.Ok(parsed) - diff --git a/mjr_am_backend/routes/core/response.py b/mjr_am_backend/routes/core/response.py index eda6be91..e376ede0 100644 --- a/mjr_am_backend/routes/core/response.py +++ b/mjr_am_backend/routes/core/response.py @@ -1,4 +1,4 @@ -""" +""" Response utilities for route handlers. """ from typing import Optional @@ -59,4 +59,3 @@ def _json_response(result: Result, status: Optional[int] = None): pass return response - diff --git a/mjr_am_backend/routes/core/security.py b/mjr_am_backend/routes/core/security.py index 0207a09b..05a87051 100644 --- a/mjr_am_backend/routes/core/security.py +++ b/mjr_am_backend/routes/core/security.py @@ -1,4 +1,4 @@ -""" +""" Security utilities: CSRF protection and rate limiting. """ @@ -100,10 +100,7 @@ def _require_operation_enabled(operation: str, *, prefs: Mapping[str, Any] | Non from ...shared import Result op = str(operation or "").strip().lower() - if prefs and "safe_mode" in prefs: - safe_mode = bool(prefs.get("safe_mode")) - else: - safe_mode = _safe_mode_enabled() + safe_mode = bool(prefs.get("safe_mode")) if (prefs and "safe_mode" in prefs) else _safe_mode_enabled() def _pref_truthy(key: str, env_var: str) -> bool: if prefs is not None and key in prefs: @@ -113,50 +110,16 @@ def _require_operation_enabled(operation: str, *, prefs: Mapping[str, Any] | Non pass return _env_truthy(env_var) - if op in ("reset_index",): - if _pref_truthy("allow_reset_index", "MAJOOR_ALLOW_RESET_INDEX"): + static_gates = _operation_static_gates() + for ops, pref_key, env_var, message in static_gates: + if op not in ops: + continue + if _pref_truthy(pref_key, env_var): return Result.Ok(True, operation=op, safe_mode=safe_mode) - return Result.Err( - "FORBIDDEN", - "Reset index is disabled by default. Enable 'allow_reset_index' in settings or set MAJOOR_ALLOW_RESET_INDEX=1.", - operation=op, - safe_mode=safe_mode, - ) - - if op in ("delete", "asset_delete", "assets_delete"): - if _pref_truthy("allow_delete", "MAJOOR_ALLOW_DELETE"): - return Result.Ok(True, operation=op, safe_mode=safe_mode) - return Result.Err( - "FORBIDDEN", - "Delete is disabled by default. Set MAJOOR_ALLOW_DELETE=1 to enable asset deletion.", - operation=op, - safe_mode=safe_mode, - ) - - if op in ("rename", "asset_rename"): - if _pref_truthy("allow_rename", "MAJOOR_ALLOW_RENAME"): - return Result.Ok(True, operation=op, safe_mode=safe_mode) - return Result.Err( - "FORBIDDEN", - "Rename is disabled by default. Set MAJOOR_ALLOW_RENAME=1 to enable asset renaming.", - operation=op, - safe_mode=safe_mode, - ) - - if op in ("open_in_folder", "open-in-folder"): - if _pref_truthy("allow_open_in_folder", "MAJOOR_ALLOW_OPEN_IN_FOLDER"): - return Result.Ok(True, operation=op, safe_mode=safe_mode) - return Result.Err( - "FORBIDDEN", - "Open-in-folder is disabled by default. Set MAJOOR_ALLOW_OPEN_IN_FOLDER=1 to enable it.", - operation=op, - safe_mode=safe_mode, - ) + return Result.Err("FORBIDDEN", message, operation=op, safe_mode=safe_mode) if op in ("write", "rating", "tags", "asset_rating", "asset_tags"): - if not safe_mode: - return Result.Ok(True, operation=op, safe_mode=safe_mode) - if _pref_truthy("allow_write", "MAJOOR_ALLOW_WRITE"): + if _write_allowed_in_context(safe_mode=safe_mode, pref_truthy=_pref_truthy): return Result.Ok(True, operation=op, safe_mode=safe_mode) return Result.Err( "FORBIDDEN", @@ -176,6 +139,41 @@ def _require_operation_enabled(operation: str, *, prefs: Mapping[str, Any] | Non return Result.Ok(True, operation=op, safe_mode=safe_mode) +def _operation_static_gates() -> tuple[tuple[tuple[str, ...], str, str, str], ...]: + return ( + ( + ("reset_index",), + "allow_reset_index", + "MAJOOR_ALLOW_RESET_INDEX", + "Reset index is disabled by default. Enable 'allow_reset_index' in settings or set MAJOOR_ALLOW_RESET_INDEX=1.", + ), + ( + ("delete", "asset_delete", "assets_delete"), + "allow_delete", + "MAJOOR_ALLOW_DELETE", + "Delete is disabled by default. Set MAJOOR_ALLOW_DELETE=1 to enable asset deletion.", + ), + ( + ("rename", "asset_rename"), + "allow_rename", + "MAJOOR_ALLOW_RENAME", + "Rename is disabled by default. Set MAJOOR_ALLOW_RENAME=1 to enable asset renaming.", + ), + ( + ("open_in_folder", "open-in-folder"), + "allow_open_in_folder", + "MAJOOR_ALLOW_OPEN_IN_FOLDER", + "Open-in-folder is disabled by default. Set MAJOOR_ALLOW_OPEN_IN_FOLDER=1 to enable it.", + ), + ) + + +def _write_allowed_in_context(*, safe_mode: bool, pref_truthy) -> bool: + if not safe_mode: + return True + return bool(pref_truthy("allow_write", "MAJOOR_ALLOW_WRITE")) + + async def _resolve_security_prefs(services: Mapping[str, Any] | None) -> Mapping[str, Any] | None: """ Extract stored security preferences from the services container safely. @@ -278,32 +276,46 @@ def _resolve_client_ip(peer_ip: str, headers: Mapping[str, str]) -> str: return "unknown" if _is_trusted_proxy(peer): - try: - forwarded_for = str(headers.get("X-Forwarded-For") or "").strip() - except Exception: - forwarded_for = "" - if forwarded_for: - cand = forwarded_for.split(",")[0].strip() - try: - ipaddress.ip_address(cand) - return cand - except Exception: - return peer - - try: - real_ip = str(headers.get("X-Real-IP") or "").strip() - except Exception: - real_ip = "" + forwarded_ip = _forwarded_for_ip(headers) + if forwarded_ip: + return forwarded_ip + real_ip = _real_ip_from_header(headers) if real_ip: - try: - ipaddress.ip_address(real_ip) - return real_ip - except Exception: - return peer + return real_ip return peer +def _forwarded_for_ip(headers: Mapping[str, str]) -> str: + forwarded_for = _header_value(headers, "X-Forwarded-For") + if not forwarded_for: + return "" + candidate = forwarded_for.split(",")[0].strip() + return candidate if _is_valid_ip(candidate) else "" + + +def _real_ip_from_header(headers: Mapping[str, str]) -> str: + real_ip = _header_value(headers, "X-Real-IP") + if not real_ip: + return "" + return real_ip if _is_valid_ip(real_ip) else "" + + +def _header_value(headers: Mapping[str, str], key: str) -> str: + try: + return str(headers.get(key) or "").strip() + except Exception: + return "" + + +def _is_valid_ip(value: str) -> bool: + try: + ipaddress.ip_address(value) + return True + except Exception: + return False + + def _is_loopback_ip(value: str) -> bool: try: return ipaddress.ip_address(str(value or "").strip()).is_loopback @@ -395,24 +407,39 @@ def _get_comfy_user_manager(request: web.Request | None = None) -> Any: """ Best-effort resolver for ComfyUI user manager object. """ - # Prefer explicit user manager injected in app context. - if request is not None: - try: - app_mgr = request.app.get("_mjr_user_manager") - if app_mgr is not None: - return app_mgr - for key in request.app.keys(): - try: - key_name = getattr(key, "name", "") - except Exception: - key_name = "" - if key_name == "_mjr_user_manager": - cand = request.app.get(key) - if cand is not None: - return cand - except Exception: - pass + app_manager = _request_app_user_manager(request) + if app_manager is not None: + return app_manager + return _server_module_user_manager() + +def _request_app_user_manager(request: web.Request | None) -> Any: + if request is None: + return None + try: + app_mgr = request.app.get("_mjr_user_manager") + if app_mgr is not None: + return app_mgr + for key in request.app.keys(): + key_name = _app_key_name(key) + if key_name != "_mjr_user_manager": + continue + candidate = request.app.get(key) + if candidate is not None: + return candidate + except Exception: + return None + return None + + +def _app_key_name(key: Any) -> str: + try: + return str(getattr(key, "name", "") or "") + except Exception: + return "" + + +def _server_module_user_manager() -> Any: try: import sys @@ -420,9 +447,9 @@ def _get_comfy_user_manager(request: web.Request | None = None) -> Any: if server_mod is None: return None for key in ("USER_MANAGER", "user_manager"): - mgr = getattr(server_mod, key, None) - if mgr is not None: - return mgr + manager = getattr(server_mod, key, None) + if manager is not None: + return manager except Exception: return None return None @@ -663,77 +690,82 @@ def _csrf_error(request: web.Request) -> Optional[str]: 1) Require an anti-CSRF header for state-changing methods (X-Requested-With or X-CSRF-Token) 2) If Origin is present, validate it against Host (with loopback allowance) """ - method = request.method.upper() - if method not in ("POST", "PUT", "DELETE", "PATCH"): + if request.method.upper() not in ("POST", "PUT", "DELETE", "PATCH"): return None - - try: - has_xrw = bool(request.headers.get("X-Requested-With")) - has_token = bool(request.headers.get("X-CSRF-Token")) - except Exception: - has_xrw = False - has_token = False - if not has_xrw and not has_token: + if not _has_csrf_header(request): return "Missing anti-CSRF header (X-Requested-With or X-CSRF-Token)" - origin = request.headers.get("Origin") if not origin: return None if origin == "null": return "Cross-site request blocked (Origin=null)" - - host = request.headers.get("Host") + host = _resolve_effective_host(request) if not host: return "Missing Host header" + parsed = _parse_origin(origin) + if parsed is None: + return "Cross-site request blocked (invalid Origin)" + if parsed.netloc == host: + return None + if _is_loopback_origin_host_match(parsed, host): + return None + return f"Cross-site request blocked ({parsed.netloc} != {host})" - # Behind a trusted reverse proxy, Host may reflect upstream internals while - # the browser-origin host is conveyed in X-Forwarded-Host. + +def _has_csrf_header(request: web.Request) -> bool: + try: + return bool(request.headers.get("X-Requested-With")) or bool(request.headers.get("X-CSRF-Token")) + except Exception: + return False + + +def _resolve_effective_host(request: web.Request) -> str: + host = request.headers.get("Host") or "" + if not host: + return "" try: peer_ip = _extract_peer_ip(request) except Exception: peer_ip = "unknown" - if _is_trusted_proxy(peer_ip): - xf_host = str(request.headers.get("X-Forwarded-Host") or "").strip() - if xf_host: - host = xf_host.split(",")[0].strip() + if not _is_trusted_proxy(peer_ip): + return host + xf_host = str(request.headers.get("X-Forwarded-Host") or "").strip() + if not xf_host: + return host + return xf_host.split(",")[0].strip() + +def _parse_origin(origin: str): try: parsed = urlparse(origin) except Exception: - return "Cross-site request blocked (invalid Origin)" - + return None if not parsed.scheme or not parsed.netloc: - return "Cross-site request blocked (invalid Origin)" + return None + return parsed - if parsed.netloc != host: - # Allow loopback aliases (localhost, 127.0.0.1, ::1) to interoperate when the - # user opens ComfyUI via a different hostname than the server reports. - try: - origin_host = parsed.hostname or "" - origin_port = parsed.port - except Exception: - origin_host = "" - origin_port = None - try: - if ":" in host and not host.endswith("]"): - host_name, host_port_raw = host.rsplit(":", 1) - host_port = int(host_port_raw) - else: - host_name = host - host_port = None - except Exception: - host_name = host - host_port = None +def _is_loopback_origin_host_match(parsed_origin: Any, host: str) -> bool: + try: + origin_host = parsed_origin.hostname or "" + origin_port = parsed_origin.port + except Exception: + return False + host_name, host_port = _split_host_port(host) + loopback = {"localhost", "127.0.0.1", "::1"} + if origin_host not in loopback or host_name not in loopback: + return False + return origin_port is None or host_port is None or origin_port == host_port - loopback = {"localhost", "127.0.0.1", "::1"} - if origin_host in loopback and host_name in loopback: - if origin_port is None or host_port is None or origin_port == host_port: - return None - return f"Cross-site request blocked ({parsed.netloc} != {host})" - - return None +def _split_host_port(host: str) -> tuple[str, Optional[int]]: + try: + if ":" in host and not host.endswith("]"): + host_name, host_port_raw = host.rsplit(":", 1) + return host_name, int(host_port_raw) + except Exception: + return host, None + return host, None def _reset_security_state_for_tests() -> None: @@ -742,4 +774,3 @@ def _reset_security_state_for_tests() -> None: _rate_limit_state.clear() except Exception: pass - diff --git a/mjr_am_backend/routes/core/services.py b/mjr_am_backend/routes/core/services.py index 91929de8..a4dc1314 100644 --- a/mjr_am_backend/routes/core/services.py +++ b/mjr_am_backend/routes/core/services.py @@ -1,4 +1,4 @@ -""" +""" Service management and initialization. """ import asyncio @@ -86,4 +86,3 @@ async def _require_services() -> tuple[dict[str, Any] | None, Result[Any] | None def get_services_error(): """Get the current services error if any.""" return _services_error - diff --git a/mjr_am_backend/routes/handlers/assets.py b/mjr_am_backend/routes/handlers/assets.py index a86c1b73..928c1d20 100644 --- a/mjr_am_backend/routes/handlers/assets.py +++ b/mjr_am_backend/routes/handlers/assets.py @@ -1,4 +1,4 @@ -""" +""" Asset management endpoints: ratings, tags, service retry. """ from aiohttp import web @@ -10,6 +10,7 @@ import subprocess import sys import mimetypes from pathlib import Path +from typing import Any, Optional from mjr_am_backend.shared import Result, get_logger, sanitize_error_message as _safe_error_message from mjr_am_backend.config import TO_THREAD_TIMEOUT_S, get_runtime_output_root @@ -56,36 +57,57 @@ _WINDOWS_RESERVED = { def _validate_filename(name: str) -> tuple[bool, str]: """Validate a filename for safety.""" - if not name or not name.strip(): + normalized = _normalize_filename(name) + if not normalized: return False, "Filename cannot be empty" - - name = name.strip() - - # Check for path separators - if "/" in name or "\\" in name: - return False, "Filename cannot contain path separators" - - # Check for null bytes - if "\x00" in name: - return False, "Filename cannot contain null bytes" - - # Check for control characters - for char in name: - if ord(char) < 32: - return False, "Filename cannot contain control characters" + separator_error = _filename_separator_error(normalized) + if separator_error: + return False, separator_error + char_error = _filename_char_error(normalized) + if char_error: + return False, char_error + boundary_error = _filename_boundary_error(normalized) + if boundary_error: + return False, boundary_error + reserved_error = _filename_reserved_error(normalized) + if reserved_error: + return False, reserved_error + return True, "" - # Check for leading/trailing dots or spaces (Windows) + +def _normalize_filename(name: str) -> str: + if not name: + return "" + return name.strip() + + +def _filename_separator_error(name: str) -> str: + if "/" in name or "\\" in name: + return "Filename cannot contain path separators" + return "" + + +def _filename_char_error(name: str) -> str: + if "\x00" in name: + return "Filename cannot contain null bytes" + if any(ord(char) < 32 for char in name): + return "Filename cannot contain control characters" + return "" + + +def _filename_boundary_error(name: str) -> str: if name.startswith('.') or name.startswith(' '): - return False, "Filename cannot start with dot or space" + return "Filename cannot start with dot or space" if name.endswith('.') or name.endswith(' '): - return False, "Filename cannot end with dot or space" - - # Check for reserved names (Windows) + return "Filename cannot end with dot or space" + return "" + + +def _filename_reserved_error(name: str) -> str: base = name.split('.')[0].upper() if base in _WINDOWS_RESERVED: - return False, "Filename uses a reserved Windows name" - - return True, "" + return "Filename uses a reserved Windows name" + return "" def _is_resolved_path_allowed(path: Path) -> bool: """ @@ -1266,74 +1288,69 @@ async def download_asset(request: web.Request) -> web.StreamResponse: Query param: filepath (absolute path, must be within allowed roots) """ - # Optional preview mode for in-app media display (viewer/grid fallback URLs). - # Preview mode can trigger many concurrent thumbnail requests from the grid, so it - # needs a much higher bucket than user-triggered downloads. - preview = str(request.query.get("preview", "")).strip().lower() in ("1", "true", "yes", "on") - - # Rate limiting - if preview: - allowed, retry_after = _check_rate_limit( - request, - "download_asset_preview", - max_requests=2000, - window_seconds=60, - ) - else: - allowed, retry_after = _check_rate_limit( - request, - "download_asset", - max_requests=30, - window_seconds=60, - ) - if not allowed: - return web.Response(status=429, text=f"Rate limit exceeded. Retry after {retry_after}s") + preview = _is_preview_download_request(request) + rate_limited = _download_rate_limit_response_or_none(request, preview=preview) + if rate_limited is not None: + return rate_limited filepath = request.query.get("filepath") if not filepath: return web.Response(status=400, text="Missing 'filepath' parameter") - # Normalize and validate path + resolved_path = _resolve_download_path(filepath) + if not isinstance(resolved_path, Path): + return resolved_path + return _build_download_response(resolved_path, preview=preview) + + +def _is_preview_download_request(request: web.Request) -> bool: + return str(request.query.get("preview", "")).strip().lower() in ("1", "true", "yes", "on") + + +def _download_rate_limit_response_or_none(request: web.Request, *, preview: bool) -> Optional[web.Response]: + key = "download_asset_preview" if preview else "download_asset" + max_requests = 2000 if preview else 30 + allowed, retry_after = _check_rate_limit(request, key, max_requests=max_requests, window_seconds=60) + if allowed: + return None + return web.Response(status=429, text=f"Rate limit exceeded. Retry after {retry_after}s") + + +def _resolve_download_path(filepath: Any) -> Path | web.Response: candidate = _normalize_path(filepath) if not candidate: return web.Response(status=400, text="Invalid filepath") - - # Strict resolution to prevent symlink escapes try: resolved = candidate.resolve(strict=True) except (OSError, RuntimeError, ValueError): return web.Response(status=404, text="File not found") if not _is_resolved_path_allowed(resolved): return web.Response(status=403, text="Path is not within allowed roots") - if not resolved.is_file(): return web.Response(status=404, text="File not found") + return resolved - # Determine content type + +def _build_download_response(resolved: Path, *, preview: bool) -> web.StreamResponse: mime_type, _ = mimetypes.guess_type(str(resolved)) - if mime_type is None: - mime_type = "application/octet-stream" - - # Sanitize filename for Content-Disposition header - filename = resolved.name.replace('"', '').replace('\r', '').replace('\n', '')[:255] - - # Use FileResponse so aiohttp handles efficient file streaming and HTTP range - # semantics for media previews. - response = web.FileResponse(path=str(resolved)) + safe_mime = mime_type or "application/octet-stream" + filename = _safe_download_filename(resolved.name) disposition = "inline" if preview else "attachment" + + response = web.FileResponse(path=str(resolved)) try: - response.headers["Content-Type"] = mime_type + response.headers["Content-Type"] = safe_mime response.headers["Content-Disposition"] = f'{disposition}; filename="{filename}"' response.headers["X-Content-Type-Options"] = "nosniff" - if preview: - response.headers["Cache-Control"] = "private, max-age=60" - else: - response.headers["Cache-Control"] = "private, no-cache" + response.headers["Cache-Control"] = "private, max-age=60" if preview else "private, no-cache" except Exception: pass return response +def _safe_download_filename(name: str) -> str: + return str(name or "").replace('"', "").replace("\r", "").replace("\n", "")[:255] + + def register_download_routes(routes: web.RouteTableDef): routes.get("/mjr/am/download")(download_asset) - diff --git a/mjr_am_backend/routes/handlers/batch_zip.py b/mjr_am_backend/routes/handlers/batch_zip.py index ecf8d1f4..78b8c423 100644 --- a/mjr_am_backend/routes/handlers/batch_zip.py +++ b/mjr_am_backend/routes/handlers/batch_zip.py @@ -1,4 +1,4 @@ -""" +""" Batch ZIP builder for drag-out (AssetsManager -> OS). Used by the frontend to support multi-selection drag-out via "DownloadURL". @@ -94,87 +94,116 @@ def _sanitize_token(token: Any) -> str: def _cleanup_batch_zips() -> None: now = time.time() - stale: List[str] = [] with _BATCH_LOCK: - for token, entry in list(_BATCH_CACHE.items()): - created_at = float(entry.get("created_at") or 0) - if created_at and now - created_at > _BATCH_TTL_SECONDS: - stale.append(token) - for token in stale: - entry = _BATCH_CACHE.pop(token, {}) - path = entry.get("path") if isinstance(entry, dict) else None - _release_token_lock(token) - if isinstance(path, Path): - try: - if path.exists(): - path.unlink() - logger.debug("Cleaned up stale batch zip: %s", path.name) - except Exception as exc: - logger.warning("Failed to cleanup batch zip %s: %s", path.name if path else token, exc) + _cleanup_stale_batch_entries(now) + _cleanup_batch_cache_cap() - # Cap cache size in case of unexpected usage. - if len(_BATCH_CACHE) > _BATCH_MAX: - items = sorted(_BATCH_CACHE.items(), key=lambda kv: float(kv[1].get("created_at") or 0)) - to_drop = items[: max(0, len(items) - _BATCH_MAX)] - for token, entry in to_drop: - _BATCH_CACHE.pop(token, None) - _release_token_lock(token) - path = entry.get("path") if isinstance(entry, dict) else None - if isinstance(path, Path): - try: - if path.exists(): - path.unlink() - logger.debug("Cleaned up batch zip (cache cap): %s", path.name) - except Exception as exc: - logger.warning("Failed to cleanup batch zip %s: %s", path.name if path else token, exc) + +def _cleanup_stale_batch_entries(now: float) -> None: + stale: List[str] = [] + for token, entry in list(_BATCH_CACHE.items()): + created_at = float(entry.get("created_at") or 0) + if created_at and now - created_at > _BATCH_TTL_SECONDS: + stale.append(token) + for token in stale: + entry = _BATCH_CACHE.pop(token, {}) + _release_token_lock(token) + _delete_batch_zip_path(entry, token, "Cleaned up stale batch zip: %s") + + +def _cleanup_batch_cache_cap() -> None: + if len(_BATCH_CACHE) <= _BATCH_MAX: + return + items = sorted(_BATCH_CACHE.items(), key=lambda kv: float(kv[1].get("created_at") or 0)) + to_drop = items[: max(0, len(items) - _BATCH_MAX)] + for token, entry in to_drop: + _BATCH_CACHE.pop(token, None) + _release_token_lock(token) + _delete_batch_zip_path(entry, token, "Cleaned up batch zip (cache cap): %s") + + +def _delete_batch_zip_path(entry: Any, token: str, success_log: str) -> None: + path = entry.get("path") if isinstance(entry, dict) else None + if not isinstance(path, Path): + return + try: + if path.exists(): + path.unlink() + logger.debug(success_log, path.name) + except Exception as exc: + logger.warning("Failed to cleanup batch zip %s: %s", path.name if path else token, exc) def _resolve_item_path(item: Dict[str, Any]) -> Optional[Path]: + filename_rel, subfolder_rel, typ = _normalized_batch_zip_item_parts(item) + if filename_rel is None or subfolder_rel is None: + return None + base_dir = _resolve_base_dir(item, typ) + if base_dir is None: + return None + candidate = _resolve_candidate_path(base_dir, subfolder_rel, filename_rel) + if not _is_valid_candidate_path(candidate, base_dir): + return None + return candidate + + +def _normalized_batch_zip_item_parts(item: Dict[str, Any]) -> tuple[Optional[Path], Optional[Path], str]: filename_raw = (item or {}).get("filename") subfolder_raw = (item or {}).get("subfolder", "") or "" typ = str((item or {}).get("type") or "output").lower() - filename_rel = _safe_rel_path(str(filename_raw or "")) if not filename_rel or len(filename_rel.parts) != 1: - return None - + return None, None, typ subfolder_rel = _safe_rel_path(str(subfolder_raw)) if subfolder_rel is None: - return None + return None, None, typ + return filename_rel, subfolder_rel, typ - base_dir: Optional[Path] = None + +def _is_valid_candidate_path(candidate: Optional[Path], base_dir: Path) -> bool: + if candidate is None: + return False + if not _is_within_root(candidate, base_dir): + return False + return candidate.exists() and candidate.is_file() + + +def _resolve_base_dir(item: Dict[str, Any], typ: str) -> Optional[Path]: if typ == "input": - try: - if folder_paths is not None: - base_dir = Path(str(folder_paths.get_input_directory())).resolve(strict=True) - except Exception: - base_dir = None - if base_dir is None: - return None - elif typ == "custom": - root_id = (item or {}).get("root_id") or (item or {}).get("rootId") or (item or {}).get("custom_root_id") or "" - root_res = resolve_custom_root(str(root_id)) - if not root_res.ok or not isinstance(root_res.data, Path): - return None - try: - base_dir = root_res.data.resolve(strict=True) - except Exception: - return None - else: - try: - base_dir = OUTPUT_ROOT_PATH.resolve(strict=True) - except Exception: - return None - + return _resolve_input_base_dir() + if typ == "custom": + return _resolve_custom_base_dir(item) try: - candidate = (base_dir / subfolder_rel / filename_rel).resolve(strict=True) + return OUTPUT_ROOT_PATH.resolve(strict=True) except Exception: return None - if not _is_within_root(candidate, base_dir): + + +def _resolve_input_base_dir() -> Optional[Path]: + try: + if folder_paths is None: + return None + return Path(str(folder_paths.get_input_directory())).resolve(strict=True) + except Exception: return None - if not candidate.exists() or not candidate.is_file(): + + +def _resolve_custom_base_dir(item: Dict[str, Any]) -> Optional[Path]: + root_id = (item or {}).get("root_id") or (item or {}).get("rootId") or (item or {}).get("custom_root_id") or "" + root_res = resolve_custom_root(str(root_id)) + if not root_res.ok or not isinstance(root_res.data, Path): + return None + try: + return root_res.data.resolve(strict=True) + except Exception: + return None + + +def _resolve_candidate_path(base_dir: Path, subfolder_rel: Path, filename_rel: Path) -> Optional[Path]: + try: + return (base_dir / subfolder_rel / filename_rel).resolve(strict=True) + except Exception: return None - return candidate def register_batch_zip_routes(routes: web.RouteTableDef) -> None: @@ -427,4 +456,3 @@ def register_batch_zip_routes(routes: web.RouteTableDef) -> None: except Exception as exc: logger.debug("FileResponse failed: %s", exc) return _json_response(Result.Err("IO_ERROR", "Failed to serve zip"), status=500) - diff --git a/mjr_am_backend/routes/handlers/calendar.py b/mjr_am_backend/routes/handlers/calendar.py index 3da30d4a..c10d50bf 100644 --- a/mjr_am_backend/routes/handlers/calendar.py +++ b/mjr_am_backend/routes/handlers/calendar.py @@ -1,4 +1,4 @@ -""" +""" Calendar / date histogram endpoints. Used by the UI to mark days that have assets (per month). @@ -128,4 +128,3 @@ def register_calendar_routes(routes: web.RouteTableDef) -> None: return _json_response(res) return _json_response(Result.Ok({"month": month, "days": res.data or {}})) - diff --git a/mjr_am_backend/routes/handlers/collections.py b/mjr_am_backend/routes/handlers/collections.py index 81601f8f..c5ac9287 100644 --- a/mjr_am_backend/routes/handlers/collections.py +++ b/mjr_am_backend/routes/handlers/collections.py @@ -1,4 +1,4 @@ -""" +""" Collections endpoints. Collections are small JSON files that store a user-curated list of assets (by filepath + basic fields). @@ -46,11 +46,7 @@ def _minimal_asset_from_item(item: Dict[str, Any]) -> Dict[str, Any]: asset_type = str(item.get("type") or "output").lower() root_id = item.get("root_id") or item.get("rootId") or item.get("custom_root_id") or None - stat = None - try: - stat = p.stat() - except Exception: - stat = None + stat_size, stat_mtime = _safe_file_stat_fields(p) return { "id": None, @@ -59,8 +55,8 @@ def _minimal_asset_from_item(item: Dict[str, Any]) -> Dict[str, Any]: "filepath": fp, "kind": kind, "ext": p.suffix.lower(), - "size": int(getattr(stat, "st_size", 0) or 0) if stat else 0, - "mtime": int(getattr(stat, "st_mtime", 0) or 0) if stat else 0, + "size": stat_size, + "mtime": stat_mtime, "width": None, "height": None, "duration": None, @@ -73,6 +69,14 @@ def _minimal_asset_from_item(item: Dict[str, Any]) -> Dict[str, Any]: } +def _safe_file_stat_fields(path: Path) -> tuple[int, int]: + try: + stat = path.stat() + return int(getattr(stat, "st_size", 0) or 0), int(getattr(stat, "st_mtime", 0) or 0) + except Exception: + return 0, 0 + + def register_collections_routes(routes: web.RouteTableDef) -> None: """Register collection management routes.""" @routes.get("/mjr/am/collections") @@ -199,5 +203,3 @@ def register_collections_routes(routes: web.RouteTableDef) -> None: logger.debug("Collections assets enrichment skipped: %s", exc) return _json_response(Result.Ok({"id": cid, "name": collection.get("name"), "assets": assets})) - - diff --git a/mjr_am_backend/routes/handlers/custom_roots.py b/mjr_am_backend/routes/handlers/custom_roots.py index 80ab79a9..63f99ce5 100644 --- a/mjr_am_backend/routes/handlers/custom_roots.py +++ b/mjr_am_backend/routes/handlers/custom_roots.py @@ -1,4 +1,4 @@ -""" +""" Custom roots management endpoints. """ import os @@ -8,6 +8,7 @@ import asyncio import shutil import ipaddress from pathlib import Path +from typing import Optional from aiohttp import web from mjr_am_backend.config import OUTPUT_ROOT from mjr_am_backend.custom_roots import ( @@ -58,38 +59,17 @@ def _compute_folder_stats(folder_path: Path, *, max_entries: int = 200000) -> di stack = [folder_path] while stack: cur = stack.pop() - try: - with os.scandir(cur) as it: - for entry in it: - scanned += 1 - if scanned > max_entries: - truncated = True - break - try: - if entry.is_symlink(): - continue - except Exception: - continue - try: - if entry.is_dir(follow_symlinks=False): - folders += 1 - try: - stack.append(Path(entry.path)) - except Exception: - continue - continue - if entry.is_file(follow_symlinks=False): - files += 1 - try: - total_size += int(entry.stat(follow_symlinks=False).st_size or 0) - except Exception: - continue - except Exception: - continue - except Exception: - continue if truncated: break + files, folders, total_size, scanned, truncated = _scan_single_folder( + cur, + stack=stack, + files=files, + folders=folders, + total_size=total_size, + scanned=scanned, + max_entries=max_entries, + ) try: st = folder_path.stat() @@ -112,6 +92,65 @@ def _compute_folder_stats(folder_path: Path, *, max_entries: int = 200000) -> di } +def _scan_single_folder( + cur: Path, + *, + stack: list[Path], + files: int, + folders: int, + total_size: int, + scanned: int, + max_entries: int, +) -> tuple[int, int, int, int, bool]: + truncated = False + try: + with os.scandir(cur) as it: + for entry in it: + scanned += 1 + if scanned > max_entries: + truncated = True + break + if _is_symlink_entry(entry): + continue + entry_stats = _folder_entry_stats(entry) + if entry_stats is None: + continue + files += entry_stats.get("files", 0) + folders += entry_stats.get("folders", 0) + total_size += entry_stats.get("size", 0) + next_path = entry_stats.get("next_path") + if isinstance(next_path, Path): + stack.append(next_path) + except Exception: + return files, folders, total_size, scanned, False + return files, folders, total_size, scanned, truncated + + +def _is_symlink_entry(entry) -> bool: + try: + return bool(entry.is_symlink()) + except Exception: + return True + + +def _folder_entry_stats(entry) -> Optional[dict]: + try: + if entry.is_dir(follow_symlinks=False): + try: + return {"files": 0, "folders": 1, "size": 0, "next_path": Path(entry.path)} + except Exception: + return {"files": 0, "folders": 1, "size": 0, "next_path": None} + if entry.is_file(follow_symlinks=False): + try: + size = int(entry.stat(follow_symlinks=False).st_size or 0) + except Exception: + size = 0 + return {"files": 1, "folders": 0, "size": size, "next_path": None} + except Exception: + return None + return None + + def register_custom_roots_routes(routes: web.RouteTableDef) -> None: """Register custom root directory routes.""" @@ -603,4 +642,3 @@ def register_custom_roots_routes(routes: web.RouteTableDef) -> None: return _json_response(Result.Err("DELETE_FAILED", sanitize_error_message(exc, "Failed to delete folder"))) return _json_response(Result.Err("INVALID_INPUT", "Unsupported folder operation")) - diff --git a/mjr_am_backend/routes/handlers/db_maintenance.py b/mjr_am_backend/routes/handlers/db_maintenance.py index 9908290c..0ae235cc 100644 --- a/mjr_am_backend/routes/handlers/db_maintenance.py +++ b/mjr_am_backend/routes/handlers/db_maintenance.py @@ -1,4 +1,4 @@ -""" +""" Database maintenance endpoints (safe, opt-in). """ @@ -135,6 +135,45 @@ async def _replace_db_from_backup(src: Path, dst: Path) -> None: await asyncio.to_thread(shutil.copy2, src, dst) +def _normalize_asset_row_for_case_cleanup(row: dict) -> tuple[int, str]: + try: + asset_id = int(row.get("id") or 0) + except Exception: + return 0, "" + if asset_id <= 0: + return 0, "" + key = str(row.get("filepath") or "").strip().lower() + if not key: + return 0, "" + return asset_id, key + + +def _collect_case_duplicate_ids(rows: list[dict]) -> tuple[int, list[int], int]: + keep_ids: set[int] = set() + delete_ids: list[int] = [] + groups = 0 + current_key = None + seen_in_group = 0 + + for row in rows: + asset_id, key = _normalize_asset_row_for_case_cleanup(row) + if asset_id <= 0 or not key: + continue + if key != current_key: + if seen_in_group > 1: + groups += 1 + current_key = key + seen_in_group = 1 + keep_ids.add(asset_id) + continue + delete_ids.append(asset_id) + seen_in_group += 1 + + if seen_in_group > 1: + groups += 1 + return groups, delete_ids, len(keep_ids) + + async def _cleanup_assets_case_duplicates(db) -> Result[dict[str, int]]: """ Remove duplicate assets that differ only by filepath casing. @@ -158,44 +197,17 @@ async def _cleanup_assets_case_duplicates(db) -> Result[dict[str, int]]: if not rows: return Result.Ok({"groups": 0, "deleted": 0, "kept": 0}) - keep_ids: set[int] = set() - delete_ids: list[int] = [] - groups = 0 - current_key = None - seen_in_group = 0 - - for row in rows: - try: - asset_id = int(row.get("id") or 0) - except Exception: - continue - if asset_id <= 0: - continue - key = str(row.get("filepath") or "").strip().lower() - if not key: - continue - if key != current_key: - if seen_in_group > 1: - groups += 1 - current_key = key - seen_in_group = 0 - keep_ids.add(asset_id) - seen_in_group = 1 - else: - delete_ids.append(asset_id) - seen_in_group += 1 - if seen_in_group > 1: - groups += 1 + groups, delete_ids, kept_count = _collect_case_duplicate_ids(rows) if not delete_ids: - return Result.Ok({"groups": 0, "deleted": 0, "kept": len(keep_ids)}) + return Result.Ok({"groups": 0, "deleted": 0, "kept": kept_count}) placeholders = ",".join("?" for _ in delete_ids) del_res = await db.aexecute(f"DELETE FROM assets WHERE id IN ({placeholders})", tuple(delete_ids)) if not del_res.ok: return Result.Err("DB_ERROR", del_res.error or "Failed to delete duplicate assets") - return Result.Ok({"groups": int(groups), "deleted": len(delete_ids), "kept": len(keep_ids)}) + return Result.Ok({"groups": int(groups), "deleted": len(delete_ids), "kept": kept_count}) except Exception as exc: return Result.Err("DB_ERROR", safe_error_message(exc, "Failed to cleanup case duplicates")) @@ -670,4 +682,3 @@ def register_db_maintenance_routes(routes: web.RouteTableDef) -> None: } ) ) - diff --git a/mjr_am_backend/routes/handlers/duplicates.py b/mjr_am_backend/routes/handlers/duplicates.py index 7e551c64..4fc43f58 100644 --- a/mjr_am_backend/routes/handlers/duplicates.py +++ b/mjr_am_backend/routes/handlers/duplicates.py @@ -1,4 +1,4 @@ -""" +""" Duplicate and similarity detection endpoints. """ from pathlib import Path @@ -158,4 +158,3 @@ def register_duplicates_routes(routes: web.RouteTableDef) -> None: return _json_response(Result.Err("INVALID_INPUT", "merge_asset_ids must be a list")) result = await dup.merge_tags_for_group(keep_asset_id, merge_asset_ids) return _json_response(result) - diff --git a/mjr_am_backend/routes/handlers/filesystem.py b/mjr_am_backend/routes/handlers/filesystem.py index be0b0bd8..550e726b 100644 --- a/mjr_am_backend/routes/handlers/filesystem.py +++ b/mjr_am_backend/routes/handlers/filesystem.py @@ -1,4 +1,4 @@ -""" +""" Filesystem operations: background scanning and file listing. """ @@ -10,7 +10,7 @@ import os import time from collections import OrderedDict from pathlib import Path -from typing import Any, Mapping, Optional +from typing import Any, Dict, Mapping, Optional from mjr_am_backend.shared import Result, classify_file, get_logger, sanitize_error_message from mjr_am_backend.config import ( @@ -182,33 +182,13 @@ async def _kickoff_background_scan( try: key = f"{source}|{str(root_id or '')}|{normalized_dir}" now_mono = time.monotonic() - - def _peek_entry(k: str) -> tuple[float, bool]: - try: - e = _BACKGROUND_SCAN_LAST.get(k) - if isinstance(e, dict): - last_v = float(e.get("last_mono") or e.get("last") or 0.0) - in_prog = bool(e.get("in_progress")) - return last_v, in_prog - if isinstance(e, (int, float)): - return float(e), False - except Exception: - pass - return 0.0, False - - last_seen, in_progress = _peek_entry(key) - if in_progress: - return - if now_mono - last_seen < min_interval_seconds: + if _should_skip_scan_enqueue(key, now_mono, min_interval_seconds): return scan_lock = await _get_background_scan_lock(key) async with scan_lock: # Double-check under lock - last_seen, in_progress = _peek_entry(key) - if in_progress: - return - if now_mono - last_seen < min_interval_seconds: + if _should_skip_scan_enqueue(key, now_mono, min_interval_seconds): return job = { @@ -224,14 +204,8 @@ async def _kickoff_background_scan( _ensure_worker() async with _SCAN_PENDING_LOCK: - if key in _SCAN_PENDING: + if not _enqueue_background_scan_job(key, job): return - # Move to end if exists (refresh priority) - _SCAN_PENDING[key] = job - - # Enforce max pending size - while len(_SCAN_PENDING) > _SCAN_PENDING_MAX: - _SCAN_PENDING.popitem(last=False) # Only mark as enqueued once we've successfully updated _SCAN_PENDING. _set_background_entry(key, {"last_mono": now_mono, "in_progress": False, "last_wall": time.time()}) @@ -243,7 +217,7 @@ def _ensure_worker() -> None: global _SCAN_TASK if _SCAN_TASK and not _SCAN_TASK.done(): return - + _SCAN_TASK = asyncio.create_task(_worker_loop()) @@ -259,17 +233,9 @@ def _set_background_entry(key: str, value: dict[str, Any]) -> None: async def _worker_loop() -> None: while True: - task = None - async with _SCAN_PENDING_LOCK: - if _SCAN_PENDING: - _, task = _SCAN_PENDING.popitem(last=False) - else: - task = None - + task = await _pop_pending_scan_task() if not task: await asyncio.sleep(0.5) - # Check exit condition or sleep again - # We keep running to allow new tasks to arrive continue if is_db_maintenance_active(): await asyncio.sleep(0.5) @@ -282,10 +248,6 @@ async def _worker_loop() -> None: entry["in_progress"] = True entry["last_mono"] = time.monotonic() entry["last_wall"] = time.time() - - # Import strictly locally or use the core helper to avoid circular imports if possible - # But here we are effectively inside backend.routes.handlers - # We can use the helper from core. svc, error_result = await _require_services() if error_result: _record_scan_failure(str(task.get("directory")), str(task.get("source")), "SERVICE_UNAVAILABLE", str(error_result.error or "")) @@ -297,34 +259,10 @@ async def _worker_loop() -> None: continue try: - try: - await svc["index"].scan_directory( - str(task.get("directory")), - bool(task.get("recursive")), - bool(task.get("incremental")), - source=str(task.get("source") or "output"), - root_id=task.get("root_id"), - fast=bool(task.get("fast")), - background_metadata=bool(task.get("background_metadata")), - ) - except TypeError: - # Fallback in case of call signature mismatch - await svc["index"].scan_directory( - str(task.get("directory")), - bool(task.get("recursive")), - bool(task.get("incremental")), - source=str(task.get("source") or "output"), - root_id=task.get("root_id"), - ) + await _run_scan_directory_task(svc, task) except Exception as exc: _record_scan_failure(str(task.get("directory")), str(task.get("source")), "SCAN_FAILED", str(exc)) logger.warning("Background scan failed: %s", exc) - - except Exception: - try: - _record_scan_failure(str(task.get("directory")), str(task.get("source")), "EXCEPTION", "Unhandled exception") - except Exception: - pass finally: if entry is not None: entry["in_progress"] = False @@ -334,6 +272,66 @@ async def _worker_loop() -> None: await asyncio.sleep(0) +async def _pop_pending_scan_task() -> Optional[dict[str, Any]]: + async with _SCAN_PENDING_LOCK: + if not _SCAN_PENDING: + return None + _, task = _SCAN_PENDING.popitem(last=False) + return task + + +def _peek_background_entry(key: str) -> tuple[float, bool]: + try: + entry = _BACKGROUND_SCAN_LAST.get(key) + if isinstance(entry, dict): + last_v = float(entry.get("last_mono") or entry.get("last") or 0.0) + in_prog = bool(entry.get("in_progress")) + return last_v, in_prog + if isinstance(entry, (int, float)): + return float(entry), False + except Exception: + pass + return 0.0, False + + +def _should_skip_scan_enqueue(key: str, now_mono: float, min_interval_seconds: float) -> bool: + last_seen, in_progress = _peek_background_entry(key) + if in_progress: + return True + return (now_mono - last_seen) < min_interval_seconds + + +def _enqueue_background_scan_job(key: str, job: dict[str, Any]) -> bool: + if key in _SCAN_PENDING: + return False + _SCAN_PENDING[key] = job + while len(_SCAN_PENDING) > _SCAN_PENDING_MAX: + _SCAN_PENDING.popitem(last=False) + return True + + +async def _run_scan_directory_task(svc: dict[str, Any], task: dict[str, Any]) -> None: + try: + await svc["index"].scan_directory( + str(task.get("directory")), + bool(task.get("recursive")), + bool(task.get("incremental")), + source=str(task.get("source") or "output"), + root_id=task.get("root_id"), + fast=bool(task.get("fast")), + background_metadata=bool(task.get("background_metadata")), + ) + except TypeError: + # Fallback in case of call signature mismatch + await svc["index"].scan_directory( + str(task.get("directory")), + bool(task.get("recursive")), + bool(task.get("incremental")), + source=str(task.get("source") or "output"), + root_id=task.get("root_id"), + ) + + def _record_scan_failure(directory: str, source: str, code: str, error: str) -> None: # No async lock needed for simple append/pop in main thread _BACKGROUND_SCAN_FAILURES.insert(0, { @@ -439,67 +437,123 @@ def _collect_filesystem_entries_window( total = 0 start = max(0, int(offset or 0)) lim = int(limit or 0) - end = (start + lim) if lim > 0 else None + window_size = lim if lim > 0 else None for entry in os.scandir(target_dir_resolved): - try: - if not entry.is_file(): - continue - except OSError: + decision = _filesystem_entry_window_decision( + entry, + query_lower=query_lower, + browse_all=browse_all, + filter_kind=filter_kind, + filter_extensions=filter_extensions, + ) + if decision is None: continue - filename = str(entry.name or "") - if not filename: - continue - kind = classify_file(filename) - if kind == "unknown": - continue - if filter_kind and kind != filter_kind: - continue - if not browse_all and query_lower not in filename.lower(): - continue - ext = _normalize_extension(Path(filename).suffix.lower()) - if filter_extensions and ext and ext not in filter_extensions: - continue - + filename, kind = decision idx = total total += 1 if idx < start: continue - if end is not None and len(out) >= (end - start): + if window_size is not None and len(out) >= window_size: continue - try: - stat = entry.stat() - except OSError: - continue - try: - rel_to_root = Path(entry.path).parent.relative_to(base) - sub = "" if str(rel_to_root) == "." else str(rel_to_root).replace("\\", "/") - except Exception: - sub = "" - out.append( - { - "id": None, - "filename": filename, - "subfolder": sub, - "filepath": str(entry.path), - "kind": kind, - "ext": Path(filename).suffix.lower(), - "size": stat.st_size, - "mtime": int(stat.st_mtime), - "width": None, - "height": None, - "duration": None, - "rating": 0, - "tags": [], - "has_workflow": None, - "has_generation_data": None, - "type": asset_type, - "root_id": root_id, - } - ) + row = _build_filesystem_window_row(entry, filename, kind, base, asset_type, root_id) + if row is not None: + out.append(row) return out, total +def _filesystem_entry_window_decision( + entry: os.DirEntry[str], + *, + query_lower: str, + browse_all: bool, + filter_kind: str, + filter_extensions: list[str], +) -> Optional[tuple[str, str]]: + if not _is_regular_dir_entry_file(entry): + return None + filename = str(entry.name or "") + if not filename: + return None + kind = classify_file(filename) + if not _filesystem_entry_matches_filters( + filename=filename, + kind=kind, + query_lower=query_lower, + browse_all=browse_all, + filter_kind=filter_kind, + filter_extensions=filter_extensions, + ): + return None + return filename, kind + + +def _filesystem_entry_matches_filters( + *, + filename: str, + kind: str, + query_lower: str, + browse_all: bool, + filter_kind: str, + filter_extensions: list[str], +) -> bool: + if kind == "unknown": + return False + if filter_kind and kind != filter_kind: + return False + if not browse_all and query_lower not in filename.lower(): + return False + ext = _normalize_extension(Path(filename).suffix.lower()) + if filter_extensions and ext and ext not in filter_extensions: + return False + return True + + +def _is_regular_dir_entry_file(entry: os.DirEntry[str]) -> bool: + try: + return bool(entry.is_file()) + except OSError: + return False + + +def _build_filesystem_window_row( + entry: os.DirEntry[str], + filename: str, + kind: str, + base: Path, + asset_type: str, + root_id: Optional[str], +) -> Optional[dict[str, Any]]: + try: + stat = entry.stat() + except OSError: + return None + try: + rel_to_root = Path(entry.path).parent.relative_to(base) + sub = "" if str(rel_to_root) == "." else str(rel_to_root).replace("\\", "/") + except Exception: + sub = "" + return { + "id": None, + "filename": filename, + "subfolder": sub, + "filepath": str(entry.path), + "kind": kind, + "ext": Path(filename).suffix.lower(), + "size": stat.st_size, + "mtime": int(stat.st_mtime), + "width": None, + "height": None, + "duration": None, + "rating": 0, + "tags": [], + "has_workflow": None, + "has_generation_data": None, + "type": asset_type, + "root_id": root_id, + } + + async def _fs_cache_get_or_build( base: Path, target_dir_resolved: Path, @@ -507,37 +561,20 @@ async def _fs_cache_get_or_build( root_id: Optional[str], ) -> Result[dict[str, Any]]: # Made async to allow async locking if needed, though mostly sync logic - try: - dir_mtime_ns = target_dir_resolved.stat().st_mtime_ns - except OSError as exc: - return Result.Err("DIR_NOT_FOUND", sanitize_error_message(exc, "Directory not found")) - try: - ensure_fs_list_cache_watching(str(base)) - except Exception: - pass - try: - watch_token = int(get_fs_list_cache_token(str(base))) - except Exception: - watch_token = 0 - + dir_stat = _filesystem_dir_cache_state(base, target_dir_resolved) + if not dir_stat.ok: + return dir_stat + state = dir_stat.data if isinstance(dir_stat.data, dict) else {} + dir_mtime_ns = int(state.get("dir_mtime_ns") or 0) + watch_token = int(state.get("watch_token") or 0) cache_key = f"{str(base)}|{str(target_dir_resolved)}|{asset_type}|{str(root_id or '')}" - + async with _FS_LIST_CACHE_LOCK: cached = _FS_LIST_CACHE.get(cache_key) - if ( - cached - and cached.get("dir_mtime_ns") == dir_mtime_ns - and int(cached.get("watch_token") or 0) == watch_token - and isinstance(cached.get("entries"), list) - ): - try: - now = time.monotonic() - cached_at = float(cached.get("cached_at_mono") or cached.get("cached_at") or 0.0) - if cached_at and (now - cached_at) <= float(FS_LIST_CACHE_TTL_SECONDS): - _FS_LIST_CACHE.move_to_end(cache_key) - return Result.Ok({"entries": cached["entries"], "dir_mtime_ns": dir_mtime_ns}) - except Exception: - pass + if _cache_entry_matches_dir_state(cached, dir_mtime_ns=dir_mtime_ns, watch_token=watch_token): + if _cache_entry_is_fresh(cached): + _FS_LIST_CACHE.move_to_end(cache_key) + return Result.Ok({"entries": cached["entries"], "dir_mtime_ns": dir_mtime_ns}) _FS_LIST_CACHE.move_to_end(cache_key) try: @@ -554,7 +591,6 @@ async def _fs_cache_get_or_build( ) async with _FS_LIST_CACHE_LOCK: - # `watch_token` may be implemented later (fs watch integration); keep None for now. _FS_LIST_CACHE[cache_key] = { "dir_mtime_ns": dir_mtime_ns, "watch_token": watch_token, @@ -569,18 +605,46 @@ async def _fs_cache_get_or_build( return Result.Ok({"entries": entries, "dir_mtime_ns": dir_mtime_ns}) -async def _list_filesystem_assets( +def _filesystem_dir_cache_state(base: Path, target_dir_resolved: Path) -> Result[dict[str, int]]: + try: + dir_mtime_ns = target_dir_resolved.stat().st_mtime_ns + except OSError as exc: + return Result.Err("DIR_NOT_FOUND", sanitize_error_message(exc, "Directory not found")) + try: + ensure_fs_list_cache_watching(str(base)) + except Exception: + pass + try: + watch_token = int(get_fs_list_cache_token(str(base))) + except Exception: + watch_token = 0 + return Result.Ok({"dir_mtime_ns": int(dir_mtime_ns), "watch_token": int(watch_token)}) + + +def _cache_entry_matches_dir_state( + cached: Optional[dict[str, Any]], *, dir_mtime_ns: int, watch_token: int +) -> bool: + return bool( + cached + and cached.get("dir_mtime_ns") == dir_mtime_ns + and int(cached.get("watch_token") or 0) == watch_token + and isinstance(cached.get("entries"), list) + ) + + +def _cache_entry_is_fresh(cached: dict[str, Any]) -> bool: + try: + now = time.monotonic() + cached_at = float(cached.get("cached_at_mono") or cached.get("cached_at") or 0.0) + return bool(cached_at and (now - cached_at) <= float(FS_LIST_CACHE_TTL_SECONDS)) + except Exception: + return False + + +def _resolve_filesystem_listing_target( root_dir: Path, subfolder: str, - query: str, - limit: int, - offset: int, - asset_type: str, - root_id: Optional[str] = None, - filters: Optional[Mapping[str, Any]] = None, - index_service: Optional[Any] = None, - sort: Optional[str] = None, -) -> Result[dict[str, Any]]: +) -> Result[dict[str, Path]]: try: base = root_dir.resolve() except OSError as exc: @@ -601,104 +665,224 @@ async def _list_filesystem_assets( if not _is_within_root(target_dir_resolved, base): return Result.Err("INVALID_INPUT", "Subfolder outside root") + return Result.Ok( + { + "base": base, + "target_dir_resolved": target_dir_resolved, + } + ) + + +def _parse_filesystem_listing_filters( + query: str, + filters: Optional[Mapping[str, Any]], + sort: Optional[str], +) -> dict[str, Any]: q = (query or "*").strip() q_lower = q.lower() - browse_all = q == "*" or q == "" + browse_all = q in ("*", "") + source = filters if filters is not None else {} - filter_kind = str(filters.get("kind") or "").strip().lower() if filters is not None else "" - filter_min_rating = int(filters.get("min_rating") or 0) if filters is not None else 0 - filter_workflow_only = bool(filters.get("has_workflow")) if filters is not None else False - filter_extensions = _normalize_extensions(filters.get("extensions") if filters is not None else None) + filter_kind = str(source.get("kind") or "").strip().lower() + filter_min_rating = int(source.get("min_rating") or 0) + filter_workflow_only = bool(source.get("has_workflow")) + filter_extensions = _normalize_extensions(source.get("extensions")) - mtime_start_raw = filters.get("mtime_start") if filters is not None else None - mtime_end_raw = filters.get("mtime_end") if filters is not None else None + mtime_start_raw = source.get("mtime_start") + mtime_end_raw = source.get("mtime_end") filter_mtime_start = int(mtime_start_raw) if mtime_start_raw is not None else None filter_mtime_end = int(mtime_end_raw) if mtime_end_raw is not None else None - sort_key = _normalize_sort_key(sort) - # Streaming fast path for massive folders: avoids stat/materialization for all entries. - if ( - sort_key == "none" - and filter_min_rating <= 0 - and not filter_workflow_only - and filter_mtime_start is None - and filter_mtime_end is None - ): - try: - entries_window, total_window = await asyncio.to_thread( - _collect_filesystem_entries_window, - target_dir_resolved, - base, - asset_type, - root_id, - query_lower=q_lower, - browse_all=browse_all, - filter_kind=filter_kind, - filter_extensions=filter_extensions, - offset=int(offset or 0), - limit=int(limit or 0), - ) - # Keep DB hydration behavior consistent with non-fast path. - try: - lookup = getattr(index_service, "lookup_assets_by_filepaths", None) - if callable(lookup) and entries_window: - filepaths = [str(a.get("filepath") or "") for a in entries_window if isinstance(a, dict)] - filepaths = [p for p in filepaths if p] - if filepaths: - enrich_result = await lookup(filepaths) - if enrich_result and getattr(enrich_result, "ok", False) and isinstance(enrich_result.data, dict): - mapping = enrich_result.data - for asset in entries_window: - if not isinstance(asset, dict): - continue - fp = str(asset.get("filepath") or "") - db_row = mapping.get(fp) - if not isinstance(db_row, dict): - continue - asset["id"] = db_row.get("id") - asset["rating"] = int(db_row.get("rating") or 0) - asset["tags"] = db_row.get("tags") or [] - asset["has_workflow"] = db_row.get("has_workflow") - asset["has_generation_data"] = db_row.get("has_generation_data") - if db_row.get("root_id"): - asset["root_id"] = db_row.get("root_id") - except Exception as exc: - logger.debug("Filesystem DB enrichment skipped (sort=none path): %s", exc) - return Result.Ok( - { - "assets": entries_window, - "total": int(total_window), - "limit": limit, - "offset": offset, - "query": q, - "sort": sort_key, - } - ) - except OSError as exc: - return Result.Err("LIST_FAILED", sanitize_error_message(exc, "Failed to list directory")) + return { + "q": q, + "q_lower": q_lower, + "browse_all": browse_all, + "filter_kind": filter_kind, + "filter_min_rating": filter_min_rating, + "filter_workflow_only": filter_workflow_only, + "filter_extensions": filter_extensions, + "filter_mtime_start": filter_mtime_start, + "filter_mtime_end": filter_mtime_end, + "sort_key": _normalize_sort_key(sort), + } - # Changed to await the async cache builder - cache_result = await _fs_cache_get_or_build(base, target_dir_resolved, asset_type, root_id) - if not cache_result.ok: - return cache_result - all_entries = cache_result.data.get("entries") if isinstance(cache_result.data, dict) else None - if not isinstance(all_entries, list): - return Result.Err("LIST_FAILED", "Failed to list directory") - entries = [] +def _can_use_listing_fast_path(opts: Mapping[str, Any]) -> bool: + return ( + str(opts.get("sort_key") or "") == "none" + and int(opts.get("filter_min_rating") or 0) <= 0 + and not bool(opts.get("filter_workflow_only")) + and opts.get("filter_mtime_start") is None + and opts.get("filter_mtime_end") is None + ) + + +def _build_filesystem_listing_payload( + assets: list[dict[str, Any]], + total: int, + limit: int, + offset: int, + query: str, + sort_key: str, +) -> dict[str, Any]: + return { + "assets": assets, + "total": int(total), + "limit": limit, + "offset": offset, + "query": query, + "sort": sort_key, + } + + +async def _enrich_filesystem_entries_from_db( + index_service: Optional[Any], + entries: list[dict[str, Any]], + *, + log_label: str, +) -> None: + try: + lookup, filepaths = _resolve_filesystem_lookup(index_service, entries) + if lookup is None or not filepaths: + return + enrich_result = await lookup(filepaths) + mapping = _extract_enrichment_mapping(enrich_result) + if mapping is None: + return + for asset in entries: + _apply_db_enrichment_row(asset, mapping) + except Exception as exc: + logger.debug("%s: %s", log_label, exc) + + +def _resolve_filesystem_lookup( + index_service: Optional[Any], + entries: list[dict[str, Any]], +) -> tuple[Optional[Any], list[str]]: + lookup = getattr(index_service, "lookup_assets_by_filepaths", None) + if not callable(lookup) or not entries: + return None, [] + filepaths = [str(a.get("filepath") or "") for a in entries if isinstance(a, dict)] + return lookup, [p for p in filepaths if p] + + +def _extract_enrichment_mapping(enrich_result: Any) -> Optional[dict]: + if not enrich_result or not getattr(enrich_result, "ok", False): + return None + if not isinstance(getattr(enrich_result, "data", None), dict): + return None + return enrich_result.data + + +def _apply_db_enrichment_row(asset: Any, mapping: dict) -> None: + if not isinstance(asset, dict): + return + fp = str(asset.get("filepath") or "") + db_row = mapping.get(fp) + if not isinstance(db_row, dict): + return + asset["id"] = db_row.get("id") + asset["rating"] = int(db_row.get("rating") or 0) + asset["tags"] = db_row.get("tags") or [] + asset["has_workflow"] = db_row.get("has_workflow") + asset["has_generation_data"] = db_row.get("has_generation_data") + if db_row.get("root_id"): + asset["root_id"] = db_row.get("root_id") + + +async def _list_filesystem_assets_fast_path( + base: Path, + target_dir_resolved: Path, + asset_type: str, + root_id: Optional[str], + *, + q: str, + q_lower: str, + browse_all: bool, + filter_kind: str, + filter_extensions: list[str], + sort_key: str, + limit: int, + offset: int, + index_service: Optional[Any], +) -> Result[dict[str, Any]]: + try: + entries_window, total_window = await asyncio.to_thread( + _collect_filesystem_entries_window, + target_dir_resolved, + base, + asset_type, + root_id, + query_lower=q_lower, + browse_all=browse_all, + filter_kind=filter_kind, + filter_extensions=filter_extensions, + offset=int(offset or 0), + limit=int(limit or 0), + ) + except OSError as exc: + return Result.Err( + "LIST_FAILED", sanitize_error_message(exc, "Failed to list directory") + ) + + await _enrich_filesystem_entries_from_db( + index_service, + entries_window, + log_label="Filesystem DB enrichment skipped (sort=none path)", + ) + return Result.Ok( + _build_filesystem_listing_payload( + entries_window, + total_window, + limit, + offset, + q, + sort_key, + ) + ) + + +def _prefilter_cached_filesystem_entries( + all_entries: list[Any], + *, + filter_kind: str, + browse_all: bool, + q_lower: str, +) -> list[dict[str, Any]]: + entries: list[dict[str, Any]] = [] for item in all_entries: - if not isinstance(item, dict): - continue - filename = str(item.get("filename") or "") - kind = str(item.get("kind") or "") - if not filename or kind == "unknown": - continue - if filter_kind and kind != filter_kind: - continue - if not browse_all and q_lower not in filename.lower(): + if not _is_listing_prefilter_match( + item, + filter_kind=filter_kind, + browse_all=browse_all, + q_lower=q_lower, + ): continue entries.append(item) + return entries + +def _is_listing_prefilter_match( + item: Any, + *, + filter_kind: str, + browse_all: bool, + q_lower: str, +) -> bool: + if not isinstance(item, dict): + return False + filename = str(item.get("filename") or "") + kind = str(item.get("kind") or "") + if not filename or kind == "unknown": + return False + if filter_kind and kind != filter_kind: + return False + if not browse_all and q_lower not in filename.lower(): + return False + return True + + +def _sort_filesystem_entries(entries: list[dict[str, Any]], sort_key: str) -> None: if sort_key == "name_asc": entries.sort(key=lambda x: str(x.get("filename") or "").lower()) elif sort_key == "name_desc": @@ -708,35 +892,89 @@ async def _list_filesystem_assets( else: entries.sort(key=lambda x: _safe_mtime_int(x.get("mtime")), reverse=True) - # Enrich with DB fields (workflow/generation/tags/rating) when available. - try: - lookup = getattr(index_service, "lookup_assets_by_filepaths", None) - if callable(lookup) and entries: - filepaths = [str(a.get("filepath") or "") for a in entries if isinstance(a, dict)] - filepaths = [p for p in filepaths if p] - if filepaths: - enrich_result = await lookup(filepaths) - if enrich_result and getattr(enrich_result, "ok", False) and isinstance(enrich_result.data, dict): - mapping = enrich_result.data - for asset in entries: - if not isinstance(asset, dict): - continue - fp = str(asset.get("filepath") or "") - db_row = mapping.get(fp) - if not isinstance(db_row, dict): - continue - asset["id"] = db_row.get("id") - asset["rating"] = int(db_row.get("rating") or 0) - asset["tags"] = db_row.get("tags") or [] - asset["has_workflow"] = db_row.get("has_workflow") - asset["has_generation_data"] = db_row.get("has_generation_data") - # Prefer stored root_id if present (custom) - if db_row.get("root_id"): - asset["root_id"] = db_row.get("root_id") - except Exception as exc: - logger.debug("Filesystem DB enrichment skipped: %s", exc) - # Apply filters after enrichment. +def _entry_passes_listing_post_filters( + item: dict[str, Any], + *, + filter_extensions: list[str], + filter_kind: str, + filter_min_rating: int, + filter_workflow_only: bool, + filter_mtime_start: Optional[int], + filter_mtime_end: Optional[int], + browse_all: bool, + q_lower: str, +) -> bool: + if not _passes_extension_filter(item, filter_extensions): + return False + if not _passes_kind_filter(item, filter_kind): + return False + if not _passes_rating_filter(item, filter_min_rating): + return False + if not _passes_workflow_filter(item, filter_workflow_only): + return False + if not _passes_mtime_window(item, filter_mtime_start, filter_mtime_end): + return False + if not _passes_name_query(item, browse_all=browse_all, q_lower=q_lower): + return False + return True + + +def _passes_extension_filter(item: dict[str, Any], filter_extensions: list[str]) -> bool: + entry_ext = _normalize_extension(item.get("ext")) + if filter_extensions and entry_ext and entry_ext not in filter_extensions: + return False + return True + + +def _passes_kind_filter(item: dict[str, Any], filter_kind: str) -> bool: + return not (filter_kind and item.get("kind") != filter_kind) + + +def _passes_rating_filter(item: dict[str, Any], filter_min_rating: int) -> bool: + return not (filter_min_rating > 0 and int(item.get("rating", 0)) < filter_min_rating) + + +def _passes_workflow_filter(item: dict[str, Any], filter_workflow_only: bool) -> bool: + return not (filter_workflow_only and not _is_truthy_boolish(item.get("has_workflow"))) + + +def _item_mtime_or_zero(item: dict[str, Any]) -> int: + try: + return int(item.get("mtime") or 0) + except Exception: + return 0 + + +def _passes_mtime_window(item: dict[str, Any], start: Optional[int], end: Optional[int]) -> bool: + entry_mtime = _item_mtime_or_zero(item) + if start is not None and entry_mtime < start: + return False + if end is not None and entry_mtime >= end: + return False + return True + + +def _passes_name_query(item: dict[str, Any], *, browse_all: bool, q_lower: str) -> bool: + if browse_all: + return True + return q_lower in str(item.get("filename", "")).lower() + + +def _paginate_filesystem_listing_entries( + entries: list[dict[str, Any]], + *, + filter_extensions: list[str], + filter_kind: str, + filter_min_rating: int, + filter_workflow_only: bool, + filter_mtime_start: Optional[int], + filter_mtime_end: Optional[int], + browse_all: bool, + q_lower: str, + limit: int, + offset: int, +) -> tuple[list[dict[str, Any]], int]: total = 0 paged: list[dict[str, Any]] = [] start = max(0, int(offset or 0)) @@ -744,31 +982,186 @@ async def _list_filesystem_assets( end = start + limit_int if limit_int > 0 else None for item in entries: - entry_ext = _normalize_extension(item.get("ext")) - if filter_extensions and entry_ext and entry_ext not in filter_extensions: + if not _entry_passes_listing_post_filters( + item, + filter_extensions=filter_extensions, + filter_kind=filter_kind, + filter_min_rating=filter_min_rating, + filter_workflow_only=filter_workflow_only, + filter_mtime_start=filter_mtime_start, + filter_mtime_end=filter_mtime_end, + browse_all=browse_all, + q_lower=q_lower, + ): continue - if filter_kind and item.get("kind") != filter_kind: - continue - if filter_min_rating > 0 and int(item.get("rating", 0)) < filter_min_rating: - continue - if filter_workflow_only and not _is_truthy_boolish(item.get("has_workflow")): - continue - try: - entry_mtime = int(item.get("mtime") or 0) - except Exception: - entry_mtime = 0 - if filter_mtime_start is not None and entry_mtime < filter_mtime_start: - continue - if filter_mtime_end is not None and entry_mtime >= filter_mtime_end: - continue - if not browse_all and q_lower not in str(item.get("filename", "")).lower(): - continue - - # This entry passed all filters; count it for total. if total >= start: if end is None or len(paged) < (end - start): paged.append(item) total += 1 + return paged, total - return Result.Ok({"assets": paged, "total": total, "limit": limit, "offset": offset, "query": q, "sort": sort_key}) +async def _list_filesystem_assets_cached_path( + base: Path, + target_dir_resolved: Path, + asset_type: str, + root_id: Optional[str], + *, + q: str, + q_lower: str, + browse_all: bool, + filter_kind: str, + filter_min_rating: int, + filter_workflow_only: bool, + filter_extensions: list[str], + filter_mtime_start: Optional[int], + filter_mtime_end: Optional[int], + sort_key: str, + limit: int, + offset: int, + index_service: Optional[Any], +) -> Result[dict[str, Any]]: + cache_result = await _fs_cache_get_or_build(base, target_dir_resolved, asset_type, root_id) + if not cache_result.ok: + return cache_result + all_entries = cache_result.data.get("entries") if isinstance(cache_result.data, dict) else None + if not isinstance(all_entries, list): + return Result.Err("LIST_FAILED", "Failed to list directory") + + entries = _prefilter_cached_filesystem_entries( + all_entries, + filter_kind=filter_kind, + browse_all=browse_all, + q_lower=q_lower, + ) + _sort_filesystem_entries(entries, sort_key) + + await _enrich_filesystem_entries_from_db( + index_service, + entries, + log_label="Filesystem DB enrichment skipped", + ) + + paged, total = _paginate_filesystem_listing_entries( + entries, + filter_extensions=filter_extensions, + filter_kind=filter_kind, + filter_min_rating=filter_min_rating, + filter_workflow_only=filter_workflow_only, + filter_mtime_start=filter_mtime_start, + filter_mtime_end=filter_mtime_end, + browse_all=browse_all, + q_lower=q_lower, + limit=limit, + offset=offset, + ) + return Result.Ok( + _build_filesystem_listing_payload( + paged, + total, + limit, + offset, + q, + sort_key, + ) + ) + + +async def _list_filesystem_assets( + root_dir: Path, + subfolder: str, + query: str, + limit: int, + offset: int, + asset_type: str, + root_id: Optional[str] = None, + filters: Optional[Mapping[str, Any]] = None, + index_service: Optional[Any] = None, + sort: Optional[str] = None, +) -> Result[dict[str, Any]]: + target_result = _resolve_filesystem_listing_target(root_dir, subfolder) + if not target_result.ok: + return target_result + target_data = target_result.data if isinstance(target_result.data, dict) else {} + base = target_data.get("base") + target_dir_resolved = target_data.get("target_dir_resolved") + if not isinstance(base, Path) or not isinstance(target_dir_resolved, Path): + return Result.Err("LIST_FAILED", "Failed to resolve listing directory") + + opts = _parse_filesystem_listing_filters(query, filters, sort) + listing_args = _listing_args_from_opts(opts) + return await _dispatch_filesystem_listing_path( + base, + target_dir_resolved, + asset_type, + root_id, + listing_args=listing_args, + limit=limit, + offset=offset, + index_service=index_service, + ) + + +def _listing_args_from_opts(opts: Mapping[str, Any]) -> Dict[str, Any]: + return { + "q": str(opts.get("q") or "*"), + "q_lower": str(opts.get("q_lower") or ""), + "browse_all": bool(opts.get("browse_all")), + "filter_kind": str(opts.get("filter_kind") or ""), + "filter_min_rating": int(opts.get("filter_min_rating") or 0), + "filter_workflow_only": bool(opts.get("filter_workflow_only")), + "filter_extensions": list(opts.get("filter_extensions") or []), + "filter_mtime_start": opts.get("filter_mtime_start"), + "filter_mtime_end": opts.get("filter_mtime_end"), + "sort_key": str(opts.get("sort_key") or "mtime_desc"), + "opts": opts, + } + + +async def _dispatch_filesystem_listing_path( + base: Path, + target_dir_resolved: Path, + asset_type: str, + root_id: Optional[str], + *, + listing_args: Dict[str, Any], + limit: int, + offset: int, + index_service: Optional[Any], +) -> Result[dict[str, Any]]: + opts = listing_args.get("opts") if isinstance(listing_args.get("opts"), Mapping) else {} + if _can_use_listing_fast_path(opts): + return await _list_filesystem_assets_fast_path( + base, + target_dir_resolved, + asset_type, + root_id, + q=listing_args["q"], + q_lower=listing_args["q_lower"], + browse_all=listing_args["browse_all"], + filter_kind=listing_args["filter_kind"], + filter_extensions=listing_args["filter_extensions"], + sort_key=listing_args["sort_key"], + limit=limit, + offset=offset, + index_service=index_service, + ) + return await _list_filesystem_assets_cached_path( + base, + target_dir_resolved, + asset_type, + root_id, + q=listing_args["q"], + q_lower=listing_args["q_lower"], + browse_all=listing_args["browse_all"], + filter_kind=listing_args["filter_kind"], + filter_min_rating=listing_args["filter_min_rating"], + filter_workflow_only=listing_args["filter_workflow_only"], + filter_extensions=listing_args["filter_extensions"], + filter_mtime_start=listing_args["filter_mtime_start"], + filter_mtime_end=listing_args["filter_mtime_end"], + sort_key=listing_args["sort_key"], + limit=limit, + offset=offset, + index_service=index_service, + ) diff --git a/mjr_am_backend/routes/handlers/metadata.py b/mjr_am_backend/routes/handlers/metadata.py index 44a76e90..19a1b2fd 100644 --- a/mjr_am_backend/routes/handlers/metadata.py +++ b/mjr_am_backend/routes/handlers/metadata.py @@ -1,4 +1,4 @@ -""" +""" Metadata extraction endpoint. """ import asyncio @@ -124,4 +124,3 @@ def register_metadata_routes(routes: web.RouteTableDef) -> None: ) return _json_response(result) - diff --git a/mjr_am_backend/routes/handlers/releases.py b/mjr_am_backend/routes/handlers/releases.py index 6f6d2d53..ed146a65 100644 --- a/mjr_am_backend/routes/handlers/releases.py +++ b/mjr_am_backend/routes/handlers/releases.py @@ -1,4 +1,4 @@ -""" +""" Expose repository refs (tags/branches) and a ZIP template for programmatic installs. Endpoint: GET /mjr/am/releases @@ -82,4 +82,3 @@ def register_releases_routes(routes: web.RouteTableDef) -> None: }, ) ) - diff --git a/mjr_am_backend/routes/handlers/scan.py b/mjr_am_backend/routes/handlers/scan.py index 4572f0e1..e709e3be 100644 --- a/mjr_am_backend/routes/handlers/scan.py +++ b/mjr_am_backend/routes/handlers/scan.py @@ -1,4 +1,4 @@ -""" +""" Directory scanning and file indexing endpoints. """ import asyncio @@ -23,9 +23,9 @@ except Exception: from mjr_am_backend.adapters.db.schema import rebuild_fts from mjr_am_backend.config import ( - OUTPUT_ROOT, - TO_THREAD_TIMEOUT_S, - INDEX_DIR_PATH, + OUTPUT_ROOT, + TO_THREAD_TIMEOUT_S, + INDEX_DIR_PATH, COLLECTIONS_DIR_PATH ) from mjr_am_backend.custom_roots import resolve_custom_root @@ -116,15 +116,23 @@ def _is_db_malformed_result(res: Any) -> bool: def _allowed_upload_exts() -> set[str]: + allowed = _default_allowed_upload_exts() + _add_env_upload_extensions(allowed) + return allowed + + +def _default_allowed_upload_exts() -> set[str]: try: from mjr_am_backend.shared import EXTENSIONS # type: ignore + allowed: set[str] = set() for exts in (EXTENSIONS or {}).values(): for ext in exts or []: allowed.add(str(ext).lower()) allowed.update({".json", ".txt", ".csv"}) + return allowed except Exception: - allowed = { + return { ".png", ".jpg", ".jpeg", @@ -150,18 +158,27 @@ def _allowed_upload_exts() -> set[str]: ".csv", } + +def _add_env_upload_extensions(allowed: set[str]) -> None: try: extra = os.environ.get("MJR_UPLOAD_EXTRA_EXT", "") - if extra: - for item in extra.split(","): - e = item.strip().lower() - if e and not e.startswith("."): - e = "." + e - if e: - allowed.add(e) except Exception: - pass - return allowed + return + if not extra: + return + for item in extra.split(","): + normalized = _normalize_upload_extension(item) + if normalized: + allowed.add(normalized) + + +def _normalize_upload_extension(value: str) -> str: + ext = str(value or "").strip().lower() + if not ext: + return "" + if not ext.startswith("."): + ext = "." + ext + return ext async def _write_multipart_file_atomic(dest_dir: Path, filename: str, field) -> Result[Path]: @@ -177,55 +194,75 @@ async def _write_multipart_file_atomic(dest_dir: Path, filename: str, field) -> sanitize_error_message(exc, "Cannot create destination directory"), ) - safe_name = Path(str(filename)).name - if not safe_name or "\x00" in safe_name or safe_name.startswith(".") or ".." in safe_name: + safe_name = _validate_upload_filename(filename) + if safe_name is None: return Result.Err("INVALID_INPUT", "Invalid filename") ext = Path(safe_name).suffix.lower() if ext not in _allowed_upload_exts(): return Result.Err("INVALID_INPUT", "File extension not allowed") - total = 0 fd = None tmp_path = None try: fd, tmp_path = tempfile.mkstemp(dir=str(dest_dir), prefix=".upload_", suffix=ext) - with os.fdopen(fd, "wb") as f: + with os.fdopen(fd, "wb") as handle: fd = None - while True: - chunk = await field.read_chunk(size=_UPLOAD_READ_CHUNK_BYTES) - if not chunk: - break - total += len(chunk) - if total > _MAX_UPLOAD_SIZE: - raise ValueError(f"File exceeds maximum size ({_MAX_UPLOAD_SIZE} bytes)") - # Avoid blocking the event loop on large writes. - await asyncio.to_thread(f.write, chunk) + await _write_upload_chunks(field, handle) tmp_obj = Path(str(tmp_path)) - final = dest_dir / safe_name - counter = 0 - while final.exists(): - counter += 1 - if counter > _MAX_RENAME_ATTEMPTS: - raise RuntimeError("Too many rename attempts") - final = dest_dir / f"{Path(safe_name).stem}_{counter}{ext}" + final = _unique_upload_destination(dest_dir, safe_name, ext) tmp_obj.replace(final) return Result.Ok(final) except Exception as exc: - try: - if fd is not None: - os.close(fd) - except Exception: - pass - try: - if tmp_path: - Path(str(tmp_path)).unlink(missing_ok=True) - except Exception: - pass + _cleanup_temp_upload_file(fd, tmp_path) return Result.Err("UPLOAD_FAILED", safe_error_message(exc, "Upload failed")) +def _validate_upload_filename(filename: str) -> Optional[str]: + safe_name = Path(str(filename)).name + if not safe_name or "\x00" in safe_name or safe_name.startswith(".") or ".." in safe_name: + return None + return safe_name + + +async def _write_upload_chunks(field, handle) -> None: + total = 0 + while True: + chunk = await field.read_chunk(size=_UPLOAD_READ_CHUNK_BYTES) + if not chunk: + break + total += len(chunk) + if total > _MAX_UPLOAD_SIZE: + raise ValueError(f"File exceeds maximum size ({_MAX_UPLOAD_SIZE} bytes)") + # Avoid blocking the event loop on large writes. + await asyncio.to_thread(handle.write, chunk) + + +def _unique_upload_destination(dest_dir: Path, safe_name: str, ext: str) -> Path: + final = dest_dir / safe_name + counter = 0 + while final.exists(): + counter += 1 + if counter > _MAX_RENAME_ATTEMPTS: + raise RuntimeError("Too many rename attempts") + final = dest_dir / f"{Path(safe_name).stem}_{counter}{ext}" + return final + + +def _cleanup_temp_upload_file(fd: Optional[int], tmp_path: Optional[str]) -> None: + try: + if fd is not None: + os.close(fd) + except Exception: + pass + try: + if tmp_path: + Path(str(tmp_path)).unlink(missing_ok=True) + except Exception: + pass + + def _schedule_index_task(fn) -> None: async def _runner(): try: @@ -245,35 +282,48 @@ def _files_equal_content(src: Path, dst: Path, *, chunk_size: int = _FILE_COMPAR Used to avoid creating *_1 duplicates when re-staging the exact same file. """ + if not _files_exist_and_are_regular(src, dst): + return False + + if _paths_are_same_file(src, dst): + return True + if not _file_sizes_equal(src, dst): + return False + + return _stream_content_equal(src, dst, chunk_size) + + +def _files_exist_and_are_regular(src: Path, dst: Path) -> bool: try: - if not src.exists() or not dst.exists(): - return False - if not src.is_file() or not dst.is_file(): - return False + return src.exists() and dst.exists() and src.is_file() and dst.is_file() except Exception: return False + +def _paths_are_same_file(src: Path, dst: Path) -> bool: try: import os - if os.path.samefile(str(src), str(dst)): - return True - except Exception: - pass - - try: - if src.stat().st_size != dst.stat().st_size: - return False + return os.path.samefile(str(src), str(dst)) except Exception: return False + +def _file_sizes_equal(src: Path, dst: Path) -> bool: + try: + return src.stat().st_size == dst.stat().st_size + except Exception: + return False + + +def _stream_content_equal(src: Path, dst: Path, chunk_size: int) -> bool: try: with open(src, "rb") as fs, open(dst, "rb") as fd: while True: - bs = fs.read(chunk_size) - bd = fd.read(chunk_size) - if bs != bd: + left = fs.read(chunk_size) + right = fd.read(chunk_size) + if left != right: return False - if not bs: + if not left: return True except Exception: return False @@ -331,31 +381,11 @@ def _resolve_scan_root(directory: str) -> Result[Path]: async def _run_consistency_check(db: Any) -> None: if not db: return - try: - res = await db.aquery( - "SELECT id, filepath FROM assets WHERE filepath IS NOT NULL ORDER BY RANDOM() LIMIT ?", - (_DB_CONSISTENCY_SAMPLE,), - ) - except Exception as exc: - logger.debug("Consistency check query failed: %s", exc) - return - + res = await _query_consistency_sample(db) if not res.ok or not isinstance(res.data, list): return - missing = [] - for row in res.data: - if not isinstance(row, dict): - continue - asset_id = row.get("id") - filepath = row.get("filepath") - if asset_id is None or not filepath: - continue - candidate = _normalize_path(str(filepath)) - if not candidate or candidate.exists(): - continue - missing.append((asset_id, str(filepath))) - + missing = _collect_missing_asset_rows(res.data) if not missing: return @@ -363,7 +393,43 @@ async def _run_consistency_check(db: Any) -> None: "Consistency check removed %d stale asset records", len(missing), ) + await _delete_missing_asset_rows(db, missing) + +async def _query_consistency_sample(db: Any) -> Any: + try: + return await db.aquery( + "SELECT id, filepath FROM assets WHERE filepath IS NOT NULL ORDER BY RANDOM() LIMIT ?", + (_DB_CONSISTENCY_SAMPLE,), + ) + except Exception as exc: + logger.debug("Consistency check query failed: %s", exc) + return Result.Err("QUERY_FAILED", str(exc)) + + +def _collect_missing_asset_rows(rows: list[Any]) -> list[tuple[Any, str]]: + missing: list[tuple[Any, str]] = [] + for row in rows: + pair = _missing_asset_row(row) + if pair is not None: + missing.append(pair) + return missing + + +def _missing_asset_row(row: Any) -> Optional[tuple[Any, str]]: + if not isinstance(row, dict): + return None + asset_id = row.get("id") + filepath = row.get("filepath") + if asset_id is None or not filepath: + return None + candidate = _normalize_path(str(filepath)) + if not candidate or candidate.exists(): + return None + return asset_id, str(filepath) + + +async def _delete_missing_asset_rows(db: Any, missing: list[tuple[Any, str]]) -> None: try: async with db.atransaction(mode="immediate"): for asset_id, filepath in missing: @@ -480,12 +546,12 @@ def register_scan_routes(routes: web.RouteTableDef) -> None: fast, background_metadata, ) - + results = await asyncio.wait_for( asyncio.gather(out_coro, in_coro, return_exceptions=True), timeout=TO_THREAD_TIMEOUT_S ) - + out_res = results[0] in_res = results[1] @@ -496,7 +562,7 @@ def register_scan_routes(routes: web.RouteTableDef) -> None: if isinstance(in_res, Exception): logger.error("Input scan failed: %s", in_res) return _json_response(Result.Err("SCAN_FAILED", safe_error_message(in_res, "Input scan failed"))) - + except asyncio.TimeoutError: return _json_response(Result.Err("TIMEOUT", "Scan timed out")) if not out_res.ok: @@ -851,7 +917,7 @@ def register_scan_routes(routes: web.RouteTableDef) -> None: if not fname: continue s_name = item.get("subfolder") or "" s_src = (item.get("type") or "output").lower() - + # Check duplication in params list # (Simple dedupe by tuple key) key = (fname, str(s_name), s_src) @@ -2160,5 +2226,3 @@ def register_scan_routes(routes: web.RouteTableDef) -> None: await new_watcher.start(watch_paths, loop) svc["watcher"] = new_watcher return _json_response(Result.Ok({"enabled": True, "directories": new_watcher.watched_directories, "scope": scope})) - - diff --git a/mjr_am_backend/routes/handlers/search.py b/mjr_am_backend/routes/handlers/search.py index 630f6942..f94c3834 100644 --- a/mjr_am_backend/routes/handlers/search.py +++ b/mjr_am_backend/routes/handlers/search.py @@ -1,4 +1,4 @@ -""" +""" Search and list endpoints. """ import datetime @@ -6,6 +6,7 @@ import asyncio import re import os import json +from typing import Any from pathlib import Path from aiohttp import web @@ -100,25 +101,24 @@ def _norm_fp(fp: str) -> str: return str(fp or "") -async def _hydrate_browser_assets_from_db(svc: dict | None, assets: list[dict]) -> list[dict]: - if not isinstance(assets, list) or not assets: - return assets or [] - db = (svc or {}).get("db") if isinstance(svc, dict) else None - if not db: - return assets +def _is_folder_asset(asset: dict) -> bool: + return str((asset or {}).get("kind") or "").lower() == "folder" - filepaths = [] - for a in assets: - if not isinstance(a, dict): + +def _collect_hydration_filepaths(assets: list[dict]) -> list[str]: + filepaths: list[str] = [] + for asset in assets: + if not isinstance(asset, dict): continue - if str((a or {}).get("kind") or "").lower() == "folder": + if _is_folder_asset(asset): continue - fp = str((a or {}).get("filepath") or "").strip() + fp = str(asset.get("filepath") or "").strip() if fp: filepaths.append(fp) - if not filepaths: - return assets + return filepaths + +async def _query_browser_asset_rows(db: Any, filepaths: list[str]) -> list[dict] | None: try: rows_res = await db.aquery_in( """ @@ -130,12 +130,15 @@ async def _hydrate_browser_assets_from_db(svc: dict | None, assets: list[dict]) "a.filepath", filepaths, ) - if not rows_res.ok: - return assets - rows = rows_res.data or [] except Exception: - return assets + return None + if not rows_res.ok: + return None + rows = rows_res.data or [] + return rows if isinstance(rows, list) else [] + +def _index_rows_by_filepath(rows: list[dict]) -> dict[str, dict]: by_fp: dict[str, dict] = {} for row in rows: try: @@ -145,34 +148,66 @@ async def _hydrate_browser_assets_from_db(svc: dict | None, assets: list[dict]) by_fp[key] = row or {} except Exception: continue + return by_fp - for a in assets: + +def _coerce_browser_tags(raw_tags: Any) -> list[str]: + if isinstance(raw_tags, str): try: - if str((a or {}).get("kind") or "").lower() == "folder": - continue - key = _norm_fp(str((a or {}).get("filepath") or "")) - row = by_fp.get(key) - if not row: - continue - rid = (row or {}).get("id") - if rid is not None: - a["id"] = int(rid) - a["rating"] = int((row or {}).get("rating") or 0) - raw_tags = (row or {}).get("tags") - tags = [] - if isinstance(raw_tags, str): - try: - parsed = json.loads(raw_tags) - if isinstance(parsed, list): - tags = [str(t) for t in parsed if isinstance(t, str)] - except Exception: - tags = [] - elif isinstance(raw_tags, list): - tags = [str(t) for t in raw_tags if isinstance(t, str)] - a["tags"] = tags + parsed = json.loads(raw_tags) + except Exception: + return [] + if isinstance(parsed, list): + return [str(t) for t in parsed if isinstance(t, str)] + return [] + if isinstance(raw_tags, list): + return [str(t) for t in raw_tags if isinstance(t, str)] + return [] + + +def _hydrate_browser_asset_from_row(asset: dict, by_fp: dict[str, dict]) -> None: + if _is_folder_asset(asset): + return + key = _norm_fp(str(asset.get("filepath") or "")) + row = by_fp.get(key) + if not row: + return + rid = row.get("id") + if rid is not None: + asset["id"] = int(rid) + asset["rating"] = int(row.get("rating") or 0) + asset["tags"] = _coerce_browser_tags(row.get("tags")) + + +async def _hydrate_browser_assets_from_db(svc: dict | None, assets: list[dict]) -> list[dict]: + normalized_assets = assets if isinstance(assets, list) else [] + if not normalized_assets: + return normalized_assets + db = _search_db_from_services(svc) + if not db: + return normalized_assets + filepaths = _collect_hydration_filepaths(normalized_assets) + if not filepaths: + return normalized_assets + rows = await _query_browser_asset_rows(db, filepaths) + if rows is None: + return normalized_assets + _apply_browser_hydration_rows(normalized_assets, rows) + return normalized_assets + + +def _search_db_from_services(svc: dict | None) -> Any: + return (svc or {}).get("db") if isinstance(svc, dict) else None + + +def _apply_browser_hydration_rows(assets: list[dict], rows: list[dict]) -> None: + by_fp = _index_rows_by_filepath(rows) + for asset in assets: + try: + if isinstance(asset, dict): + _hydrate_browser_asset_from_row(asset, by_fp) except Exception: continue - return assets async def _runtime_output_root(svc: dict | None) -> str: @@ -286,37 +321,55 @@ def _parse_inline_query_filters(raw_query): cleaned = [] filters = {} for token in tokens: - if ":" not in token: + consumed = _consume_inline_filter_token(token, filters) + if not consumed: cleaned.append(token) - continue - key, _, value = token.partition(":") - key = (key or "").strip().lower() - value = (value or "").strip().strip(",;") - if not key or not value: - cleaned.append(token) - continue - if key in ("ext", "extension"): - ext = _normalize_extension(value) - if ext: - filters.setdefault("extensions", []) - if ext not in filters["extensions"]: - filters["extensions"].append(ext) - continue - if key == "kind": - kind_value = (value or "").strip().lower() - if kind_value in VALID_KIND_FILTERS: - filters["kind"] = kind_value - continue - if key == "rating": - match = re.match(r"^(\d+)", value) - if match: - rating_val = max(0, min(MAX_RATING, int(match.group(1)))) - filters["min_rating"] = rating_val - continue - cleaned.append(token) return " ".join(cleaned).strip(), filters +def _consume_inline_filter_token(token: str, filters: dict) -> bool: + if ":" not in token: + return False + key, _, value = token.partition(":") + key = (key or "").strip().lower() + value = (value or "").strip().strip(",;") + if not key or not value: + return False + if key in ("ext", "extension"): + return _consume_inline_extension(value, filters) + if key == "kind": + return _consume_inline_kind(value, filters) + if key == "rating": + return _consume_inline_rating(value, filters) + return False + + +def _consume_inline_extension(value: str, filters: dict) -> bool: + ext = _normalize_extension(value) + if not ext: + return False + filters.setdefault("extensions", []) + if ext not in filters["extensions"]: + filters["extensions"].append(ext) + return True + + +def _consume_inline_kind(value: str, filters: dict) -> bool: + kind_value = (value or "").strip().lower() + if kind_value not in VALID_KIND_FILTERS: + return False + filters["kind"] = kind_value + return True + + +def _consume_inline_rating(value: str, filters: dict) -> bool: + match = re.match(r"^(\d+)", value) + if not match: + return False + filters["min_rating"] = max(0, min(MAX_RATING, int(match.group(1)))) + return True + + def _normalize_sort_key(value: str | None) -> str: s = str(value or "").strip().lower() if s in VALID_SORT_KEYS: @@ -330,7 +383,7 @@ def register_search_routes(routes: web.RouteTableDef) -> None: async def autocomplete_assets(request): """ Autocomplete tags/terms. - + Query params: q: prefix to complete limit: max results (default 10) @@ -1265,4 +1318,3 @@ def register_search_routes(routes: web.RouteTableDef) -> None: pass return _json_response(result) - diff --git a/mjr_am_backend/routes/handlers/version.py b/mjr_am_backend/routes/handlers/version.py index bad973ee..64973fdb 100644 --- a/mjr_am_backend/routes/handlers/version.py +++ b/mjr_am_backend/routes/handlers/version.py @@ -1,4 +1,4 @@ -""" +""" Version reporting endpoint. """ from aiohttp import web @@ -18,4 +18,3 @@ def register_version_routes(routes: web.RouteTableDef) -> None: routes.get("/mjr/am/version")(_get_version) routes.get("/majoor/version")(_get_version) - diff --git a/mjr_am_backend/routes/handlers/viewer.py b/mjr_am_backend/routes/handlers/viewer.py index 385a9d1f..3aa9a102 100644 --- a/mjr_am_backend/routes/handlers/viewer.py +++ b/mjr_am_backend/routes/handlers/viewer.py @@ -1,4 +1,4 @@ -""" +""" Viewer helper endpoints. These endpoints provide lightweight media info for the frontend viewer without @@ -93,4 +93,3 @@ def register_viewer_routes(routes: web.RouteTableDef) -> None: info["mime"] = None return _json_response(Result.Ok(info)) - diff --git a/mjr_am_backend/routes/registry.py b/mjr_am_backend/routes/registry.py index 419457c3..59ca9e4e 100644 --- a/mjr_am_backend/routes/registry.py +++ b/mjr_am_backend/routes/registry.py @@ -1,4 +1,4 @@ -""" +""" Route registration system. Coordinates all route handlers and registers them with PromptServer or aiohttp app. """ @@ -10,7 +10,7 @@ from mjr_am_backend.shared import get_logger from mjr_am_backend.observability import ensure_observability from .core import _json_response, _require_authenticated_user from mjr_am_backend.shared import Result -from typing import Awaitable, Callable, ClassVar, Protocol, cast +from typing import Any, Awaitable, Callable, ClassVar, Optional, Protocol, cast # --- CONFIGURATION --- API_PREFIX = "/mjr/am/" @@ -203,31 +203,47 @@ async def auth_required_middleware( """ Require authenticated ComfyUI user for sensitive Majoor routes when auth is enabled. """ - try: - path = request.path or "" - method = str(request.method or "GET").upper() - except Exception: - path = "" - method = "GET" - - if path.startswith(API_PREFIX) and method in _SENSITIVE_METHODS: - auth = _require_authenticated_user(request) - if not auth.ok: - return _json_response( - Result.Err( - auth.code or "AUTH_REQUIRED", - auth.error or "Authentication required. Please sign in first.", - ), - status=401, - ) - try: - request["mjr_user_id"] = str(auth.data or "").strip() - except Exception: - pass + path, method = _request_path_and_method(request) + if _requires_auth(path, method): + failure = _auth_error_response_or_none(request) + if failure is not None: + return failure return await handler(request) +def _request_path_and_method(request: web.Request) -> tuple[str, str]: + try: + return request.path or "", str(request.method or "GET").upper() + except Exception: + return "", "GET" + + +def _requires_auth(path: str, method: str) -> bool: + return path.startswith(API_PREFIX) and method in _SENSITIVE_METHODS + + +def _auth_error_response_or_none(request: web.Request) -> Optional[web.StreamResponse]: + auth = _require_authenticated_user(request) + if not auth.ok: + return _json_response( + Result.Err( + auth.code or "AUTH_REQUIRED", + auth.error or "Authentication required. Please sign in first.", + ), + status=401, + ) + _store_request_user_id(request, auth.data) + return None + + +def _store_request_user_id(request: web.Request, user_id: Any) -> None: + try: + request["mjr_user_id"] = str(user_id or "").strip() + except Exception: + pass + + def _install_security_middlewares(app: web.Application) -> None: try: if app.get(_APP_KEY_SECURITY_MIDDLEWARES_INSTALLED): @@ -380,4 +396,3 @@ def _install_observability_on_prompt_server() -> None: except Exception as exc: logger.debug("Observability not installed on PromptServer app: %s", exc) return - diff --git a/mjr_am_backend/settings.py b/mjr_am_backend/settings.py index 84523099..d68f13a7 100644 --- a/mjr_am_backend/settings.py +++ b/mjr_am_backend/settings.py @@ -104,22 +104,11 @@ class AppSettings: return token_hash async def _get_or_create_api_token_locked(self) -> str: - token = "" - try: - token = (os.environ.get("MAJOOR_API_TOKEN") or os.environ.get("MJR_API_TOKEN") or "").strip() - except Exception: - token = "" + token = self._env_api_token() if token: - token_hash = self._hash_api_token(token) - write_res = await self._write_setting(_SECURITY_API_TOKEN_HASH_KEY, token_hash) - if not write_res.ok: - logger.warning("Failed to persist API token hash: %s", write_res.error) + token_hash = await self._persist_api_token_hash_with_warning(token, "Failed to persist API token hash") await self._delete_setting(_SECURITY_API_TOKEN_KEY) - try: - os.environ["MAJOOR_API_TOKEN_HASH"] = token_hash - os.environ["MJR_API_TOKEN_HASH"] = token_hash - except Exception: - pass + self._set_api_token_env(token, token_hash, include_plain=False) return token token_hash = await self._get_stored_api_token_hash_locked() @@ -127,34 +116,39 @@ class AppSettings: # Hash-only persistence means we cannot recover the previous plaintext token. # Regenerate a fresh token so authenticated clients can bootstrap automatically. token = self._generate_api_token() - new_hash = self._hash_api_token(token) - write_res = await self._write_setting(_SECURITY_API_TOKEN_HASH_KEY, new_hash) - if not write_res.ok: - logger.warning("Failed to persist regenerated API token hash: %s", write_res.error) - try: - os.environ["MAJOOR_API_TOKEN"] = token - os.environ["MJR_API_TOKEN"] = token - os.environ["MAJOOR_API_TOKEN_HASH"] = new_hash - os.environ["MJR_API_TOKEN_HASH"] = new_hash - except Exception: - pass + new_hash = await self._persist_api_token_hash_with_warning(token, "Failed to persist regenerated API token hash") + self._set_api_token_env(token, new_hash, include_plain=True) return token token = self._generate_api_token() + token_hash = await self._persist_api_token_hash_with_warning(token, "Failed to persist auto-generated API token hash") + await self._delete_setting(_SECURITY_API_TOKEN_KEY) + self._set_api_token_env(token, token_hash, include_plain=True) + return token + + def _env_api_token(self) -> str: + try: + return (os.environ.get("MAJOOR_API_TOKEN") or os.environ.get("MJR_API_TOKEN") or "").strip() + except Exception: + return "" + + async def _persist_api_token_hash_with_warning(self, token: str, warning_message: str) -> str: token_hash = self._hash_api_token(token) write_res = await self._write_setting(_SECURITY_API_TOKEN_HASH_KEY, token_hash) if not write_res.ok: - logger.warning("Failed to persist auto-generated API token hash: %s", write_res.error) - await self._delete_setting(_SECURITY_API_TOKEN_KEY) + logger.warning("%s: %s", warning_message, write_res.error) + return token_hash + @staticmethod + def _set_api_token_env(token: str, token_hash: str, *, include_plain: bool) -> None: try: - os.environ["MAJOOR_API_TOKEN"] = token - os.environ["MJR_API_TOKEN"] = token + if include_plain: + os.environ["MAJOOR_API_TOKEN"] = token + os.environ["MJR_API_TOKEN"] = token os.environ["MAJOOR_API_TOKEN_HASH"] = token_hash os.environ["MJR_API_TOKEN_HASH"] = token_hash except Exception: pass - return token async def _read_setting(self, key: str) -> Optional[str]: result = await self._db.aquery("SELECT value FROM metadata WHERE key = ?", (key,)) @@ -200,54 +194,75 @@ class AppSettings: return await self._get_security_prefs_locked(include_secret=include_secret) async def set_security_prefs(self, prefs: Mapping[str, Any]) -> Result[dict[str, Any]]: + to_write = self._extract_security_prefs_to_write(prefs) + token_in_payload = self._extract_token_from_prefs_payload(prefs) + if not to_write and token_in_payload is None: + return Result.Err("INVALID_INPUT", "No security settings provided") + async with self._lock: + write_err = await self._persist_security_pref_flags(to_write) + if write_err is not None: + return write_err + if token_in_payload is not None: + token_err = await self._persist_security_api_token(token_in_payload) + if token_err is not None: + return token_err + await self._warn_if_bump_fails("Failed to bump settings version") + return Result.Ok(await self._get_security_prefs_locked(include_secret=False)) + + def _extract_security_prefs_to_write(self, prefs: Mapping[str, Any]) -> dict[str, bool]: to_write: dict[str, bool] = {} for key in _SECURITY_PREFS_INFO: if key in prefs: to_write[key] = parse_bool(prefs[key], False) - token_in_payload = ( - prefs.get("api_token") - if isinstance(prefs, Mapping) and "api_token" in prefs - else prefs.get("apiToken") - if isinstance(prefs, Mapping) and "apiToken" in prefs - else None - ) - if not to_write and token_in_payload is None: - return Result.Err("INVALID_INPUT", "No security settings provided") - async with self._lock: - for key, value in to_write.items(): - res = await self._write_setting(key, "1" if value else "0") - if not res.ok: - return Result.Err("DB_ERROR", res.error or f"Failed to persist {key}") - try: - info = _SECURITY_PREFS_INFO.get(key) or {} - env_var = str(info.get("env") or "").strip() - if env_var: - os.environ[env_var] = "1" if value else "0" - except Exception: - pass - if token_in_payload is not None: - token = str(token_in_payload or "").strip() - if not token: - token = self._generate_api_token() - token_hash = self._hash_api_token(token) - res = await self._write_setting(_SECURITY_API_TOKEN_HASH_KEY, token_hash) - if not res.ok: - return Result.Err("DB_ERROR", res.error or "Failed to persist api_token") - await self._delete_setting(_SECURITY_API_TOKEN_KEY) - try: - os.environ["MAJOOR_API_TOKEN"] = token - os.environ["MJR_API_TOKEN"] = token - os.environ["MAJOOR_API_TOKEN_HASH"] = token_hash - os.environ["MJR_API_TOKEN_HASH"] = token_hash - except Exception: - pass - bump = await self._bump_settings_version_locked() - if not bump.ok: - try: - logger.warning("Failed to bump settings version: %s", bump.error) - except Exception: - pass - return Result.Ok(await self._get_security_prefs_locked(include_secret=False)) + return to_write + + def _extract_token_from_prefs_payload(self, prefs: Mapping[str, Any]) -> Any: + if isinstance(prefs, Mapping) and "api_token" in prefs: + return prefs.get("api_token") + if isinstance(prefs, Mapping) and "apiToken" in prefs: + return prefs.get("apiToken") + return None + + async def _persist_security_pref_flags(self, to_write: Mapping[str, bool]) -> Optional[Result[dict[str, Any]]]: + for key, value in to_write.items(): + res = await self._write_setting(key, "1" if value else "0") + if not res.ok: + return Result.Err("DB_ERROR", res.error or f"Failed to persist {key}") + self._set_security_pref_env_var(key, value) + return None + + def _set_security_pref_env_var(self, key: str, value: bool) -> None: + try: + info = _SECURITY_PREFS_INFO.get(key) or {} + env_var = str(info.get("env") or "").strip() + if env_var: + os.environ[env_var] = "1" if value else "0" + except Exception: + return + + async def _persist_security_api_token(self, token_payload: Any) -> Optional[Result[dict[str, Any]]]: + token = str(token_payload or "").strip() or self._generate_api_token() + token_hash = self._hash_api_token(token) + res = await self._write_setting(_SECURITY_API_TOKEN_HASH_KEY, token_hash) + if not res.ok: + return Result.Err("DB_ERROR", res.error or "Failed to persist api_token") + await self._delete_setting(_SECURITY_API_TOKEN_KEY) + try: + os.environ["MAJOOR_API_TOKEN"] = token + os.environ["MJR_API_TOKEN"] = token + os.environ["MAJOOR_API_TOKEN_HASH"] = token_hash + os.environ["MJR_API_TOKEN_HASH"] = token_hash + except Exception: + pass + return None + + async def _warn_if_bump_fails(self, message: str) -> None: + bump = await self._bump_settings_version_locked() + if not bump.ok: + try: + logger.warning("%s: %s", message, bump.error) + except Exception: + return async def rotate_api_token(self) -> Result[dict[str, str]]: async with self._lock: @@ -322,24 +337,36 @@ class AppSettings: """Return the configured media probe backend mode.""" async with self._lock: current_version = await self._get_settings_version() - cached = self._cache.get(_PROBE_BACKEND_KEY) + cached = self._cached_probe_backend(current_version) if cached: - try: - ts = float(self._cache_at.get(_PROBE_BACKEND_KEY) or 0.0) - except Exception: - ts = 0.0 - cached_ver = int(self._cache_version.get(_PROBE_BACKEND_KEY) or 0) - if cached_ver == int(current_version or 0) and ts and (time.monotonic() - ts) < self._cache_ttl_s: - return cached + return cached mode_raw = await self._read_setting(_PROBE_BACKEND_KEY) mode = str(mode_raw or "").strip().lower() if mode not in _VALID_PROBE_MODES: mode = self._default_probe_mode - self._cache[_PROBE_BACKEND_KEY] = mode - self._cache_at[_PROBE_BACKEND_KEY] = time.monotonic() - self._cache_version[_PROBE_BACKEND_KEY] = int(current_version or 0) + self._store_probe_backend_cache(mode, current_version) return mode + def _cached_probe_backend(self, current_version: int) -> str: + cached = str(self._cache.get(_PROBE_BACKEND_KEY) or "") + if not cached: + return "" + try: + ts = float(self._cache_at.get(_PROBE_BACKEND_KEY) or 0.0) + except Exception: + ts = 0.0 + cached_ver = int(self._cache_version.get(_PROBE_BACKEND_KEY) or 0) + if cached_ver != int(current_version or 0): + return "" + if not ts or (time.monotonic() - ts) >= self._cache_ttl_s: + return "" + return cached + + def _store_probe_backend_cache(self, mode: str, current_version: int) -> None: + self._cache[_PROBE_BACKEND_KEY] = mode + self._cache_at[_PROBE_BACKEND_KEY] = time.monotonic() + self._cache_version[_PROBE_BACKEND_KEY] = int(current_version or 0) + async def set_probe_backend(self, mode: str) -> Result[str]: """Persist the media probe backend mode and bump the settings version.""" normalized = (mode or "").strip().lower() @@ -372,33 +399,59 @@ class AppSettings: """ async with self._lock: current_version = await self._get_settings_version() - out: dict[str, bool] = {} - defaults = { - "image": self._default_metadata_fallback_image, - "media": self._default_metadata_fallback_media, - } - key_map = { - "image": _METADATA_FALLBACK_IMAGE_KEY, - "media": _METADATA_FALLBACK_MEDIA_KEY, - } - for logical_key, storage_key in key_map.items(): - cached = self._cache.get(storage_key) - if cached is not None: - try: - ts = float(self._cache_at.get(storage_key) or 0.0) - except Exception: - ts = 0.0 - cached_ver = int(self._cache_version.get(storage_key) or 0) - if cached_ver == int(current_version or 0) and ts and (time.monotonic() - ts) < self._cache_ttl_s: - out[logical_key] = parse_bool(cached, defaults[logical_key]) - continue - raw = await self._read_setting(storage_key) - parsed = parse_bool(raw, defaults[logical_key]) if raw is not None else defaults[logical_key] - out[logical_key] = parsed - self._cache[storage_key] = "1" if parsed else "0" - self._cache_at[storage_key] = time.monotonic() - self._cache_version[storage_key] = int(current_version or 0) - return out + return await self._read_metadata_fallback_prefs_locked(current_version) + + async def _read_metadata_fallback_prefs_locked(self, current_version: int) -> dict[str, bool]: + out: dict[str, bool] = {} + defaults = { + "image": self._default_metadata_fallback_image, + "media": self._default_metadata_fallback_media, + } + key_map = { + "image": _METADATA_FALLBACK_IMAGE_KEY, + "media": _METADATA_FALLBACK_MEDIA_KEY, + } + for logical_key, storage_key in key_map.items(): + cached_pref = self._cached_metadata_fallback_pref(storage_key, defaults[logical_key], current_version) + if cached_pref is not None: + out[logical_key] = cached_pref + continue + parsed = await self._read_and_cache_metadata_fallback_pref(storage_key, defaults[logical_key], current_version) + out[logical_key] = parsed + return out + + def _cached_metadata_fallback_pref( + self, + storage_key: str, + default: bool, + current_version: int, + ) -> Optional[bool]: + cached = self._cache.get(storage_key) + if cached is None: + return None + try: + ts = float(self._cache_at.get(storage_key) or 0.0) + except Exception: + ts = 0.0 + cached_ver = int(self._cache_version.get(storage_key) or 0) + if cached_ver != int(current_version or 0): + return None + if not ts or (time.monotonic() - ts) >= self._cache_ttl_s: + return None + return parse_bool(cached, default) + + async def _read_and_cache_metadata_fallback_pref( + self, + storage_key: str, + default: bool, + current_version: int, + ) -> bool: + raw = await self._read_setting(storage_key) + parsed = parse_bool(raw, default) if raw is not None else default + self._cache[storage_key] = "1" if parsed else "0" + self._cache_at[storage_key] = time.monotonic() + self._cache_version[storage_key] = int(current_version or 0) + return parsed async def set_metadata_fallback_prefs( self, @@ -409,19 +462,14 @@ class AppSettings: """ Persist metadata fallback preferences and bump settings version. """ - to_write: dict[str, bool] = {} - if image is not None: - to_write[_METADATA_FALLBACK_IMAGE_KEY] = parse_bool(image, self._default_metadata_fallback_image) - if media is not None: - to_write[_METADATA_FALLBACK_MEDIA_KEY] = parse_bool(media, self._default_metadata_fallback_media) + to_write = self._normalize_metadata_fallback_write_payload(image=image, media=media) if not to_write: return Result.Err("INVALID_INPUT", "No metadata fallback settings provided") async with self._lock: - for key, value in to_write.items(): - res = await self._write_setting(key, "1" if value else "0") - if not res.ok: - return Result.Err("DB_ERROR", res.error or f"Failed to persist {key}") + write_error = await self._write_metadata_fallback_payload(to_write) + if write_error: + return write_error bump = await self._bump_settings_version_locked() if not bump.ok: @@ -436,12 +484,33 @@ class AppSettings: self._cache_at[key] = time.monotonic() self._cache_version[key] = current_version - return Result.Ok( - { - "image": parse_bool(self._cache.get(_METADATA_FALLBACK_IMAGE_KEY), self._default_metadata_fallback_image), - "media": parse_bool(self._cache.get(_METADATA_FALLBACK_MEDIA_KEY), self._default_metadata_fallback_media), - } - ) + return Result.Ok(self._current_metadata_fallback_prefs_from_cache()) + + def _normalize_metadata_fallback_write_payload( + self, + *, + image: Any, + media: Any, + ) -> dict[str, bool]: + to_write: dict[str, bool] = {} + if image is not None: + to_write[_METADATA_FALLBACK_IMAGE_KEY] = parse_bool(image, self._default_metadata_fallback_image) + if media is not None: + to_write[_METADATA_FALLBACK_MEDIA_KEY] = parse_bool(media, self._default_metadata_fallback_media) + return to_write + + async def _write_metadata_fallback_payload(self, to_write: dict[str, bool]) -> Optional[Result[Any]]: + for key, value in to_write.items(): + res = await self._write_setting(key, "1" if value else "0") + if not res.ok: + return Result.Err("DB_ERROR", res.error or f"Failed to persist {key}") + return None + + def _current_metadata_fallback_prefs_from_cache(self) -> dict[str, bool]: + return { + "image": parse_bool(self._cache.get(_METADATA_FALLBACK_IMAGE_KEY), self._default_metadata_fallback_image), + "media": parse_bool(self._cache.get(_METADATA_FALLBACK_MEDIA_KEY), self._default_metadata_fallback_media), + } async def get_output_directory(self) -> str | None: """Return persisted output directory override, or None when unset.""" @@ -457,105 +526,100 @@ class AppSettings: async def set_output_directory(self, path: str) -> Result[str]: """Persist output directory override and bump settings version.""" - def _apply_comfy_output_directory(target: str) -> None: - """ - Best-effort runtime sync with ComfyUI output directory. - """ - try: - import folder_paths # type: ignore - except Exception: - return - try: - normalized_target = str(target or "").strip() - if not normalized_target: - return - setter = getattr(folder_paths, "set_output_directory", None) - if callable(setter): - try: - setter(normalized_target) - return - except Exception: - pass - # Fallback for Comfy versions exposing a module-level variable. - try: - setattr(folder_paths, "output_directory", normalized_target) - except Exception: - pass - except Exception: - return - - def _get_current_comfy_output_directory() -> str: - try: - import folder_paths # type: ignore - getter = getattr(folder_paths, "get_output_directory", None) - if callable(getter): - cur = str(getter() or "").strip() - if cur: - return cur - except Exception: - pass - return "" - normalized = str(path or "").strip() if not normalized: - async with self._lock: - result = await self._delete_setting(_OUTPUT_DIRECTORY_KEY) - if not result.ok: - return Result.Err("DB_ERROR", result.error or "Failed to clear output directory") - try: - os.environ.pop("MAJOOR_OUTPUT_DIRECTORY", None) - os.environ.pop("MJR_AM_OUTPUT_DIRECTORY", None) - except Exception: - pass - try: - original = str(os.environ.get(_ORIGINAL_OUTPUT_DIRECTORY_ENV) or "").strip() - except Exception: - original = "" - try: - restore_target = original or str(OUTPUT_ROOT or "").strip() - if restore_target: - _apply_comfy_output_directory(restore_target) - except Exception: - pass - try: - os.environ.pop(_ORIGINAL_OUTPUT_DIRECTORY_ENV, None) - except Exception: - pass - bump = await self._bump_settings_version_locked() - if not bump.ok: - try: - logger.warning("Failed to bump settings version: %s", bump.error) - except Exception: - pass - return Result.Ok("") + return await self._clear_output_directory_override() try: from pathlib import Path p = Path(normalized).expanduser().resolve() normalized = str(p) except Exception: return Result.Err("INVALID_INPUT", "Invalid output directory path") + return await self._persist_output_directory_override(normalized) + + async def _clear_output_directory_override(self) -> Result[str]: + async with self._lock: + result = await self._delete_setting(_OUTPUT_DIRECTORY_KEY) + if not result.ok: + return Result.Err("DB_ERROR", result.error or "Failed to clear output directory") + self._clear_output_directory_env_vars() + restore_target = self._restore_output_directory_target() + if restore_target: + self._apply_comfy_output_directory(restore_target) + self._clear_original_output_directory_env() + await self._warn_if_bump_fails("Failed to bump settings version") + return Result.Ok("") + + async def _persist_output_directory_override(self, normalized: str) -> Result[str]: async with self._lock: result = await self._write_setting(_OUTPUT_DIRECTORY_KEY, normalized) if not result.ok: return Result.Err("DB_ERROR", result.error or "Failed to persist output directory") - try: - if not str(os.environ.get(_ORIGINAL_OUTPUT_DIRECTORY_ENV) or "").strip(): - current = _get_current_comfy_output_directory() - if current: - os.environ[_ORIGINAL_OUTPUT_DIRECTORY_ENV] = current - os.environ["MAJOOR_OUTPUT_DIRECTORY"] = normalized - os.environ["MJR_AM_OUTPUT_DIRECTORY"] = normalized - except Exception: - pass - try: - _apply_comfy_output_directory(normalized) - except Exception: - pass - bump = await self._bump_settings_version_locked() - if not bump.ok: - try: - logger.warning("Failed to bump settings version: %s", bump.error) - except Exception: - pass + self._set_output_directory_env_vars(normalized) + self._apply_comfy_output_directory(normalized) + await self._warn_if_bump_fails("Failed to bump settings version") return Result.Ok(normalized) + def _clear_output_directory_env_vars(self) -> None: + try: + os.environ.pop("MAJOOR_OUTPUT_DIRECTORY", None) + os.environ.pop("MJR_AM_OUTPUT_DIRECTORY", None) + except Exception: + return + + def _restore_output_directory_target(self) -> str: + try: + original = str(os.environ.get(_ORIGINAL_OUTPUT_DIRECTORY_ENV) or "").strip() + except Exception: + original = "" + return original or str(OUTPUT_ROOT or "").strip() + + def _clear_original_output_directory_env(self) -> None: + try: + os.environ.pop(_ORIGINAL_OUTPUT_DIRECTORY_ENV, None) + except Exception: + return + + def _set_output_directory_env_vars(self, normalized: str) -> None: + try: + if not str(os.environ.get(_ORIGINAL_OUTPUT_DIRECTORY_ENV) or "").strip(): + current = self._get_current_comfy_output_directory() + if current: + os.environ[_ORIGINAL_OUTPUT_DIRECTORY_ENV] = current + os.environ["MAJOOR_OUTPUT_DIRECTORY"] = normalized + os.environ["MJR_AM_OUTPUT_DIRECTORY"] = normalized + except Exception: + return + + def _apply_comfy_output_directory(self, target: str) -> None: + try: + import folder_paths # type: ignore + except Exception: + return + normalized_target = str(target or "").strip() + if not normalized_target: + return + setter = getattr(folder_paths, "set_output_directory", None) + if callable(setter): + try: + setter(normalized_target) + return + except Exception: + pass + try: + setattr(folder_paths, "output_directory", normalized_target) + except Exception: + return + + def _get_current_comfy_output_directory(self) -> str: + try: + import folder_paths # type: ignore + + getter = getattr(folder_paths, "get_output_directory", None) + if callable(getter): + cur = str(getter() or "").strip() + if cur: + return cur + except Exception: + return "" + return "" diff --git a/mjr_am_backend/shared.py b/mjr_am_backend/shared.py index fc7e5f1f..b1a7b664 100644 --- a/mjr_am_backend/shared.py +++ b/mjr_am_backend/shared.py @@ -1,4 +1,4 @@ -"""Backend-facing alias for shared utilities. +"""Backend-facing alias for shared utilities. This file exists because ComfyUI can load custom nodes either: - as a pseudo-package via file paths under `custom_nodes/`, or @@ -57,4 +57,3 @@ __all__ = _root_shared.__all__ + [ "EXTENSIONS", "sanitize_error_message", ] - diff --git a/mjr_am_backend/tool_detect.py b/mjr_am_backend/tool_detect.py index 0504f9a1..e6963b5a 100644 --- a/mjr_am_backend/tool_detect.py +++ b/mjr_am_backend/tool_detect.py @@ -1,4 +1,4 @@ -""" +""" Tool detection helpers for ExifTool and FFprobe. Cached detection to avoid repeated subprocess calls. """ @@ -178,4 +178,3 @@ def reset_tool_cache(): global _TOOL_CACHE, _TOOL_VERSIONS _TOOL_CACHE = {"exiftool": None, "ffprobe": None} _TOOL_VERSIONS = {"exiftool": None, "ffprobe": None} - diff --git a/mjr_am_shared/scan_throttle.py b/mjr_am_shared/scan_throttle.py index c566a4bb..4d1ce9a1 100644 --- a/mjr_am_shared/scan_throttle.py +++ b/mjr_am_shared/scan_throttle.py @@ -1,4 +1,4 @@ -""" +""" Throttle helpers so background scans can skip directories that were just indexed. """ from __future__ import annotations @@ -109,4 +109,3 @@ def _cleanup_locked(now: float) -> None: keys_to_remove = [k for k, v in _RECENT_SCAN_TIMES.items() if v < cutoff] for k in keys_to_remove: _RECENT_SCAN_TIMES.pop(k, None) - diff --git a/scripts/check_no_bom.py b/scripts/check_no_bom.py new file mode 100644 index 00000000..b6351a63 --- /dev/null +++ b/scripts/check_no_bom.py @@ -0,0 +1,82 @@ +"""Fail if any Python source file is encoded with UTF-8 BOM.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + + +UTF8_BOM = b"\xef\xbb\xbf" +SKIP_DIRS = { + ".git", + ".venv", + "node_modules", + "__pycache__", + ".pytest_cache", +} + + +def _iter_python_files(root: Path): + for path in root.rglob("*.py"): + if any(part in SKIP_DIRS for part in path.parts): + continue + yield path + + +def _has_utf8_bom(path: Path) -> bool: + try: + with path.open("rb") as fh: + return fh.read(3) == UTF8_BOM + except OSError: + return False + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Check that Python files do not contain a UTF-8 BOM." + ) + parser.add_argument( + "paths", + nargs="*", + default=["."], + help="Files or directories to scan (default: repository root).", + ) + return parser.parse_args() + + +def main() -> int: + args = _parse_args() + cwd = Path.cwd().resolve() + offenders: list[Path] = [] + + for raw in args.paths: + target = Path(raw) + if not target.exists(): + continue + + if target.is_file(): + if target.suffix == ".py" and _has_utf8_bom(target): + offenders.append(target) + continue + + for path in _iter_python_files(target): + if _has_utf8_bom(path): + offenders.append(path) + + offenders = sorted({p.resolve() for p in offenders}) + if offenders: + print(f"Found UTF-8 BOM in {len(offenders)} Python file(s):") + for path in offenders: + try: + rel = path.relative_to(cwd) + except ValueError: + rel = path + print(f"- {rel.as_posix()}") + return 1 + + print("No UTF-8 BOM detected in Python files.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/install-requirements.py b/scripts/install-requirements.py index 902e7f87..2ba9d90f 100644 --- a/scripts/install-requirements.py +++ b/scripts/install-requirements.py @@ -1,4 +1,4 @@ -"""Install Majoor Assets Manager runtime dependencies explicitly. +"""Install Majoor Assets Manager runtime dependencies explicitly. Usage: python scripts/install-requirements.py diff --git a/tests/compat/test_import_smoke.py b/tests/compat/test_import_smoke.py index 9e1e3e82..050ae3e8 100644 --- a/tests/compat/test_import_smoke.py +++ b/tests/compat/test_import_smoke.py @@ -1,4 +1,4 @@ -import importlib.util +import importlib.util from pathlib import Path diff --git a/tests/conftest.py b/tests/conftest.py index d64b1e5d..b2103e87 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,4 +1,4 @@ -import sys +import sys import pytest_asyncio from .repo_root import REPO_ROOT @@ -22,4 +22,3 @@ async def services(tmp_path): await svc.get("db").aclose() except Exception: pass - diff --git a/tests/core/test_routes.py b/tests/core/test_routes.py index f5e7e846..33a536a8 100644 --- a/tests/core/test_routes.py +++ b/tests/core/test_routes.py @@ -1,4 +1,4 @@ -""" +""" Test routes registration without starting full ComfyUI. """ import sys @@ -37,4 +37,3 @@ async def test_routes(): if __name__ == "__main__": import asyncio asyncio.run(test_routes()) - diff --git a/tests/database/test_schema_heal.py b/tests/database/test_schema_heal.py index 3f60c64c..b6e6b4ef 100644 --- a/tests/database/test_schema_heal.py +++ b/tests/database/test_schema_heal.py @@ -1,4 +1,4 @@ -import pytest +import pytest import sqlite3 @@ -132,5 +132,3 @@ async def test_sqlite_query_missing_column_triggers_self_heal(tmp_path): await db.aclose() except Exception: pass - - diff --git a/tests/metadata/test_extractors.py b/tests/metadata/test_extractors.py index d8627020..61dc985f 100644 --- a/tests/metadata/test_extractors.py +++ b/tests/metadata/test_extractors.py @@ -1,4 +1,4 @@ -import pytest +import pytest from mjr_am_backend.features.metadata.parsing_utils import looks_like_comfyui_prompt_graph @@ -12,4 +12,3 @@ async def test_prompt_graph_accepts_colon_ids(): } assert looks_like_comfyui_prompt_graph(prompt_graph) - diff --git a/tests/parser/test_geninfo_flux.py b/tests/parser/test_geninfo_flux.py index 68450848..6c67bf89 100644 --- a/tests/parser/test_geninfo_flux.py +++ b/tests/parser/test_geninfo_flux.py @@ -1,4 +1,4 @@ -from mjr_am_backend.features.geninfo.parser import parse_geninfo_from_prompt +from mjr_am_backend.features.geninfo.parser import parse_geninfo_from_prompt def test_geninfo_flux_standard_advanced_sampler(): """ @@ -8,7 +8,7 @@ def test_geninfo_flux_standard_advanced_sampler(): # 1. Output Sink "100": { "class_type": "SaveImage", - "inputs": {"images": [90, 0]} + "inputs": {"images": [90, 0]} }, # 2. VAE Decode "90": { @@ -69,7 +69,7 @@ def test_geninfo_flux_standard_advanced_sampler(): "10": { "class_type": "VAELoader", "inputs": {"vae_name": "ae.safetensors"}}, "20": { "class_type": "LoadDiffusionModel", "inputs": {"unet_name": "flux1-dev.safetensors"} }, "25": { "class_type": "DualCLIPLoader", "inputs": {"clip_name1": "t5XXL.fp16.safetensors", "clip_name2": "clip_l.safetensors"} }, - + # 9. Prompt "30": { "class_type": "CLIPTextEncode", @@ -80,20 +80,20 @@ def test_geninfo_flux_standard_advanced_sampler(): }, "15": { "class_type": "EmptyLatentImage", "inputs": {"width": 1024, "height": 1024} } } - + result = parse_geninfo_from_prompt(prompt_graph) assert result.ok, result.error data = result.data assert data is not None - + # Check sampler params assert data["steps"]["value"] == 25 assert data["seed"]["value"] == 99999 assert data["cfg"]["value"] == 3.5 # Extracted from FluxGuidance - + # Check Prompt assert data["positive"]["value"] == "A photo of a cat" - + def test_geninfo_flux_gguf_fp8_variant(): """ Test variant with UnetLoaderGGUF and potentially distinct structure. @@ -146,20 +146,18 @@ def test_geninfo_flux_gguf_fp8_variant(): result = parse_geninfo_from_prompt(prompt_graph) assert result.ok data = result.data - + # Parser maps UNET/Diffusion models to 'checkpoint' field for compatibility assert data["checkpoint"]["name"] == "flux1-dev-Q8_0" - - # CFG might be missing if no FluxGuidance? + + # CFG might be missing if no FluxGuidance? # Or BasicGuider doesn't have CFG. - + # The parser currently doesn't default CFG if missing. # But it should extract seed, steps. assert data["steps"]["value"] == 20 assert data["seed"]["value"] == 12345 - + # Verify Prompt was extracted # "GGUF cat" via BasicGuider -> conditioning assert data["positive"]["value"] == "GGUF cat" - - diff --git a/tests/rating_tags/test_rating_tags_import.py b/tests/rating_tags/test_rating_tags_import.py index afe75d4a..82db898c 100644 --- a/tests/rating_tags/test_rating_tags_import.py +++ b/tests/rating_tags/test_rating_tags_import.py @@ -1,4 +1,4 @@ -""" +""" Tests for importing ratings and tags logic, ensuring no data loss or incorrect overwrites. """ import pytest @@ -68,4 +68,3 @@ async def test_does_not_override_existing_db_rating_tags(tmp_path: Path): assert row["tags_text"] == "keep" finally: await db.aclose() - diff --git a/tests/security/test_write_access_guard.py b/tests/security/test_write_access_guard.py index 2ad31517..b100afd1 100644 --- a/tests/security/test_write_access_guard.py +++ b/tests/security/test_write_access_guard.py @@ -1,4 +1,4 @@ -import pytest +import pytest from mjr_am_backend.routes.core.security import _check_write_access @@ -92,4 +92,3 @@ async def test_require_auth_without_token_blocks(monkeypatch): res = _check_write_access(peer_ip="127.0.0.1", headers={}) assert not res.ok assert res.code == "AUTH_REQUIRED" -