diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 17676b8..3a5210c 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -97,6 +97,11 @@ jobs: PM_COMFY_ROOT: ${{ github.workspace }}/.ci-comfyui PM_TEST_PYTHON: python run: npm run --prefix tools/vendor test:comfyui + - name: Real cached and headless execution history + env: + PM_COMFY_ROOT: ${{ github.workspace }}/.ci-comfyui + PM_TEST_PYTHON: python + run: npm run --prefix tools/vendor test:executions - name: Independent and combined pack compatibility env: PM_COMFY_ROOT: ${{ github.workspace }}/.ci-comfyui diff --git a/README.md b/README.md index 1e1db44..e259362 100644 --- a/README.md +++ b/README.md @@ -531,7 +531,7 @@ Keep your database optimized with built-in maintenance tools: ### Database Location -By default, the database is saved as `example_prompts.db` in the node directory. This file contains all your prompts and linked images and can be backed up or shared. +By default, the database is saved as `prompts.db` in the extension directory. It contains your prompt library and image links. See the [database path compatibility note](#database-path-compatibility-211) before upgrading a deployment with a custom path. ### Searching Prompts @@ -773,6 +773,32 @@ Relative paths are resolved against the extension directory. Absolute paths are } ``` +### Database Path Compatibility (#211) + +PromptManager now honors the `database.default_path` setting in `config.json` +when loaded as a ComfyUI custom node. Relative paths resolve against the extension +directory, not ComfyUI's working directory; absolute paths are used as-is. Earlier +versions could ignore this setting and use `/prompts.db` instead. +Honoring a different configured path can therefore show an empty or different +library. PromptManager does not automatically move or merge database files. + +Before upgrading or changing this setting: + +1. Identify the database your current installation actually uses and confirm its + library contents. Check both the configured path and the prior fallback + `/prompts.db`; keep track of which existing file you intend to use. +2. Preserve a consistent backup of the active database with the dashboard's + Backup/Restore backup action or SQLite's online backup API. Do not copy only + the main file of a live SQLite database: committed data may still be in its + `-wal` file. For a filesystem copy, stop ComfyUI cleanly first and preserve any + remaining `-wal` file together with the main database. +3. Set `database.default_path` to the existing intended database file. An absolute + path makes the selection explicit; do not move or overwrite either database + while choosing the path. +4. Restart ComfyUI, then verify that the expected library appears and the active + database path resolves to the file you selected. Keep the backup until both + checks succeed. + ### Gallery & Monitoring Settings Configure the automatic image detection and gallery system: @@ -844,7 +870,7 @@ print(f"Average rating: {info['average_rating']}") # Create backup db.model.backup_database("backup_prompts.db") -# The database file can be copied directly for backup +# The online backup includes committed data still in the WAL. ``` ## Development diff --git a/TESTING.md b/TESTING.md index dc1f5c9..2fed74e 100644 --- a/TESTING.md +++ b/TESTING.md @@ -10,7 +10,7 @@ reviewed. Add a regression test for each bug; prove it fails before the fix. | Python tests | Domain logic, SQLite, API handlers, filesystem checks | ComfyUI host and ML runtimes where needed | | JavaScript tests | Frontend modules and controller methods | DOM/network fixtures for focused edge cases | | Viewer checks | Production Viewer.js and frontend methods in Chromium desktop/touch | Controlled images and page fixtures | -| Real ComfyUI | Pinned CPU host, actual extension, node definitions, widgets, routes and SQLite | Inference is not queued; browser startup settings and host preference writes are isolated | +| Real ComfyUI | Pinned CPU host, actual extension, node definitions, widgets, routes, SQLite and real execution lifecycle | Stock UI cases isolate host writes; the execution suite queues deterministic test-only nodes without models | | Page workflows | Production HTML, startup, click handlers, HTTP routes, SQLite, PNG files, thumbnails | ComfyUI host services and the tagging model | Page workflows cover save/edit/reload, duplicate-edit rejection, gallery image @@ -106,7 +106,7 @@ and failure screenshots/traces. Open traces with Playwright's `show-trace` tool. The real-host suite covers both Classic and Nodes 2.0 startup, switching both ways, actual database prompt selection, execution-text serialization, workflow reload, widget/style cleanup, node schemas and standalone-page opening. It does -not queue inference or certify every existing feature, arbitrary host versions, +not certify every existing feature, arbitrary host versions, GPU/model changes, cross-pack installation or visual design changes. Those still need targeted validation. @@ -197,3 +197,61 @@ Viewer, including nested modal focus, mobile layouts and per-role sending. `tests/comfyui/prompt-sender.spec.mjs` sends through the actual workspace UI into real Classic and Nodes 2.0 widgets and verifies workflow serialization without loading a workflow or queueing generation. + +## Durable execution ledger + +`npm run --prefix tools/vendor test:executions` adds ten real execution cases +(five each in Classic and Nodes 2.0) to the existing 78 stock host cases and six +pack compatibility cases. Run these suites serially because they share report +cleanup and owned host resources. Use the same pinned disposable `PM_COMFY_ROOT`, +`PM_TEST_PYTHON` and Playwright browser installation as `test:comfyui`: + +```sh +PM_COMFY_ROOT=/path/to/disposable/ComfyUI \ +PM_TEST_PYTHON=/path/to/cpu-env/bin/python \ +PLAYWRIGHT_BROWSERS_PATH=/path/to/playwright-browsers \ +npm run --prefix tools/vendor test:executions +``` + +This suite reserves `http://127.0.0.1:8196/` and never attaches to an existing +server. Each renderer worker launches a CPU/offline host with independent +configuration, database, user and output directories. The shared launcher installs +tracked snapshots of PromptManager and a tiny temporary Git-tracked fixture pack. +Only that owned test pack provides cached output, deliberate failure, an +interruptible queue gate and gate-control routes; production has no test nodes or +endpoints. Cleanup always releases the gate, stops the host and removes both +owned directories. Logs and Playwright failures appear under +`tools/vendor/test-results/executions/` and the existing CI artifact upload. + +The request-only case creates no browser or WebSocket client. Other cases verify +real `execution_cached` node lists after a failed job, both renderer switches and +workflow reload, repeated public UUIDs, interruption, graph rejection, and queued +removal before dequeue. Actual host history is checked alongside finite polling +of SQLite after the background writer. The read helper uses SQLite read-only +mode, exact UUID job/session filters and a 256-row limit; absent ledger schema is +explicit, while missing/corrupt databases and other read errors fail the suite. +Sensitive sentinels in graph inputs, extra data and error messages must never +appear in any stored row. + +For a meaningful historical failure, create a separate disposable detached +checkout at `fdba6fd0126a8a40cbd6f3e50ed4c5de8991efd4` and pass its path to +`launchExecutionHost({ repository })`. That base also predates the #211 config +import fix: after confirming an actual headless job succeeded, use +`readExecutionRuns` on its explicitly constructed owned path +`host.stateDirectory/custom_nodes/promptmanager/prompts.db`. It reports +`schemaAvailable: false` and empty rows, failing the expected succeeded-row +assertion. This historical inspection stays outside the ordinary suite, which +requires the configured database path. The launcher also accepts +`PM_EXECUTION_REPOSITORY`; omit it for the implementation worktree. Stage new +files first because snapshots use `git ls-files`. Never reset the implementation +checkout or point this harness at a live ComfyUI checkout. + +The package-qualified config import also corrects #211: the configured database +path is honored when PromptManager runs as a ComfyUI custom node. Follow the +[database path compatibility and safe backup sequence](README.md#database-path-compatibility-211) +before restarting an existing deployment. + +The ledger records only new dequeued jobs and observed outcomes. Existing +counters, prompts, pairs and image links are preserved, with no reconstructed old +runs. Prompt contexts, output associations and the chronological timeline remain +future work under #209. Deploying the Python changes requires a ComfyUI restart. diff --git a/__init__.py b/__init__.py index c1e1dd5..2663b6a 100644 --- a/__init__.py +++ b/__init__.py @@ -158,6 +158,34 @@ except Exception as e: f"Failed to register usage hook: {e}" ) + +def _history_db_factory() -> "PromptDatabase": + """Initialize history on its worker even if monitoring initialization failed.""" + if _global_db is not None: + return _global_db + from .database.operations import PromptDatabase + + return PromptDatabase() + + +_history_registration = None +try: + from .py.config import server_instance + from .utils.execution_history import register_history_hook + + _history_registration = register_history_hook(server_instance, _history_db_factory) +except Exception: + try: + from .utils.logging_config import get_logger + + get_logger("prompt_manager.init").error("Failed to register execution history") + except Exception: + import logging + + logging.getLogger("prompt_manager.init").error( + "Failed to register execution history" + ) + __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] # Print startup message with loaded tools diff --git a/database/execution_history.py b/database/execution_history.py new file mode 100644 index 0000000..ff4e53c --- /dev/null +++ b/database/execution_history.py @@ -0,0 +1,194 @@ +"""Transactional attempt writes and process-local restore invalidation.""" + +import os +import sqlite3 +import threading +from contextlib import contextmanager +from typing import Callable, Iterator + +try: + from ..utils.execution_history_events import ( + RunFinish, + RunStart, + classify_evidence, + encode_command, + ) + from ..utils.logging_config import get_logger +except ImportError: + from utils.execution_history_events import ( + RunFinish, + RunStart, + classify_evidence, + encode_command, + ) + from utils.logging_config import get_logger + + +_logger = get_logger("prompt_manager.database.execution_history") +_restore_lock = threading.Lock() +_restore_observers: dict[str, dict[object, Callable[[], None]]] = {} + + +def _database_identity(db_path: str) -> str: + return os.path.normcase(os.path.realpath(db_path)) + + +def subscribe_history_restore( + db_path: str, callback: Callable[[], None] +) -> Callable[[], None]: + """Observe successful restores of the canonical database path.""" + identity = _database_identity(db_path) + token = object() + with _restore_lock: + _restore_observers.setdefault(identity, {})[token] = callback + + def unsubscribe() -> None: + with _restore_lock: + observers = _restore_observers.get(identity) + if observers is not None: + observers.pop(token, None) + if not observers: + _restore_observers.pop(identity, None) + + return unsubscribe + + +def notify_history_restore(db_path: str) -> None: + """Notify outside the registry lock; observers cannot fail a restore.""" + with _restore_lock: + callbacks = tuple( + _restore_observers.get(_database_identity(db_path), {}).values() + ) + for callback in callbacks: + try: + callback() + except Exception: + _logger.warning("Execution history restore observer failed", exc_info=True) + + +class HistoryEpochChanged(Exception): + """A command belongs to storage replaced by a completed restore.""" + + def __init__(self, storage_epoch: str): + self.storage_epoch = storage_epoch + super().__init__("Execution history storage epoch changed") + + +class ExecutionHistoryOperations: + """Mixin using the calling thread's existing managed model connection.""" + + def get_execution_epoch(self) -> str: + row = ( + self.model.get_connection() + .execute( + "SELECT storage_epoch FROM execution_history_metadata " + "WHERE singleton_id=1" + ) + .fetchone() + ) + if row is None: + raise sqlite3.DatabaseError("Missing execution history metadata") + return row[0] + + @contextmanager + def _execution_transaction( + self, expected_epoch: str + ) -> Iterator[sqlite3.Connection]: + conn = self.model.get_connection() + if conn.in_transaction: + raise sqlite3.ProgrammingError( + "Execution history requires its own transaction" + ) + timeout = conn.execute("PRAGMA busy_timeout").fetchone()[0] + conn.execute("PRAGMA busy_timeout=250") + try: + conn.execute("BEGIN IMMEDIATE") + try: + epoch = self.get_execution_epoch() + if epoch != expected_epoch: + raise HistoryEpochChanged(epoch) + yield conn + conn.commit() + except BaseException: + conn.rollback() + raise + finally: + conn.execute(f"PRAGMA busy_timeout={timeout}") + + def write_execution_start(self, command: RunStart) -> bool: + encode_command(command) + if type(command) is not RunStart: + raise ValueError("Expected execution start command") + with self._execution_transaction(command.storage_epoch) as conn: + conn.execute( + "INSERT INTO execution_runs " + "(run_id,session_id,queue_item_id,job_id,storage_epoch,record_version," + "observed_dequeue_at,outcome,diagnostics_omitted) " + "VALUES(?,?,?,?,?,1,?,'observed',0) " + "ON CONFLICT(session_id,queue_item_id) DO NOTHING", + ( + command.run_id, + command.session_id, + command.queue_item_id, + command.job_id, + command.storage_epoch, + command.observed_dequeue_at, + ), + ) + return True + + def write_execution_finish(self, command: RunFinish) -> bool: + encode_command(command) + if type(command) is not RunFinish: + raise ValueError("Expected execution finish command") + with self._execution_transaction(command.storage_epoch) as conn: + row = conn.execute( + "SELECT * FROM execution_runs WHERE session_id=? AND queue_item_id=? " + "AND storage_epoch=?", + (command.session_id, command.queue_item_id, command.storage_epoch), + ).fetchone() + if row is None: + return False + completion = classify_evidence(row["job_id"], command.evidence) + if row["outcome"] == "unknown": + return True + if row["outcome"] != "observed": + contradictory = completion.reason == "conflicting_terminal" or ( + completion.outcome in ("succeeded", "failed", "interrupted") + and completion.outcome != row["outcome"] + ) + if contradictory: + conn.execute( + "UPDATE execution_runs SET outcome='unknown'," + "reason='conflicting_terminal' WHERE run_id=?", + (row["run_id"],), + ) + return True + conn.execute( + "UPDATE execution_runs SET observed_callback_at=?," + "executor_started_at=?," + "terminal_event_at=?,outcome=?,reason=?,diagnostics_omitted=?," + "node_id=?," + "node_type=?,exception_type=? WHERE run_id=?", + ( + command.observed_callback_at, + completion.executor_started_at, + completion.terminal_event_at, + completion.outcome, + completion.reason, + int(completion.diagnostics_omitted), + completion.node_id, + completion.node_type, + completion.exception_type, + row["run_id"], + ), + ) + return True + + def get_execution_run(self, run_id: str) -> dict | None: + row = ( + self.model.get_connection() + .execute("SELECT * FROM execution_runs WHERE run_id=?", (run_id,)) + .fetchone() + ) + return dict(row) if row is not None else None diff --git a/database/execution_history_schema.py b/database/execution_history_schema.py new file mode 100644 index 0000000..562648f --- /dev/null +++ b/database/execution_history_schema.py @@ -0,0 +1,225 @@ +"""Versioned additive schema for observed execution attempts.""" + +import re +import sqlite3 +from uuid import uuid4 + +SCHEMA_VERSION = 1 + + +def _uuid_check(column: str) -> str: + return ( + f"typeof({column}) = 'text' AND length({column}) = 36 " + f"AND substr({column}, 9, 1) = '-' AND substr({column}, 14, 1) = '-' " + f"AND substr({column}, 19, 1) = '-' AND substr({column}, 24, 1) = '-' " + f"AND length(replace({column}, '-', '')) = 32 " + f"AND replace({column}, '-', '') NOT GLOB '*[^0-9a-f]*'" + ) + + +def _integer_check(column: str, *, nullable: bool = False) -> str: + expression = f"typeof({column}) = 'integer' AND {column} >= 0" + return f"{column} IS NULL OR ({expression})" if nullable else expression + + +_METADATA_CHECKS = ( + "singleton_id = 1", + "schema_version = 1", + _uuid_check("storage_epoch"), +) +_RUN_CHECKS = ( + *( + _uuid_check(column) + for column in ("run_id", "session_id", "job_id", "storage_epoch") + ), + _integer_check("queue_item_id"), + "record_version = 1", + _integer_check("observed_dequeue_at"), + *( + _integer_check(column, nullable=True) + for column in ( + "observed_callback_at", + "executor_started_at", + "terminal_event_at", + ) + ), + "outcome IN ('observed', 'succeeded', 'failed', 'interrupted', 'unknown')", + "(outcome = 'observed' AND reason IS NULL AND observed_callback_at IS NULL) OR " + "(outcome != 'observed' AND reason IS NOT NULL " + "AND observed_callback_at IS NOT NULL AND " + "((outcome = 'succeeded' AND reason = 'terminal_success') OR " + "(outcome = 'failed' AND reason = 'terminal_error') OR " + "(outcome = 'interrupted' AND reason = 'terminal_interrupted') OR " + "(outcome = 'unknown' AND reason IN ('missing_terminal', " + "'conflicting_terminal', 'invalid_evidence', 'evidence_limit'))))", + "diagnostics_omitted IN (0, 1)", + *( + f"{column} IS NULL OR (typeof({column}) = 'text' " + f"AND length({column}) <= {limit})" + for column, limit in ( + ("node_id", 128), + ("node_type", 128), + ("exception_type", 256), + ) + ), +) +_TABLE_COLUMNS = { + "execution_history_metadata": ( + ("singleton_id", "INTEGER", 1, 1), + ("schema_version", "INTEGER", 1, 0), + ("storage_epoch", "TEXT", 1, 0), + ), + "execution_runs": ( + ("run_id", "TEXT", 1, 1), + ("session_id", "TEXT", 1, 0), + ("queue_item_id", "INTEGER", 1, 0), + ("job_id", "TEXT", 1, 0), + ("storage_epoch", "TEXT", 1, 0), + ("record_version", "INTEGER", 1, 0), + ("observed_dequeue_at", "INTEGER", 1, 0), + ("observed_callback_at", "INTEGER", 0, 0), + ("executor_started_at", "INTEGER", 0, 0), + ("terminal_event_at", "INTEGER", 0, 0), + ("outcome", "TEXT", 1, 0), + ("reason", "TEXT", 0, 0), + ("diagnostics_omitted", "INTEGER", 1, 0), + ("node_id", "TEXT", 0, 0), + ("node_type", "TEXT", 0, 0), + ("exception_type", "TEXT", 0, 0), + ), +} +_TABLE_CHECKS = { + "execution_history_metadata": _METADATA_CHECKS, + "execution_runs": _RUN_CHECKS, +} +_INDEXES = { + "idx_execution_runs_identity": (True, ("session_id", "queue_item_id")), + "idx_execution_runs_observed": (False, ("observed_dequeue_at", "run_id")), +} + + +def _compact(sql: str) -> str: + return "".join( + token if token.startswith("'") else re.sub(r"\s+", "", token).lower() + for token in re.findall(r"'(?:''|[^'])*'|[^']+", sql) + ) + + +def history_schema_error( + conn: sqlite3.Connection, *, allow_missing: bool = True +) -> str | None: + """Reject partial/unsupported history schemas and invalid persisted records.""" + tables = { + row[0]: row[1] + for row in conn.execute("SELECT name,sql FROM sqlite_master WHERE type='table'") + if row[0] in _TABLE_COLUMNS + } + if not tables: + return None if allow_missing else "Missing execution history schema" + if tables.keys() != _TABLE_COLUMNS.keys(): + return "Incomplete execution history schema" + for table, expected in _TABLE_COLUMNS.items(): + columns = tuple( + (row[1], row[2].upper(), row[3], row[5]) + for row in conn.execute(f"PRAGMA table_info({table})") + ) + if columns != expected: + return f"{table} has invalid execution history columns" + sql = _compact(tables[table]) + checks = _TABLE_CHECKS[table] + if sql.count("check(") != len(checks) or any( + f"check({_compact(check)})" not in sql for check in checks + ): + return f"{table} has invalid execution history checks" + if conn.execute(f"PRAGMA foreign_key_list({table})").fetchone(): + return f"{table} has unexpected execution history foreign keys" + for index in conn.execute(f"PRAGMA index_list({table})"): + if not index[2]: + continue + if table == "execution_runs" and index[1] == "idx_execution_runs_identity": + continue + name = index[1].replace('"', '""') + keys = tuple( + row for row in conn.execute(f'PRAGMA index_xinfo("{name}")') if row[5] + ) + primary = tuple(column[0] for column in expected if column[3]) + if ( + index[3] != "pk" + or index[4] + or tuple(key[2] for key in keys) != primary + or any(key[3] or key[4] != "BINARY" for key in keys) + ): + return f"{table} has unexpected execution history uniqueness" + predicate = " AND ".join(f"({check})" for check in checks) + if conn.execute( + f"SELECT 1 FROM {table} WHERE NOT ({predicate}) LIMIT 1" + ).fetchone(): + return f"{table} contains invalid execution history records" + metadata = conn.execute( + "SELECT singleton_id,schema_version,storage_epoch " + "FROM execution_history_metadata" + ).fetchall() + if len(metadata) != 1 or metadata[0][0] != 1 or metadata[0][1] != SCHEMA_VERSION: + return "Unsupported execution history metadata/version" + indexes = {row[1]: row for row in conn.execute("PRAGMA index_list(execution_runs)")} + for name, (unique, columns) in _INDEXES.items(): + row = indexes.get(name) + if row is None or bool(row[2]) != unique or row[4]: + return f"Missing or invalid execution history index {name}" + actual = tuple(conn.execute(f"PRAGMA index_xinfo({name})")) + keys = tuple(item for item in actual if item[5]) + if tuple(item[2] for item in keys) != columns or any( + item[3] or item[4] != "BINARY" for item in keys + ): + return f"Invalid execution history index {name}" + return None + + +def migrate_execution_history(conn: sqlite3.Connection) -> None: + """Install both tables and indexes atomically, without backfilling runs.""" + conn.execute("SAVEPOINT execution_history_schema") + try: + present = conn.execute( + "SELECT 1 FROM sqlite_master WHERE name IN " + "('execution_runs','execution_history_metadata') LIMIT 1" + ).fetchone() + if not present: + for table, columns in _TABLE_COLUMNS.items(): + definitions = [ + f"{name} {kind}{' NOT NULL' if required else ''}" + f"{' PRIMARY KEY' if primary else ''}" + for name, kind, required, primary in columns + ] + definitions.extend(f"CHECK({check})" for check in _TABLE_CHECKS[table]) + conn.execute(f"CREATE TABLE {table} ({', '.join(definitions)})") + for name, (unique, columns) in _INDEXES.items(): + conn.execute( + f"CREATE {'UNIQUE ' if unique else ''}INDEX {name} " + f"ON execution_runs({', '.join(columns)})" + ) + conn.execute( + "INSERT INTO execution_history_metadata VALUES(1,?,?)", + (SCHEMA_VERSION, str(uuid4())), + ) + error = history_schema_error(conn, allow_missing=False) + if error: + raise sqlite3.DatabaseError(error) + conn.execute("RELEASE execution_history_schema") + except BaseException: + conn.execute("ROLLBACK TO execution_history_schema") + conn.execute("RELEASE execution_history_schema") + raise + + +def rotate_history_epoch(conn: sqlite3.Connection) -> str: + """Publish a fresh staging epoch without rewriting restored run epochs.""" + epoch = str(uuid4()) + with conn: + cursor = conn.execute( + "UPDATE execution_history_metadata SET storage_epoch=? " + "WHERE singleton_id=1", + (epoch,), + ) + if cursor.rowcount != 1: + raise sqlite3.DatabaseError("Missing execution history metadata") + return epoch diff --git a/database/models.py b/database/models.py index 37a6dcd..ff0bb9a 100644 --- a/database/models.py +++ b/database/models.py @@ -16,6 +16,12 @@ from contextlib import closing from pathlib import Path from typing import Optional, Set, Tuple +from .execution_history_schema import ( + history_schema_error, + migrate_execution_history, + rotate_history_epoch, +) + # Import logging system try: from ..utils.logging_config import get_logger @@ -274,6 +280,7 @@ class PromptModel: # Image uniqueness by full path instead of basename (3.2.4) self._migrate_image_uniqueness_by_path(conn) self._migrate_prompt_pairs(conn) + migrate_execution_history(conn) def _migrate_prompt_pairs(self, conn: sqlite3.Connection) -> None: conn.execute("SAVEPOINT prompt_pair_schema") @@ -1037,6 +1044,9 @@ class PromptModel: pair_error = PromptModel._pair_schema_error(conn) if pair_error: return False, pair_error + history_error = history_schema_error(conn) + if history_error: + return False, history_error return True, "ok" except sqlite3.Error as e: return False, f"Invalid SQLite database: {e}" @@ -1112,6 +1122,7 @@ class PromptModel: staged = None try: staged = PromptModel(staged_path) + rotate_history_epoch(staged.get_connection()) ok, reason = self.verify_database_file(staged_path) if not ok: raise ValueError(reason) @@ -1127,6 +1138,9 @@ class PromptModel: with PromptModel._init_lock: PromptModel._initialized_paths.add(self._schema_key()) self.logger.info(f"Database restored from {src_path}") + from .execution_history import notify_history_restore + + notify_history_restore(self.db_path) return backup_path finally: if staged is not None: diff --git a/database/operations.py b/database/operations.py index 588c07e..f927196 100644 --- a/database/operations.py +++ b/database/operations.py @@ -11,6 +11,7 @@ from typing import Optional, List, Dict, Any from .models import PromptModel, normalize_image_path, utc_now_iso from .pairs import PairDependencyError, PairOperations from .library import LibraryOperations +from .execution_history import ExecutionHistoryOperations # Import logging system try: @@ -86,7 +87,10 @@ def _resolve_db_path(db_path: Optional[str] = None) -> str: if db_path is None: try: - from py.config import PromptManagerConfig + try: + from ..py.config import PromptManagerConfig + except ImportError: + from py.config import PromptManagerConfig db_path = PromptManagerConfig.DEFAULT_DB_PATH except Exception: @@ -148,7 +152,7 @@ def _coerce_file_info(metadata: Any) -> tuple: ) -class PromptDatabase(PairOperations, LibraryOperations): +class PromptDatabase(PairOperations, LibraryOperations, ExecutionHistoryOperations): """Database operations class for managing prompts.""" def __init__(self, db_path: Optional[str] = None): diff --git a/docs/application-redesign.md b/docs/application-redesign.md index f1fa0f0..3107e9a 100644 --- a/docs/application-redesign.md +++ b/docs/application-redesign.md @@ -9,8 +9,8 @@ code. This plan proposes implementation order and acceptance gates. - Foundation merged in [#198](https://github.com/ComfyAssets/ComfyUI_PromptManager/pull/198). - [#200](https://github.com/ComfyAssets/ComfyUI_PromptManager/issues/200) delivers native Images browsing, physical-file details, infinite scroll, Grid/List/Masonry, Viewer.js and video playback. The workspace opens on Images; original application entry points remain available. - Shared role-aware sending and the physical-file prompt inspector merged in [#203](https://github.com/ComfyAssets/ComfyUI_PromptManager/pull/203) and [#206](https://github.com/ComfyAssets/ComfyUI_PromptManager/pull/206). -- [#208](https://github.com/ComfyAssets/ComfyUI_PromptManager/issues/208) adds the native Prompt Library: one positive parent, saved negative variants, full selected text and a shared filtered image gallery. It does not invent chronological runs from counters or scan dates. [Durable history storage](prompt-history-storage.md) must be designed and tested before the execution timeline. -- Execution history, tagging/review and System consolidation remain subsequent slices under #196. +- [#208](https://github.com/ComfyAssets/ComfyUI_PromptManager/issues/208) adds the native Prompt Library: one positive parent, saved negative variants, full selected text and a shared filtered image gallery. It does not invent chronological runs from counters or scan dates. [Durable history storage](prompt-history-storage.md) now records new job outcomes; prompt contexts, output associations and the execution timeline remain pending under #209. +- Prompt execution contexts and timeline, tagging/review and System consolidation remain subsequent slices under #196; #209 now records durable job-level outcomes. ## Scope and compatibility diff --git a/docs/prompt-history-storage.md b/docs/prompt-history-storage.md index 79b86ac..2e5440c 100644 --- a/docs/prompt-history-storage.md +++ b/docs/prompt-history-storage.md @@ -1,46 +1,62 @@ # Durable Prompt History -Status: future proposal tracked in [#209](https://github.com/ComfyAssets/ComfyUI_PromptManager/issues/209); the native library does not migrate execution storage. +Status: the durable job ledger is implemented as the first slice of +[#209](https://github.com/ComfyAssets/ComfyUI_PromptManager/issues/209). Prompt +contexts, output associations and the chronological timeline remain pending; +#209 stays open. -## What the current database proves +## Recorded jobs -- `prompts.run_count` and `prompt_pairs.run_count` are aggregate usage counters. - `utils/usage_tracking.py` counts known recipes when ComfyUI dequeues a job, - including cached node executions. Dequeue does not establish successful completion. -- `last_used_at` records the latest aggregate update. The legacy prompt migration - estimated usage from linked images, with a minimum of one; no flag distinguishes - those estimates from newer recorded uses (`database/models.py`). -- `generated_images` identifies linked files; distinct image IDs are image counts, - never run counts. Its default `generation_time` is database insertion time. -- `image_prompt_pairs` stores immutable association provenance. More than one - provenance row can describe the same image and recipe. Its `created_at` is the - association time, potentially an old-image scan. -- `prompt_pairs.created_at` is the saved recipe date. Execution PNG evidence in - `utils/pair_execution.py` records distinct text/composition values without a - durable job ID or timestamp. `PromptTracker` execution IDs remain process-local. +New jobs observed at the ComfyUI queue's dequeue boundary receive an immutable +run identity and a host-session/queue-item identity. The public ComfyUI job UUID +is a correlation field: submitting it twice records two attempts. Rejected graphs +and jobs removed from the queue before dequeue create no execution rows. Cached +jobs are recorded even when no PromptManager node runs, and headless jobs do not +require a browser or WebSocket client. -The current UI can show saved negative variants, linked images and aggregate uses -with an estimate caveat. It cannot honestly show chronological execution dots, -repeat-run branches or completed-run counts. +The additive `execution_runs` table stores dequeue time, observed callback time, +optional executor event times and an explicit outcome: `observed`, `succeeded`, +`failed`, `interrupted` or `unknown`. Matching stored terminal events determine +outcomes. Missing or conflicting terminal evidence remains unknown; a process +crash can leave an observed row without an invented completion time. Bounded +safe diagnostics retain node ID/type and exception type, never graph text, +arbitrary extra data, credentials, exception messages, traceback, inputs or +output payloads. A bounded background writer keeps storage failures from +changing host execution and reports capture gaps rather than inventing rows. -## Smallest useful future migration +Existing databases begin with zero execution rows. Migration and backup restore +preserve prompts, pairs, image links and aggregate counters; storage epochs keep +stale queued writes from being replayed into restored data. There is no +historical backfill from counters, scan dates, image timestamps or saved recipes. -Add a durable `run_events` table using a stable host-session and ComfyUI job ID -plus recipe identity. Create started events at dequeue, rather than node encoding, -so cached executions remain visible. Retain nullable recipe references and immutable -raw/effective positive, negative and composition snapshots. Record completion, -failure and interruption only when observed; otherwise retain unknown status. +[Testing](../TESTING.md#durable-execution-ledger) describes the real CPU host +matrix, historical failure proof, privacy checks and cleanup. Python deployment +requires restarting ComfyUI. -Add a deduplicated `run_images(run_id, image_id)` relationship. Multiple output -images can belong to one run, and one physical file must not manufacture multiple -runs through repeated metadata evidence. Preserve aggregate counters and existing -image provenance. Do not manufacture historical events from counters, file dates -or scans. +## What the existing UI proves -## Required migration and integration tests +`prompts.run_count` and `prompt_pairs.run_count` remain aggregate usage counters, +not completed-run counts. `last_used_at` is the latest aggregate update; legacy +migration may have estimated counts from linked images. `generated_images` +identifies linked files; image counts and insertion times are not run counts or +execution times. Image/pair association creation time can reflect an old-image +scan. A saved pair's creation time is its recipe date. -Verify old databases upgrade transactionally and idempotently, with rollback on -failure; unchanged counters and associations; cached repeat executions; failed or -cancelled jobs; batched outputs; duplicate callbacks; process restart; unknown -connected inputs; and metadata rescans that create no execution events. Cover real -Classic and Nodes 2.0 hosts before exposing run chronology. +The native Prompt Library shows saved negative variants, linked images and +aggregate uses. It does not expose the ledger as chronological dots or repeat-run +branches. `PromptTracker` identities and PNG recipe evidence remain separate +from durable host job identity. + +## Remaining history work + +Design bounded execution contexts with nullable recipe references, immutable +raw/effective positive, negative and composition snapshots and explicit +inferred/recorded/unknown provenance. Cached nodes must not inherit unrelated +recent prompt evidence. Define deletion and restore behavior before connecting +saved recipes to historical snapshots. + +Add exact output references and a deduplicated run/image relationship without +claiming cached outputs were newly created. Preserve existing image provenance +and same-path replacement protection. Only after those semantics are tested can +the timeline honestly show prompt-level execution chronology, batched outputs +and repeat-run branches. diff --git a/tests/comfyui/execution-fixture.mjs b/tests/comfyui/execution-fixture.mjs new file mode 100644 index 0000000..1966f65 --- /dev/null +++ b/tests/comfyui/execution-fixture.mjs @@ -0,0 +1,47 @@ +import { execFile } from 'node:child_process'; +import { promisify } from 'node:util'; +import { fileURLToPath } from 'node:url'; +import { access } from 'node:fs/promises'; +import { join } from 'node:path'; +import assert from 'node:assert/strict'; +import { test as base, expect } from '../../tools/vendor/node_modules/@playwright/test/index.mjs'; +import { launchExecutionHost } from './execution-host.mjs'; +import runtime from '../../tools/vendor/comfyui-runtime.json' with { type: 'json' }; + +export { launchExecutionHost }; +const run = promisify(execFile); +const reader = fileURLToPath(new URL('./read_execution_history.py', import.meta.url)); + +export async function readExecutionRuns(databasePath, jobId, sessionId) { + const args = [reader, databasePath, jobId]; + if (sessionId !== undefined) args.push('--session-id', sessionId); + const { stdout } = await run(process.env.PM_TEST_PYTHON || 'python', args, { timeout: 5000, maxBuffer: 1024 * 1024 }); + return JSON.parse(stdout); +} + +export const test = base.extend({ + execution: [async ({ playwright }, use, workerInfo) => { + const execution = await launchExecutionHost({ + logPath: fileURLToPath(new URL(`../../tools/vendor/test-results/executions/${workerInfo.project.name}-host.log`, import.meta.url)), + }); + let request; + try { + request = await playwright.request.newContext({ baseURL: 'http://127.0.0.1:8196/' }); + await expect.poll(async () => { + if (execution.host.child.exitCode !== null) throw new Error('Owned execution host exited before readiness'); + try { return (await request.get('/system_stats')).ok(); } catch { return false; } + }, { timeout: 90000 }).toBe(true); + const info = (await (await request.get('/system_stats')).json()).system; + expect(info.comfyui_version).toBe(runtime.serverVersion); + expect(info.pytorch_version).toBe(`${runtime.torchVersion}+cpu`); + expect(info.required_frontend_version).toBe(runtime.frontendVersion); + expect((await request.get('/prompt_manager/stats')).ok()).toBe(true); + await access(execution.databasePath); + await assert.rejects(access(join(execution.host.stateDirectory, 'custom_nodes/promptmanager/prompts.db')), { code: 'ENOENT' }); + await use({ ...execution, request }); + } finally { + try { if (request) await request.post('/pm_execution_fixture/gate/release', { timeout: 5000 }); } + finally { try { await request?.dispose(); } finally { await execution.dispose(); } } + } + }, { scope: 'worker', timeout: 120000 }], +}); diff --git a/tests/comfyui/execution-history.spec.mjs b/tests/comfyui/execution-history.spec.mjs new file mode 100644 index 0000000..f39fe00 --- /dev/null +++ b/tests/comfyui/execution-history.spec.mjs @@ -0,0 +1,168 @@ +import { randomUUID } from 'node:crypto'; +import { expect } from '../../tools/vendor/node_modules/@playwright/test/index.mjs'; +import { test, readExecutionRuns } from './execution-fixture.mjs'; + +const sensitive = ['PM_SECRET_PROMPT_209', 'PM_SECRET_TOKEN_209', 'PM_SECRET_TRACE_209']; +const cached = value => ({ '1': { class_type: 'PMHistoryCached', inputs: { value } } }); +const gate = value => ({ '1': { class_type: 'PMHistoryGate', inputs: { value } } }); + +async function submit(request, prompt, promptId = randomUUID()) { + const response = await request.post('/prompt', { data: { prompt, prompt_id: promptId, + extra_data: { token: sensitive[1], extra_pnginfo: { private: sensitive[2] } } } }); + expect(response.ok()).toBe(true); + expect((await response.json()).prompt_id).toBe(promptId); + return promptId; +} + +async function history(request, jobId, terminal = 'execution_success') { + let item; + await expect.poll(async () => { + const response = await request.get(`/history/${jobId}`); + expect(response.ok()).toBe(true); + item = (await response.json())[jobId]; + return item?.status?.messages?.some(([kind]) => kind === terminal) || false; + }, { timeout: 15000 }).toBe(true); + expect(item.prompt[1]).toBe(jobId); + return item; +} + +async function rows(execution, jobId, outcomes) { + let result; + await expect.poll(async () => { + result = await readExecutionRuns(execution.databasePath, jobId); + return { schemaAvailable: result.schemaAvailable, outcomes: result.rows.map(row => row.outcome) }; + }, { timeout: 10000, message: `Persisted exact-job outcomes for ${jobId}` }).toEqual({ schemaAvailable: true, outcomes }); + for (const sentinel of sensitive) expect(JSON.stringify(result.rows)).not.toContain(sentinel); + for (const row of result.rows) { + expect(row.observed_dequeue_at).toBeGreaterThan(0); + expect(row.observed_callback_at).toBeGreaterThanOrEqual(row.observed_dequeue_at); + expect(row.executor_started_at).toBeGreaterThan(0); + expect(row.terminal_event_at).toBeGreaterThanOrEqual(row.executor_started_at); + } + return result.rows; +} + +async function control(request, action) { + expect((await request.post(`/pm_execution_fixture/gate/${action}`)).ok()).toBe(true); +} + +async function entered(request) { + await expect.poll(async () => (await (await request.get('/pm_execution_fixture/gate')).json()).entered, + { timeout: 10000 }).toBe(true); +} + +test('headless actual queue execution persists without a browser or WebSocket client', async ({ execution }) => { + const id = await submit(execution.request, cached(`headless ${sensitive[0]}`)); + const item = await history(execution.request, id); + expect(item.status.status_str).toBe('success'); + expect((await rows(execution, id, ['succeeded']))[0].outcome).toBe('succeeded'); +}); + +test('cached reruns after failure survive renderer switches both ways and workflow reload', async ({ execution, page, context }, testInfo) => { + const request = execution.request; + await context.route('**/*', async route => { + const path = new URL(route.request().url()).pathname.replace(/^\/api\//, '/'); + if (path === '/settings' && route.request().method() === 'GET') { + const response = await route.fetch(); + return route.fulfill({ json: { ...await response.json(), + 'Comfy.VueNodes.Enabled': testInfo.project.metadata.initialRenderer, + 'Comfy.TutorialCompleted': true, 'Comfy.Workflow.Persist': false, + 'Comfy.Notification.ShowVersionUpdates': false } }); + } + return route.continue(); + }); + const errors = []; + page.on('pageerror', error => errors.push(error.message)); + await page.goto('/', { waitUntil: 'domcontentloaded' }); + await expect(page.getByRole('status', { name: 'Loading ComfyUI' })).toBeHidden({ timeout: 30000 }); + await page.waitForFunction(async () => { + const { app } = await import('/scripts/app.js'); + return !!app.graph && !!globalThis.LiteGraph?.registered_node_types?.PMHistoryCached; + }); + await expect.poll(() => page.evaluate(() => LiteGraph.vueNodesMode)).toBe(testInfo.project.metadata.initialRenderer); + await page.evaluate(async value => { + const { app } = await import('/scripts/app.js'); + app.graph.clear(); + const node = LiteGraph.createNode('PMHistoryCached'); + app.graph.add(node); + node.widgets.find(widget => widget.name === 'value').value = value; + }, `cache ${sensitive[0]}`); + const graph = await page.evaluate(async () => (await (await import('/scripts/app.js')).app.graphToPrompt()).output); + expect(Object.values(graph).map(node => node.class_type)).toEqual(['PMHistoryCached']); + const cachedNode = Object.keys(graph)[0]; + await history(request, await submit(request, graph)); + const failureGraph = { ...graph, '999': { class_type: 'PMHistoryFailure', inputs: { value: sensitive[2] } } }; + const failed = await submit(request, failureGraph); + const failureHistory = await history(request, failed, 'execution_error'); + expect(failureHistory.status.messages.find(([kind]) => kind === 'execution_error')[1].exception_message).toContain(sensitive[2]); + const failedRows = await rows(execution, failed, ['failed']); + expect(failedRows[0]).toMatchObject({ outcome: 'failed', node_id: '999', node_type: 'PMHistoryFailure', exception_type: 'ValueError' }); + for (const renderer of [!testInfo.project.metadata.initialRenderer, testInfo.project.metadata.initialRenderer]) { + await page.evaluate(async renderer => { + const { app } = await import('/scripts/app.js'); + await app.ui.settings.setSettingValue('Comfy.VueNodes.Enabled', renderer); + await app.loadGraphData(app.graph.serialize()); + }, renderer); + await expect.poll(() => page.evaluate(() => LiteGraph.vueNodesMode)).toBe(renderer); + const reloaded = await page.evaluate(async () => (await (await import('/scripts/app.js')).app.graphToPrompt()).output); + const id = await submit(request, reloaded); + const item = await history(request, id); + expect(item.status.messages.find(([kind]) => kind === 'execution_cached')[1].nodes).toContain(cachedNode); + expect((await rows(execution, id, ['succeeded']))[0].outcome).toBe('succeeded'); + } + expect(errors).toEqual([]); +}); + +test('reusing a public UUID creates two distinct persisted attempts', async ({ execution }) => { + const id = randomUUID(); + await history(execution.request, await submit(execution.request, cached('first reused UUID'), id)); + await rows(execution, id, ['succeeded']); + const graph = cached('second reused UUID'); + await submit(execution.request, graph, id); + await expect.poll(async () => { + const response = await execution.request.get(`/history/${id}`); + return (await response.json())[id]?.prompt[2]?.['1']?.inputs?.value; + }, { timeout: 15000 }).toBe('second reused UUID'); + const attempts = await rows(execution, id, ['succeeded', 'succeeded']); + expect(attempts).toHaveLength(2); + expect(new Set(attempts.map(row => row.run_id)).size).toBe(2); + expect(new Set(attempts.map(row => row.queue_item_id)).size).toBe(2); + expect(new Set(attempts.map(row => row.session_id)).size).toBe(1); +}); + +test('real worker interruption persists an interrupted outcome', async ({ execution }) => { + const request = execution.request; + await control(request, 'hold'); + try { + const id = await submit(request, gate(`interrupt ${sensitive[0]}`)); + await entered(request); + expect((await request.post('/interrupt', { data: {} })).ok()).toBe(true); + await history(request, id, 'execution_interrupted'); + expect((await rows(execution, id, ['interrupted']))[0].outcome).toBe('interrupted'); + } finally { await control(request, 'release'); } +}); + +test('rejected and removed-before-dequeue jobs create no persisted attempts', async ({ execution }) => { + const request = execution.request; + const rejected = randomUUID(); + const response = await request.post('/prompt', { data: { prompt_id: rejected, + prompt: { '1': { class_type: 'PMHistoryMissing', inputs: {} } } } }); + expect(response.status()).toBe(400); + await control(request, 'hold'); + try { + const blocking = await submit(request, gate('hold worker until queued item is deleted')); + await entered(request); + const removed = await submit(request, cached('must never dequeue')); + const queued = (await (await request.get('/queue')).json()).queue_pending; + expect(queued.some(item => item[1] === removed)).toBe(true); + expect((await request.post('/queue', { data: { delete: [removed] } })).ok()).toBe(true); + expect((await (await request.get('/queue')).json()).queue_pending.some(item => item[1] === removed)).toBe(false); + await control(request, 'release'); + await history(request, blocking); + await rows(execution, blocking, ['succeeded']); + for (const id of [rejected, removed]) { + expect((await (await request.get(`/history/${id}`)).json())[id]).toBeUndefined(); + expect((await rows(execution, id, [])).length).toBe(0); + } + } finally { await control(request, 'release'); } +}); diff --git a/tests/comfyui/execution-host.mjs b/tests/comfyui/execution-host.mjs new file mode 100644 index 0000000..0e86aa8 --- /dev/null +++ b/tests/comfyui/execution-host.mjs @@ -0,0 +1,41 @@ +import { execFileSync } from 'node:child_process'; +import { copyFile, mkdtemp, rm } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import { createServer } from 'node:net'; +import { join } from 'node:path'; +import { fileURLToPath } from 'node:url'; +import { launchSharedHost as launchComfyHost, preparePromptManagerPlugin } from '../../tools/vendor/comfyui-host.mjs'; +import runtime from '../../tools/vendor/comfyui-runtime.json' with { type: 'json' }; + +const repositoryRoot = fileURLToPath(new URL('../../', import.meta.url)); + +export async function launchExecutionHost({ checkout = process.env.PM_COMFY_ROOT, + repository = process.env.PM_EXECUTION_REPOSITORY || repositoryRoot, + python = process.env.PM_TEST_PYTHON || 'python', port = 8196, pins = runtime, + logPath = join(repositoryRoot, 'tools/vendor/test-results/executions/host.log'), forwardOutput = true } = {}) { + if (port !== 8196) throw new Error('Execution tests reserve managed loopback port 8196'); + await new Promise((resolve, reject) => { + const probe = createServer(); + probe.once('error', error => reject(new Error('Execution loopback port 8196 is already in use', { cause: error }))); + probe.listen({ host: '127.0.0.1', port, exclusive: true }, () => probe.close(error => error ? reject(error) : resolve())); + }); + const fixtureDirectory = await mkdtemp(join(tmpdir(), 'pm-execution-fixture-')); + let host; + try { + await copyFile(new URL('./fixtures/ledger_nodes.py', import.meta.url), join(fixtureDirectory, '__init__.py')); + execFileSync('git', ['init', '--quiet', fixtureDirectory]); + execFileSync('git', ['-C', fixtureDirectory, 'add', '__init__.py']); + host = await launchComfyHost({ checkout, repository, python, port, pins, logPath, forwardOutput, + pluginName: 'promptmanager', preparePlugin: preparePromptManagerPlugin, + additionalPlugins: [{ pluginName: 'pm_execution_fixture', repository: fixtureDirectory }], + }); + } catch (error) { + await rm(fixtureDirectory, { recursive: true, force: true }); + throw error; + } + host.finished.finally(() => rm(fixtureDirectory, { recursive: true, force: true })).catch(() => {}); + return { host, fixtureDirectory, databasePath: join(host.stateDirectory, 'prompts.db'), async dispose() { + try { await host.stop(); } + finally { await rm(fixtureDirectory, { recursive: true, force: true }); } + } }; +} diff --git a/tests/comfyui/fixtures/ledger_nodes.py b/tests/comfyui/fixtures/ledger_nodes.py new file mode 100644 index 0000000..785fa21 --- /dev/null +++ b/tests/comfyui/fixtures/ledger_nodes.py @@ -0,0 +1,71 @@ +"""Deterministic nodes and queue controls installed only in an owned test host.""" + +import threading +import time + +from aiohttp import web +import comfy.model_management +from server import PromptServer + +_GATE = threading.Event() +_GATE.set() +_ENTERED = threading.Event() + + +@PromptServer.instance.routes.post("/pm_execution_fixture/gate/{action}") +async def control_gate(request: web.Request) -> web.Response: + action = request.match_info["action"] + if action == "hold": + _ENTERED.clear() + _GATE.clear() + elif action == "release": + _GATE.set() + else: + raise web.HTTPBadRequest() + return web.json_response({"held": not _GATE.is_set()}) + + +@PromptServer.instance.routes.get("/pm_execution_fixture/gate") +async def gate_status(request: web.Request) -> web.Response: + return web.json_response({"entered": _ENTERED.is_set(), "held": not _GATE.is_set()}) + + +class PMHistoryCached: + @classmethod + def INPUT_TYPES(cls) -> dict: + return {"required": {"value": ("STRING", {"default": "deterministic"})}} + + RETURN_TYPES = ("STRING",) + FUNCTION = "execute" + CATEGORY = "testing/execution-history" + OUTPUT_NODE = True + + def execute(self, value: str) -> dict: + return {"ui": {"text": [value]}, "result": (value,)} + + +class PMHistoryFailure(PMHistoryCached): + def execute(self, value: str) -> dict: + raise ValueError(value) + + +class PMHistoryGate(PMHistoryCached): + @classmethod + def IS_CHANGED(cls, **kwargs: object) -> float: + return float("nan") + + def execute(self, value: str) -> dict: + _ENTERED.set() + deadline = time.monotonic() + 30 + while not _GATE.wait(0.05): + comfy.model_management.throw_exception_if_processing_interrupted() + if time.monotonic() >= deadline: + raise TimeoutError("Owned execution fixture gate timed out") + return super().execute(value) + + +NODE_CLASS_MAPPINGS = { + "PMHistoryCached": PMHistoryCached, + "PMHistoryFailure": PMHistoryFailure, + "PMHistoryGate": PMHistoryGate, +} diff --git a/tests/comfyui/read_execution_history.py b/tests/comfyui/read_execution_history.py new file mode 100644 index 0000000..f2ac22a --- /dev/null +++ b/tests/comfyui/read_execution_history.py @@ -0,0 +1,46 @@ +"""Bounded read-only SQLite inspection for the owned execution fixture.""" + +import argparse +import json +from pathlib import Path +import sqlite3 +from uuid import UUID + + +def read_runs(database: str, job_id: str, session_id: str | None = None) -> dict: + job_id = str(UUID(job_id)) + if session_id is not None: + session_id = str(UUID(session_id)) + uri = Path(database).resolve().as_uri() + "?mode=ro" + with sqlite3.connect(uri, uri=True, timeout=0.25) as connection: + connection.row_factory = sqlite3.Row + connection.execute("PRAGMA query_only=ON") + table = connection.execute( + "SELECT 1 FROM sqlite_master WHERE type='table' AND name='execution_runs'" + ).fetchone() + if table is None: + return {"schemaAvailable": False, "rows": []} + sql = "SELECT * FROM execution_runs WHERE job_id=?" + parameters = [job_id] + if session_id is not None: + sql += " AND session_id=?" + parameters.append(session_id) + rows = connection.execute( + sql + " ORDER BY rowid LIMIT 257", parameters + ).fetchall() + if len(rows) > 256: + raise ValueError("Execution query exceeds the 256-row bound") + return {"schemaAvailable": True, "rows": [dict(row) for row in rows]} + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("database") + parser.add_argument("job_id") + parser.add_argument("--session-id") + arguments = parser.parse_args() + print( + json.dumps( + read_runs(arguments.database, arguments.job_id, arguments.session_id) + ) + ) diff --git a/tests/test_api_backup_restore.py b/tests/test_api_backup_restore.py index eb55b47..d91b302 100644 --- a/tests/test_api_backup_restore.py +++ b/tests/test_api_backup_restore.py @@ -11,7 +11,11 @@ import sqlite3 import sys import tempfile import unittest -from unittest.mock import patch +from unittest.mock import MagicMock, patch + +_mock_server = MagicMock() +_mock_server.PromptServer.instance.routes = MagicMock() +sys.modules.setdefault("server", _mock_server) sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from database.models import PromptModel # noqa: E402 diff --git a/tests/test_db_path_config.py b/tests/test_db_path_config.py index 36f49e2..500cc5d 100644 --- a/tests/test_db_path_config.py +++ b/tests/test_db_path_config.py @@ -14,6 +14,7 @@ from unittest.mock import patch sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +import database.operations as operations from database.operations import _resolve_db_path EXTENSION_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) @@ -22,6 +23,37 @@ EXTENSION_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) class TestResolveDbPath(unittest.TestCase): """Test _resolve_db_path() database path resolution.""" + def test_package_qualified_config_precedes_unrelated_standalone_config( + self, + ) -> None: + """A custom-node package must use its own configured disposable database.""" + with tempfile.TemporaryDirectory() as directory: + configured = os.path.join(directory, "owned", "prompts.db") + package_config = types.ModuleType("owned_promptmanager.py.config") + package_config.PromptManagerConfig = types.SimpleNamespace( + DEFAULT_DB_PATH=configured + ) + standalone_config = types.ModuleType("py.config") + standalone_config.PromptManagerConfig = types.SimpleNamespace( + DEFAULT_DB_PATH=os.path.join(directory, "unrelated.db") + ) + with ( + patch.dict( + operations.__dict__, + {"__package__": "owned_promptmanager.database", "__spec__": None}, + ), + patch.dict( + sys.modules, + { + "owned_promptmanager.py.config": package_config, + "py.config": standalone_config, + }, + ), + ): + resolved = _resolve_db_path(None) + self.assertEqual(resolved, configured) + self.assertTrue(os.path.isdir(os.path.dirname(configured))) + def test_explicit_absolute_path_used_as_is(self): """An explicit absolute path should be returned unchanged.""" with tempfile.TemporaryDirectory() as tmpdir: diff --git a/tests/test_execution_history_database.py b/tests/test_execution_history_database.py new file mode 100644 index 0000000..4c9f37e --- /dev/null +++ b/tests/test_execution_history_database.py @@ -0,0 +1,452 @@ +"""Durable attempt identity, terminal reconciliation and additive migration.""" + +import sqlite3 +import tempfile +import unittest +from contextlib import closing +from dataclasses import replace +from pathlib import Path +from types import SimpleNamespace +from uuid import UUID + +from database.models import PromptModel +from database.operations import PromptDatabase +from utils.execution_history_events import make_finish, make_start + +SESSION = "11111111-1111-4111-8111-111111111111" +JOB = "22222222-2222-4222-8222-222222222222" + + +class HistoryDatabaseTests(unittest.TestCase): + def setUp(self): + self.temp = tempfile.TemporaryDirectory() + self.addCleanup(self.temp.cleanup) + self.addCleanup(PromptModel.reset_schema_cache) + self.databases = [] + self.addCleanup(self.close_databases) + self.path = str(Path(self.temp.name) / "live.db") + + def close_databases(self): + for database in self.databases: + database.close_all() + + def database(self, path=None): + database = PromptDatabase(path or self.path) + self.databases.append(database) + return database + + def start(self, db, item=0, time=100): + return make_start( + SESSION, ((0, JOB, {}, {}), item), db.get_execution_epoch(), time + ) + + def finish(self, start, kind="execution_success", time=200, **fields): + messages = [] if kind is None else [(kind, {"prompt_id": JOB, **fields})] + return make_finish( + start.session_id, + start.queue_item_id, + start.storage_epoch, + SimpleNamespace(messages=messages), + time, + ) + + def count(self, db): + return ( + db.model.get_connection() + .execute("SELECT COUNT(*) FROM execution_runs") + .fetchone()[0] + ) + + def test_fresh_database_has_epoch_and_zero_runs(self): + from database.execution_history_schema import history_schema_error + + db = self.database() + self.assertEqual(str(UUID(db.get_execution_epoch())), db.get_execution_epoch()) + self.assertEqual(self.count(db), 0) + self.assertIsNone( + history_schema_error(db.model.get_connection(), allow_missing=False) + ) + self.assertEqual(PromptModel.verify_database_file(self.path), (True, "ok")) + + def test_migration_preserves_all_existing_rows_and_counters(self): + from database.execution_history_schema import migrate_execution_history + + db = self.database() + conn = db.model.get_connection() + with conn: + conn.execute("INSERT INTO prompts(id,text,run_count) VALUES(7,'kept',19)") + conn.execute( + "INSERT INTO generated_images(id,prompt_id,image_path,filename," + "file_path) VALUES(4,7,'/tmp/a','a','/tmp/a')" + ) + conn.execute( + "INSERT INTO prompt_pairs(id,positive_prompt_id,run_count) " + "VALUES(2,7,31)" + ) + conn.execute( + "INSERT INTO image_prompt_pairs(image_id,pair_id,provenance) " + "VALUES(4,2,'{}')" + ) + tables = ("prompts", "generated_images", "prompt_pairs", "image_prompt_pairs") + before = { + table: [tuple(row) for row in conn.execute(f"SELECT * FROM {table}")] + for table in tables + } + for pre_pair in (False, True): + with self.subTest(pre_pair=pre_pair): + with conn: + conn.execute("DROP TABLE execution_runs") + conn.execute("DROP TABLE execution_history_metadata") + if pre_pair: + with conn: + for table in ( + "image_prompt_pairs", + "prompt_pair_tags", + "prompt_pairs", + ): + conn.execute(f"DROP TABLE {table}") + before = { + table: rows + for table, rows in before.items() + if table in ("prompts", "generated_images") + } + db.close_all() + PromptModel.reset_schema_cache() + db = self.database() + conn = db.model.get_connection() + self.assertEqual(self.count(db), 0) + for table, rows in before.items(): + self.assertEqual( + [tuple(row) for row in conn.execute(f"SELECT * FROM {table}")], + rows, + ) + epoch = db.get_execution_epoch() + command = self.start(db) + db.write_execution_start(command) + row = db.get_execution_run(command.run_id) + migrate_execution_history(conn) + self.assertEqual(db.get_execution_epoch(), epoch) + self.assertEqual(db.get_execution_run(command.run_id), row) + self.assertEqual( + conn.execute("PRAGMA foreign_key_check").fetchall(), [] + ) + self.assertEqual( + conn.execute("PRAGMA integrity_check").fetchone()[0], "ok" + ) + + def test_failed_index_creation_rolls_back_both_tables_then_retry_succeeds(self): + from database.execution_history_schema import migrate_execution_history + + class FailingConnection(sqlite3.Connection): + def execute(self, sql, parameters=()): + if "CREATE" in sql.upper() and "INDEX" in sql.upper(): + raise sqlite3.OperationalError("injected index failure") + return super().execute(sql, parameters) + + with closing(sqlite3.connect(self.path, factory=FailingConnection)) as conn: + with self.assertRaisesRegex(sqlite3.OperationalError, "injected"): + migrate_execution_history(conn) + self.assertEqual( + conn.execute( + "SELECT name FROM sqlite_master WHERE type='table'" + ).fetchall(), + [], + ) + with closing(sqlite3.connect(self.path)) as conn: + migrate_execution_history(conn) + self.assertEqual( + conn.execute("SELECT COUNT(*) FROM execution_runs").fetchone()[0], 0 + ) + + def test_replay_preserves_first_observation_and_public_job_is_nonunique(self): + db = self.database() + command = self.start(db) + self.assertTrue(db.write_execution_start(command)) + self.assertTrue( + db.write_execution_start(replace(command, observed_dequeue_at=999)) + ) + self.assertEqual( + db.get_execution_run(command.run_id)["observed_dequeue_at"], 100 + ) + self.assertTrue(db.write_execution_start(self.start(db, item=1))) + self.assertEqual(self.count(db), 2) + self.assertEqual(db.get_execution_run(command.run_id)["outcome"], "observed") + + def test_unrelated_unique_conflicts_are_not_acknowledged_as_recorded(self): + db = self.database() + first = self.start(db) + db.write_execution_start(first) + with db.model.get_connection() as conn: + conn.execute("CREATE UNIQUE INDEX unexpected_job ON execution_runs(job_id)") + with self.assertRaises(sqlite3.IntegrityError): + db.write_execution_start(self.start(db, item=1)) + self.assertEqual(self.count(db), 1) + self.assertEqual(db.get_execution_run(first.run_id)["observed_dequeue_at"], 100) + + def test_primary_key_collision_for_another_attempt_is_not_a_replay(self): + db = self.database() + first = self.start(db) + db.write_execution_start(first) + with self.assertRaises(sqlite3.IntegrityError): + db.write_execution_start( + replace(self.start(db, item=1), run_id=first.run_id) + ) + self.assertEqual(self.count(db), 1) + + def test_schema_rejects_extra_uniqueness_and_allows_nonunique_indexes(self): + from database.execution_history_schema import history_schema_error + + for table, constraint in ( + ("execution_runs", "job_id"), + ("execution_runs", "session_id, queue_item_id"), + ("execution_history_metadata", "storage_epoch"), + ): + for table_constraint in (False, True): + with self.subTest( + table=table, + table_constraint=table_constraint, + constraint=constraint, + ): + path = str( + Path(self.temp.name) / f"unique-{len(self.databases)}.db" + ) + db = self.database(path) + conn = db.model.get_connection() + with conn: + if table_constraint: + epoch = db.get_execution_epoch() + sql = conn.execute( + "SELECT sql FROM sqlite_master WHERE name=?", (table,) + ).fetchone()[0] + conn.execute(f"DROP TABLE {table}") + conn.execute(sql[:-1] + f", UNIQUE({constraint}))") + if table == "execution_runs": + conn.execute( + "CREATE UNIQUE INDEX idx_execution_runs_identity " + "ON execution_runs(session_id, queue_item_id)" + ) + conn.execute( + "CREATE INDEX idx_execution_runs_observed " + "ON execution_runs(observed_dequeue_at, run_id)" + ) + else: + conn.execute( + "INSERT INTO execution_history_metadata " + "VALUES(1,1,?)", + (epoch,), + ) + else: + conn.execute( + f"CREATE UNIQUE INDEX unexpected_unique " + f"ON {table}({constraint})" + ) + self.assertIsNotNone(history_schema_error(conn)) + db = self.database() + with db.model.get_connection() as conn: + conn.execute('CREATE INDEX "extra job lookup" ON execution_runs(job_id)') + conn.execute( + "CREATE INDEX extra_epoch_lookup " + "ON execution_history_metadata(storage_epoch)" + ) + self.assertIsNone(history_schema_error(conn)) + + def test_absent_finish_never_fabricates_a_start(self): + db = self.database() + self.assertFalse(db.write_execution_finish(self.finish(self.start(db)))) + self.assertEqual(self.count(db), 0) + self.assertIsNone(db.get_execution_run(self.start(db).run_id)) + + def test_identical_terminal_replay_preserves_all_first_timestamps(self): + db = self.database() + start = self.start(db) + db.write_execution_start(start) + finish = self.finish( + start, timestamp=150, node_id="4", exception_type="SafeType" + ) + self.assertTrue(db.write_execution_finish(finish)) + original = db.get_execution_run(start.run_id) + self.assertTrue( + db.write_execution_finish(replace(finish, observed_callback_at=999)) + ) + self.assertEqual(db.get_execution_run(start.run_id), original) + self.assertEqual(original["outcome"], "succeeded") + self.assertEqual(original["observed_callback_at"], 200) + self.assertEqual(original["terminal_event_at"], 150) + self.assertIsNone(original["executor_started_at"]) + + def test_conflicting_terminal_callbacks_are_irreversibly_unknown(self): + for initial, next_kind in ( + ("execution_error", "execution_success"), + ("execution_success", "execution_interrupted"), + ): + with self.subTest(initial=initial): + db = self.database(str(Path(self.temp.name) / f"{initial}.db")) + start = self.start(db) + db.write_execution_start(start) + db.write_execution_finish(self.finish(start, initial, timestamp=130)) + db.write_execution_finish( + self.finish(start, next_kind, time=300, timestamp=220) + ) + row = db.get_execution_run(start.run_id) + self.assertEqual( + (row["outcome"], row["reason"]), ("unknown", "conflicting_terminal") + ) + db.write_execution_finish(self.finish(start, initial)) + self.assertEqual(db.get_execution_run(start.run_id), row) + + def test_empty_malformed_and_limited_duplicates_do_not_downgrade_proven_result( + self, + ): + db = self.database() + start = self.start(db) + db.write_execution_start(start) + db.write_execution_finish(self.finish(start, "execution_error", timestamp=180)) + original = db.get_execution_run(start.run_id) + empty = self.finish(start, None) + variants = ( + empty, + replace(empty, evidence=replace(empty.evidence, invalid=True)), + replace(empty, evidence=replace(empty.evidence, limited=True)), + ) + for finish in variants: + db.write_execution_finish(finish) + self.assertEqual(db.get_execution_run(start.run_id), original) + + def test_conflicting_evidence_duplicate_downgrades_proven_result(self): + db = self.database() + start = self.start(db) + db.write_execution_start(start) + db.write_execution_finish(self.finish(start)) + finish = make_finish( + SESSION, + 0, + start.storage_epoch, + SimpleNamespace( + messages=[ + ("execution_success", {"prompt_id": JOB}), + ("execution_error", {"prompt_id": JOB}), + ] + ), + 300, + ) + db.write_execution_finish(finish) + self.assertEqual( + db.get_execution_run(start.run_id)["reason"], "conflicting_terminal" + ) + + def test_initial_unknown_is_never_promoted(self): + db = self.database() + start = self.start(db) + db.write_execution_start(start) + db.write_execution_finish(self.finish(start, None)) + row = db.get_execution_run(start.run_id) + self.assertEqual(row["reason"], "missing_terminal") + db.write_execution_finish(self.finish(start)) + self.assertEqual(db.get_execution_run(start.run_id), row) + + def test_writer_timeout_is_bounded_and_restored_for_existing_operations(self): + db = self.database() + start = self.start(db) + conn = db.model.get_connection() + timeout = conn.execute("PRAGMA busy_timeout").fetchone()[0] + with closing(sqlite3.connect(self.path)) as locker: + locker.execute("BEGIN IMMEDIATE") + import time + + before = time.monotonic() + with self.assertRaises(sqlite3.OperationalError): + db.write_execution_start(start) + self.assertLess(time.monotonic() - before, 1.5) + self.assertEqual(conn.execute("PRAGMA busy_timeout").fetchone()[0], timeout) + self.assertEqual(self.count(db), 0) + self.assertTrue(db.write_execution_start(start)) + + def test_sessions_keep_same_queue_item_and_public_job_separate(self): + db = self.database() + first = self.start(db) + second = make_start( + "33333333-3333-4333-8333-333333333333", + ((0, JOB, {}, {}), 0), + db.get_execution_epoch(), + 101, + ) + db.write_execution_start(first) + db.write_execution_start(second) + db.write_execution_finish(self.finish(second, "execution_interrupted")) + self.assertEqual(self.count(db), 2) + self.assertEqual(db.get_execution_run(first.run_id)["outcome"], "observed") + self.assertEqual(db.get_execution_run(second.run_id)["outcome"], "interrupted") + + def test_persisted_constraints_reject_invalid_identity_outcome_and_diagnostics( + self, + ): + db = self.database() + start = self.start(db) + db.write_execution_start(start) + conn = db.model.get_connection() + for assignment in ( + "queue_item_id=-1", + "job_id='not-a-uuid'", + "record_version=2", + "observed_dequeue_at=-1", + "outcome='succeeded'", + "reason='terminal_success'", + "diagnostics_omitted=2", + "node_id='" + "a" * 129 + "'", + "exception_type='" + "a" * 257 + "'", + ): + with self.subTest(assignment=assignment): + with self.assertRaises(sqlite3.IntegrityError), conn: + conn.execute(f"UPDATE execution_runs SET {assignment}") + self.assertEqual(db.get_execution_run(start.run_id)["outcome"], "observed") + with conn: + conn.execute("PRAGMA ignore_check_constraints=ON") + conn.execute("UPDATE execution_runs SET queue_item_id=-1") + self.assertFalse(PromptModel.verify_database_file(self.path)[0]) + + def test_bad_commands_cannot_bypass_transport_validation(self): + db = self.database() + start = self.start(db) + with self.assertRaises(ValueError): + db.write_execution_start(replace(start, observed_dequeue_at=-1)) + self.assertEqual(self.count(db), 0) + + def test_schema_verification_rejects_missing_tables_indexes_checks_and_versions( + self, + ): + from database.execution_history_schema import history_schema_error + + for mutation in ( + "DROP TABLE execution_runs", + "DROP TABLE execution_history_metadata", + "DELETE FROM execution_history_metadata", + "ALTER TABLE execution_runs ADD COLUMN extra_data TEXT", + "DROP INDEX idx_execution_runs_identity", + "DROP INDEX idx_execution_runs_observed", + "DROP INDEX idx_execution_runs_identity; CREATE UNIQUE INDEX " + "idx_execution_runs_identity ON execution_runs(job_id)", + "DROP INDEX idx_execution_runs_observed; CREATE INDEX " + "idx_execution_runs_observed ON " + "execution_runs(observed_dequeue_at DESC,run_id)", + "PRAGMA writable_schema=ON; UPDATE sqlite_master SET " + "sql=replace(sql, 'queue_item_id >= 0', 'queue_item_id >= -1') " + "WHERE name='execution_runs'; PRAGMA writable_schema=OFF", + "PRAGMA ignore_check_constraints=ON; " + "UPDATE execution_history_metadata SET schema_version=2", + "PRAGMA writable_schema=ON; UPDATE sqlite_master SET " + "sql=replace(sql, \"'observed'\", \"'OBSERVED'\") " + "WHERE name='execution_runs'; PRAGMA writable_schema=OFF", + ): + with self.subTest(mutation=mutation): + path = str(Path(self.temp.name) / f"schema-{len(self.databases)}.db") + db = self.database(path) + db.close_all() + with closing(sqlite3.connect(path)) as conn, conn: + conn.executescript(mutation) + with closing(sqlite3.connect(path)) as conn: + self.assertIsNotNone(history_schema_error(conn)) + self.assertFalse(PromptModel.verify_database_file(path)[0]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_execution_history_events.py b/tests/test_execution_history_events.py new file mode 100644 index 0000000..76a1474 --- /dev/null +++ b/tests/test_execution_history_events.py @@ -0,0 +1,379 @@ +"""Bounded, immutable evidence captured from real host-shaped queue callbacks.""" + +from collections import namedtuple +from dataclasses import FrozenInstanceError, dataclass, replace +import json +import unittest +from uuid import UUID, uuid5 + +from utils.execution_history_events import ( + LifecycleEvent, + RunFinish, + TerminalEvidence, + classify_evidence, + encode_command, + make_finish, + make_start, +) + +SESSION = "11111111-1111-4111-8111-111111111111" +JOB = "22222222-2222-4222-8222-222222222222" +OTHER_JOB = "33333333-3333-4333-8333-333333333333" +EPOCH = "44444444-4444-4444-8444-444444444444" +Status = namedtuple("ExecutionStatus", "status_str completed messages") + + +def queue_result(item_id=0, job_id=JOB): + return ((1.0, job_id, {"secret": "graph"}, {"token": "private"}, []), item_id) + + +def message(kind, job_id=JOB, **fields): + return (kind, {"prompt_id": job_id, **fields}) + + +def finish(messages, completed=True): + return make_finish(SESSION, 0, EPOCH, Status("success", completed, messages), 500) + + +class ExecutionHistoryEventsTests(unittest.TestCase): + def test_dequeue_identity_is_deterministic_per_session_and_item(self): + start = make_start(SESSION, queue_result(), EPOCH, 100) + repeated = make_start(SESSION, queue_result(), EPOCH, 100) + second = make_start(SESSION, queue_result(1), EPOCH, 100) + self.assertEqual(start.run_id, repeated.run_id) + self.assertNotEqual(start.run_id, second.run_id) + self.assertEqual(start.run_id, str(uuid5(UUID(SESSION), "0"))) + self.assertEqual(start.queue_item_id, 0) + self.assertEqual(start.job_id, JOB) + with self.assertRaises(FrozenInstanceError): + start.job_id = OTHER_JOB + + def test_required_identity_and_observation_inputs_are_rejected(self): + invalid = [None, "bad", JOB.upper(), "", True, 1] + # The UUID fixture contains digits only; use an upper-case UUID with letters. + invalid[2] = "AAAAAAAA-AAAA-4AAA-8AAA-AAAAAAAAAAAA" + for value in invalid: + with self.subTest(value=value): + self.assertIsNone(make_start(value, queue_result(), EPOCH, 100)) + self.assertIsNone( + make_start(SESSION, queue_result(job_id=value), EPOCH, 100) + ) + self.assertIsNone(make_start(SESSION, queue_result(), value, 100)) + self.assertIsNone(make_finish(value, 0, EPOCH, None, 100)) + for value in [-1, True, 1.5, "0", 2**63]: + with self.subTest(value=value): + self.assertIsNone(make_start(SESSION, queue_result(value), EPOCH, 100)) + self.assertIsNone(make_finish(SESSION, value, EPOCH, None, 100)) + self.assertIsNone(make_start(SESSION, queue_result(), EPOCH, value)) + self.assertIsNone(make_finish(SESSION, 0, EPOCH, None, value)) + for malformed in [None, (), ((), 0), ("bad", 0), (queue_result()[0],)]: + self.assertIsNone(make_start(SESSION, malformed, EPOCH, 0)) + + def test_terminal_messages_override_aggregate_flags_and_keep_safe_fields(self): + command = finish( + [ + message("execution_start", timestamp=100), + message( + "execution_error", + timestamp=200, + node_id="9", + node_type="Sampler", + exception_type="RuntimeError", + exception_message="SECRET", + traceback=["SECRET"], + inputs={"SECRET": 1}, + ), + ] + ) + result = classify_evidence(JOB, command.evidence) + self.assertEqual((result.outcome, result.reason), ("failed", "terminal_error")) + self.assertEqual( + (result.executor_started_at, result.terminal_event_at), (100, 200) + ) + self.assertEqual( + (result.node_id, result.node_type, result.exception_type), + ("9", "Sampler", "RuntimeError"), + ) + self.assertNotIn(b"SECRET", encode_command(command)) + self.assertNotIn(b"completed", encode_command(command)) + + def test_success_interruption_and_duplicate_terminals(self): + for kind, outcome, reason in [ + ("execution_success", "succeeded", "terminal_success"), + ("execution_interrupted", "interrupted", "terminal_interrupted"), + ]: + with self.subTest(kind=kind): + event = message(kind, timestamp=200) + result = classify_evidence(JOB, finish([event, event], False).evidence) + self.assertEqual((result.outcome, result.reason), (outcome, reason)) + result = classify_evidence( + JOB, + finish( + [message("execution_success"), message("execution_interrupted")] + ).evidence, + ) + self.assertEqual( + (result.outcome, result.reason), ("unknown", "conflicting_terminal") + ) + + def test_absent_unknown_and_foreign_messages_do_not_prove_completion(self): + for messages in [ + [], + [message("executed")], + [message("execution_success", OTHER_JOB)], + ]: + with self.subTest(messages=messages): + result = classify_evidence(JOB, finish(messages).evidence) + self.assertEqual( + (result.outcome, result.reason), ("unknown", "missing_terminal") + ) + for status in [ + None, + Status("success", True, None), + Status("success", True, {}), + ]: + command = make_finish(SESSION, 0, EPOCH, status, 100) + self.assertEqual( + classify_evidence(JOB, command.evidence).reason, "invalid_evidence" + ) + + def test_malformed_evidence_is_unknown_even_with_success(self): + for malformed in [ + None, + (), + ("execution_error",), + ("execution_error", None), + ("execution_error", {}), + message("execution_error", timestamp=-1), + message("execution_error", timestamp=True), + message("execution_error", timestamp=float("nan")), + ]: + with self.subTest(malformed=malformed): + result = classify_evidence( + JOB, finish([message("execution_success"), malformed]).evidence + ) + self.assertEqual( + (result.outcome, result.reason), ("unknown", "invalid_evidence") + ) + + def test_missing_timestamps_remain_null(self): + result = classify_evidence( + JOB, + finish( + [ + message("execution_start", timestamp=None), + message("execution_success"), + ] + ).evidence, + ) + self.assertEqual(result.outcome, "succeeded") + self.assertIsNone(result.executor_started_at) + self.assertIsNone(result.terminal_event_at) + + def test_optional_oversize_and_invalid_fields_are_omitted_without_losing_failure( + self, + ): + for value in ["x" * 129, "x" * 1_000_000, {}, 12]: + with self.subTest(value_type=type(value).__name__): + command = finish( + [ + message( + "execution_error", + node_id=value, + node_type=value, + exception_type="x" * 257, + ) + ] + ) + result = classify_evidence(JOB, command.evidence) + self.assertEqual(result.outcome, "failed") + self.assertTrue(result.diagnostics_omitted) + self.assertIsNone(result.node_id) + self.assertIsNone(result.node_type) + self.assertIsNone(result.exception_type) + + def test_unicode_diagnostics_fit_utf8_budget_and_preserve_failure(self): + command = finish( + [ + message( + "execution_error", + node_id="😀" * 128, + node_type="😀" * 128, + exception_type="😀" * 256, + ) + ] + ) + encoded = encode_command(command) + self.assertLessEqual(len(encoded), 2048) + self.assertIsInstance(json.loads(encoded), dict) + result = classify_evidence(JOB, command.evidence) + self.assertEqual(result.outcome, "failed") + self.assertTrue(result.diagnostics_omitted) + + def test_message_limit_never_trusts_partial_success(self): + command = finish([message("execution_success")] * 1025) + result = classify_evidence(JOB, command.evidence) + self.assertEqual((result.outcome, result.reason), ("unknown", "evidence_limit")) + self.assertLessEqual(len(encode_command(command)), 2048) + self.assertEqual( + classify_evidence( + JOB, finish([message("execution_success")] * 1024).evidence + ).outcome, + "succeeded", + ) + + def test_required_event_budget_exhaustion_is_not_truncated_success(self): + messages = [message("execution_success", timestamp=i) for i in range(100)] + command = finish(messages) + self.assertEqual( + classify_evidence(JOB, command.evidence).reason, "evidence_limit" + ) + self.assertLessEqual(len(encode_command(command)), 2048) + + def test_classification_uses_stored_job_id_after_capture(self): + command = finish( + [message("execution_error", OTHER_JOB), message("execution_success", JOB)] + ) + self.assertEqual(classify_evidence(JOB, command.evidence).outcome, "succeeded") + self.assertEqual( + classify_evidence(OTHER_JOB, command.evidence).outcome, "failed" + ) + + def test_encode_rejects_forged_invalid_or_unbounded_commands(self): + start = make_start(SESSION, queue_result(), EPOCH, 100) + command = finish([message("execution_success")]) + bad_event = LifecycleEvent("execution_success", JOB, -1, None, None, None) + for invalid in [ + None, + {}, + replace(start, run_id="bad"), + replace(start, job_id="x" * 129), + replace(start, queue_item_id=True), + replace(command, observed_callback_at=-1), + replace( + command, evidence=TerminalEvidence((bad_event,), False, False, False) + ), + replace(command, evidence=TerminalEvidence((), 1, False, False)), + replace( + command, + evidence=TerminalEvidence( + command.evidence.events * 100, False, False, False + ), + ), + ]: + with self.subTest(invalid=invalid): + with self.assertRaises(ValueError): + encode_command(invalid) + self.assertLessEqual(len(encode_command(start)), 2048) + self.assertIsInstance(command, RunFinish) + + def test_encoder_rejects_unknown_fields_and_nonstring_event_kinds(self): + @dataclass(frozen=True) + class ExtendedFinish(RunFinish): + secret: str = "credentials" + + command = finish([message("execution_success")]) + extended = ExtendedFinish(SESSION, 0, EPOCH, 100, command.evidence) + invalid_kind = LifecycleEvent([], JOB, None, None, None, None) + for invalid in [ + extended, + replace( + command, evidence=TerminalEvidence((invalid_kind,), False, False, False) + ), + ]: + with self.subTest(invalid=invalid): + with self.assertRaises(ValueError): + encode_command(invalid) + + def test_malformed_foreign_timestamp_cannot_discredit_matching_success(self): + command = finish( + [ + message("execution_error", OTHER_JOB, timestamp=-1), + message("execution_success", JOB, timestamp=300), + ] + ) + matching = classify_evidence(JOB, command.evidence) + self.assertEqual( + (matching.outcome, matching.reason), ("succeeded", "terminal_success") + ) + self.assertEqual( + classify_evidence(OTHER_JOB, command.evidence).reason, "invalid_evidence" + ) + self.assertLessEqual(len(encode_command(command)), 2048) + + def test_malformed_foreign_envelope_is_scoped_when_identity_is_usable(self): + command = finish( + [ + ("execution_error", {"prompt_id": OTHER_JOB}, "unexpected"), + message("execution_success", JOB), + ] + ) + self.assertEqual(classify_evidence(JOB, command.evidence).outcome, "succeeded") + self.assertEqual( + classify_evidence(OTHER_JOB, command.evidence).reason, "invalid_evidence" + ) + + def test_diagnostic_budget_reduction_retains_job_scoped_invalidity(self): + command = finish( + [ + message("execution_error", JOB, timestamp=-1), + message( + "execution_success", + JOB, + node_id="😀" * 128, + node_type="😀" * 128, + exception_type="😀" * 256, + ), + ] + ) + result = classify_evidence(JOB, command.evidence) + self.assertEqual( + (result.outcome, result.reason), ("unknown", "invalid_evidence") + ) + self.assertLessEqual(len(encode_command(command)), 2048) + + def test_job_scoped_invalidity_budget_overflow_is_evidence_limit(self): + command = finish( + [ + message("execution_error", str(UUID(int=i)), timestamp=-1) + for i in range(100) + ] + + [message("execution_success")] + ) + self.assertEqual( + classify_evidence(JOB, command.evidence).reason, "evidence_limit" + ) + self.assertLessEqual(len(encode_command(command)), 2048) + + def test_oversized_unknown_kind_is_ignored_without_hashing(self): + class OversizedKind(str): + def __hash__(self): + raise AssertionError("Oversized event kind must not be hashed") + + command = finish( + [ + message(OversizedKind("x" * 1_000_000)), + message("execution_success", timestamp=300), + ] + ) + result = classify_evidence(JOB, command.evidence) + self.assertEqual( + (result.outcome, result.reason), ("succeeded", "terminal_success") + ) + self.assertLessEqual(len(encode_command(command)), 2048) + + def test_encoder_rejects_oversized_kind_without_hashing(self): + class OversizedKind(str): + def __hash__(self): + raise AssertionError("Oversized event kind must not be hashed") + + command = finish([message("execution_success")]) + event = LifecycleEvent( + OversizedKind("x" * 1_000_000), JOB, None, None, None, None + ) + with self.assertRaises(ValueError): + encode_command( + replace( + command, evidence=TerminalEvidence((event,), False, False, False) + ) + ) diff --git a/tests/test_execution_history_hooks.py b/tests/test_execution_history_hooks.py new file mode 100644 index 0000000..f54b56f --- /dev/null +++ b/tests/test_execution_history_hooks.py @@ -0,0 +1,450 @@ +"""Actual dequeue and completion capture through the owned SQLite worker.""" + +from collections import namedtuple +import logging +import gc +import contextlib +import io +import os +import tempfile +import threading +import time +import sys +import types +import weakref +from types import SimpleNamespace +import unittest +from unittest.mock import patch + +from database import PromptDatabase +from utils import usage_tracking +from utils.execution_history_writer import HistoryWriter + +try: + from utils.execution_history import register_history_hook +except ImportError: + register_history_hook = None + +Status = namedtuple("ExecutionStatus", "status_str completed messages") +JOB = "a0c26da3-eb26-405c-b143-c662a84c3112" + + +class HostQueue: + def __init__(self): + self.results = [] + self.done_result = object() + self.done_calls = [] + + def get(self, *args, **kwargs): + return self.results.pop(0) if self.results else None + + def task_done(self, *args, **kwargs): + self.done_calls.append((args, kwargs)) + return self.done_result + + +def item(item_id, job=JOB, graph=None, extra=None, sensitive=None): + return ((0, job, graph or {}, extra or {}, [], sensitive or {}), item_id) + + +def status(kind="execution_success", job=JOB): + return Status("success", True, [(kind, {"prompt_id": job, "timestamp": 20})]) + + +class HistoryHooksTests(unittest.TestCase): + def setUp(self): + self.directory = tempfile.TemporaryDirectory() + self.addCleanup(self.directory.cleanup) + self.db = PromptDatabase(os.path.join(self.directory.name, "history.db")) + self.addCleanup(self.db.model.reset_schema_cache) + self.addCleanup(self.db.model.close_all) + self.queue = HostQueue() + self.server = SimpleNamespace(prompt_queue=self.queue) + self.registrations, self.releases = [], [] + self.addCleanup(self.cleanup_registrations) + for name in ("execution_history_writer", "execution_history"): + logger = logging.getLogger("prompt_manager." + name) + previous = logger.handlers, logger.propagate, logger.level + logger.handlers, logger.propagate, logger.level = ( + [logging.NullHandler()], + False, + logging.WARNING, + ) + self.addCleanup(self.restore_logger, logger, previous) + + def restore_logger(self, logger, previous): + logger.handlers, logger.propagate, logger.level = previous + + def cleanup_registrations(self): + for release in self.releases: + release.set() + for registration in self.registrations: + self.assertTrue(registration.dispose(3), "owned worker leaked") + + def wait_for(self, predicate): + deadline = time.monotonic() + 3 + while time.monotonic() < deadline: + if predicate(): + return + threading.Event().wait(0.005) + self.fail("finite deadline expired waiting for recorded history") + + def register(self, factory=None, ready=True): + self.assertTrue( + callable(register_history_hook), "history registration is absent" + ) + registration = register_history_hook(self.server, factory or (lambda: self.db)) + self.assertIsNotNone(registration) + self.registrations.append(registration) + if ready: + self.wait_for(lambda: registration.writer.status().accepting) + return registration + + def dequeue(self, item_id=0, **kwargs): + result = item(item_id, **kwargs) + self.queue.results.append(result) + self.assertIs(self.queue.get(), result) + return result + + def rows(self): + return [ + dict(row) + for row in self.db.model.get_connection() + .execute("SELECT * FROM execution_runs ORDER BY queue_item_id") + .fetchall() + ] + + def drain(self, registration): + self.wait_for(lambda: registration.writer.status().pending == 0) + + def test_package_history_initializes_when_monitor_and_usage_setup_fail(self): + try: + from tests.test_init_package import ( + fake_server_module, + load_package, + ) + except ImportError: + from test_init_package import ( + fake_server_module, + load_package, + ) + name = "prompt_manager_pkg_history_independent" + server = fake_server_module() + server.PromptServer.instance.prompt_queue = self.queue + server.PromptServer.instance.app = SimpleNamespace(middlewares=[]) + + def reject_usage(_): + raise RuntimeError("usage registration unavailable") + + server.PromptServer.instance.add_on_prompt_handler = reject_usage + output = os.path.join(self.directory.name, "output") + os.makedirs(output) + monitor_module = types.ModuleType(name + ".utils.image_monitor") + # Missing monitor import fails before shared database initialization. + with ( + patch.dict(sys.modules, {monitor_module.__name__: monitor_module}), + contextlib.redirect_stdout(io.StringIO()), + self.assertLogs("prompt_manager.init", level="ERROR"), + ): + with load_package(name, self.db.model.db_path, output, server) as ( + package, + _, + ): + registration = getattr(package, "_history_registration", None) + self.assertIsNotNone( + registration, "history setup depended on another hook" + ) + self.assertIsNone(package._global_db) + self.wait_for(lambda: registration.writer.status().accepting) + self.dequeue(0) + self.queue.task_done(0, {}, status()) + self.drain(registration) + self.assertEqual(self.rows()[0]["outcome"], "succeeded") + + def test_distinct_actual_attempts_for_reused_job_and_cached_rerun(self): + registration = self.register() + self.dequeue(0) + self.queue.task_done(0, {"outputs": {}}, status("execution_error")) + self.dequeue(1) + self.queue.task_done(item_id=1, history_result={}, status=status()) + self.drain(registration) + rows = self.rows() + self.assertEqual(len(rows), 2) + self.assertEqual([r["job_id"] for r in rows], [JOB, JOB]) + self.assertEqual([r["outcome"] for r in rows], ["failed", "succeeded"]) + self.assertNotEqual(rows[0]["run_id"], rows[1]["run_id"]) + + def test_later_conflicting_completion_resolves_completed_stored_identity(self): + registration = self.register() + self.dequeue(0) + self.queue.task_done(0, {}, status()) + self.drain(registration) + self.queue.task_done(0, {}, status("execution_error")) + self.drain(registration) + self.assertEqual(self.rows()[0]["outcome"], "unknown") + self.assertEqual(self.rows()[0]["reason"], "conflicting_terminal") + + def test_repeated_registration_cleanup_and_reregistration_keep_session_and_wrappers( + self, + ): + before = self.queue.get + calls = [] + + def unrelated_before(*a, **k): + calls.append("before") + return before(*a, **k) + + self.queue.get = unrelated_before + first = self.register() + self.assertIs(register_history_hook(self.server, lambda: self.db), first) + observed = self.queue.get + + def unrelated_after(*a, **k): + calls.append("after") + return observed(*a, **k) + + self.queue.get = unrelated_after + self.dequeue(0) + self.queue.task_done(0, {}, status()) + self.assertTrue(first.dispose()) + self.assertTrue(first.dispose()) + self.assertIs(self.queue.get, unrelated_after) + second = self.register() + self.dequeue(1) + self.queue.task_done(1, {}, status()) + self.drain(second) + self.assertEqual(len(self.rows()), 2) + self.assertEqual(self.rows()[0]["session_id"], self.rows()[1]["session_id"]) + self.assertEqual(calls, ["after", "before"] * 2) + + def test_concurrent_handle_capacity_never_reassigns_active_attempts(self): + registration = self.register() + self.queue.results.extend(item(index) for index in range(129)) + ready = threading.Barrier(130) + errors = [] + + def dequeue(): + try: + ready.wait(3) + self.queue.get() + except Exception as error: + errors.append(error) + + threads = [threading.Thread(target=dequeue) for _ in range(129)] + for thread in threads: + thread.start() + ready.wait(3) + for thread in threads: + thread.join(3) + self.assertFalse(thread.is_alive()) + self.assertEqual(errors, []) + self.drain(registration) + rows = self.rows() + admitted = {row["queue_item_id"] for row in rows} + self.assertEqual(len(rows), 128) + missing = (set(range(129)) - admitted).pop() + oldest = min(admitted) + self.queue.task_done(oldest, {}, status()) + self.queue.task_done(missing, {}, status()) + self.drain(registration) + by_item = {row["queue_item_id"]: row for row in self.rows()} + self.assertEqual(by_item[oldest]["outcome"], "succeeded") + self.assertNotIn(missing, by_item) + self.assertGreater(registration.writer.status().recording_gaps, 0) + self.dequeue(129) + self.queue.task_done(129, {}, status()) + self.drain(registration) + self.assertEqual(self.rows()[-1]["queue_item_id"], 129) + + def test_active_handles_do_not_retain_graph_extra_or_sensitive_envelopes(self): + class Envelope(dict): + pass + + registration = self.register() + graph, extra, sensitive = ( + Envelope({"g": object()}), + Envelope({"e": object()}), + Envelope({"s": object()}), + ) + references = [weakref.ref(value) for value in (graph, extra, sensitive)] + self.dequeue(0, graph=graph, extra=extra, sensitive=sensitive) + self.drain(registration) + del graph, extra, sensitive + gc.collect() + self.assertEqual([reference() for reference in references], [None, None, None]) + self.queue.task_done(0, {}, status()) + self.drain(registration) + self.assertEqual(self.rows()[0]["outcome"], "succeeded") + + def test_capture_uses_no_database_methods_on_the_host_thread(self): + registration = self.register() + host_thread = threading.get_ident() + original = self.db.model.get_connection + + def connection(): + if threading.get_ident() == host_thread: + self.fail("history queried SQLite on the host thread") + return original() + + with patch.object(self.db.model, "get_connection", connection): + self.dequeue(0) + self.queue.task_done(0, {}, status()) + self.drain(registration) + self.assertEqual(self.rows()[0]["outcome"], "succeeded") + + def test_capture_never_copies_or_serializes_sensitive_objects(self): + class Sensitive: + def __deepcopy__(self, memo): + raise AssertionError("sensitive copy") + + def __repr__(self): + raise AssertionError("sensitive serialization") + + registration = self.register() + secret = Sensitive() + graph, extra, sensitive = {"graph": secret}, {"extra": secret}, {"auth": secret} + result = self.dequeue(0, graph=graph, extra=extra, sensitive=sensitive) + payload = {"outputs": secret, "meta": secret} + completion = status() + completion.messages[0][1].update( + {"traceback": secret, "current_inputs": secret} + ) + self.assertIs( + self.queue.task_done(0, payload, completion), self.queue.done_result + ) + self.drain(registration) + self.assertIs(result[0][2], graph) + self.assertIs(result[0][3], extra) + self.assertIs(result[0][5], sensitive) + self.assertEqual(len(self.rows()), 1) + self.assertEqual(self.rows()[0]["outcome"], "succeeded") + self.assertEqual( + set(completion.messages[0][1]), + {"prompt_id", "timestamp", "traceback", "current_inputs"}, + ) + + def test_callbacks_return_while_writer_blocked_and_disposal_is_bounded(self): + entered, release = threading.Event(), threading.Event() + self.releases.append(release) + original = self.db.write_execution_start + + def blocked(command): + entered.set() + release.wait(5) + return original(command) + + with patch.object(self.db, "write_execution_start", blocked): + registration = self.register() + self.dequeue(0) + self.assertTrue(entered.wait(1)) + began = time.monotonic() + self.queue.task_done(0, {}, status()) + self.dequeue(1) + self.assertLess(time.monotonic() - began, 0.2) + began = time.monotonic() + self.assertFalse(registration.dispose()) + self.assertGreaterEqual(time.monotonic() - began, 1.8) + self.assertLess(time.monotonic() - began, 2.3) + release.set() + self.assertTrue(registration.dispose(2)) + + def test_restore_disables_capture_and_clears_handles_before_fresh_epoch(self): + try: + from tests.test_execution_history_writer import ManualClock + except ImportError: + from test_execution_history_writer import ManualClock + clock = ManualClock() + with patch( + "utils.execution_history.HistoryWriter", + lambda factory: HistoryWriter(factory, clock=clock), + ): + registration = self.register() + for index in range(128): + self.dequeue(index) + self.drain(registration) + old_epoch = registration.writer.status().epoch + backup = os.path.join(self.directory.name, "backup.db") + self.assertTrue(self.db.model.backup_database(backup)) + self.db.model.restore_from_file(backup) + self.assertFalse(registration.writer.status().accepting) + self.assertIsNone(registration.writer.status().epoch) + self.queue.task_done(0, {}, status()) + self.dequeue(128) + self.assertGreater(registration.writer.status().recording_gaps, 0) + clock.advance(5) + self.wait_for(lambda: registration.writer.status().accepting) + self.assertNotEqual(registration.writer.status().epoch, old_epoch) + # A late final descriptor cannot complete rows from the restored epoch. + self.queue.task_done(0, {}, status()) + self.dequeue(129) + self.queue.task_done(129, {}, status()) + self.drain(registration) + rows = self.rows() + self.assertEqual(len(rows), 129) + self.assertEqual(rows[0]["outcome"], "observed") + self.assertEqual(rows[-1]["queue_item_id"], 129) + self.assertEqual(rows[-1]["outcome"], "succeeded") + + def test_unavailable_database_and_invalid_identity_never_escape_host(self): + entered = threading.Event() + + def broken(): + entered.set() + raise RuntimeError("secret unavailable database") + + registration = self.register(broken, ready=False) + self.assertTrue(entered.wait(1)) + result = self.dequeue(0) + self.assertEqual(result[1], 0) + self.queue.task_done(0, {}, status()) + self.assertEqual(self.rows(), []) + self.assertGreater(registration.writer.status().recording_gaps, 0) + self.assertTrue(registration.dispose()) + registration = self.register() + self.dequeue(0, job="invalid id") + self.queue.task_done(True, {}, object()) + self.drain(registration) + self.assertEqual(self.rows(), []) + self.assertGreater(registration.writer.status().recording_gaps, 0) + + def test_invalid_evidence_is_recorded_unknown_and_raising_status_is_isolated(self): + registration = self.register() + self.dequeue(0) + self.queue.task_done(0, {}, object()) + self.drain(registration) + self.assertEqual(self.rows()[0]["outcome"], "unknown") + self.assertEqual(self.rows()[0]["reason"], "invalid_evidence") + + class BrokenStatus: + @property + def messages(self): + raise RuntimeError("secret failure") + + self.dequeue(1) + self.queue.task_done(1, {}, BrokenStatus()) + self.drain(registration) + self.assertEqual(self.rows()[1]["outcome"], "observed") + self.assertGreater(registration.writer.status().recording_gaps, 0) + + def test_unsupported_queue_does_not_create_writer_or_submit_fallback(self): + self.assertTrue( + callable(register_history_hook), "history registration is absent" + ) + for host in ( + object(), + SimpleNamespace(prompt_queue=SimpleNamespace(get=lambda: None)), + ): + self.assertIsNone( + register_history_hook(host, lambda: self.fail("DB touched")) + ) + + def test_existing_usage_callback_counts_once_alongside_history(self): + registration = self.register() + seen = [] + self.assertTrue(usage_tracking._wrap_queue_get(self.queue, seen.append)) + self.assertFalse(usage_tracking._wrap_queue_get(self.queue, seen.append)) + self.dequeue(0) + self.queue.task_done(0, {}, status()) + self.drain(registration) + self.assertEqual(seen, [JOB]) + self.assertEqual(len(self.rows()), 1) diff --git a/tests/test_execution_history_restore.py b/tests/test_execution_history_restore.py new file mode 100644 index 0000000..601287c --- /dev/null +++ b/tests/test_execution_history_restore.py @@ -0,0 +1,246 @@ +"""Restore publishes a fresh epoch and preserves historical attempt records.""" + +import sqlite3 +import unittest +from contextlib import closing +from pathlib import Path +from unittest.mock import patch + +from database.models import PromptModel + +try: + from tests import test_execution_history_database as history_tests +except ImportError: + import test_execution_history_database as history_tests + + +class HistoryRestoreTests(unittest.TestCase): + setUp = history_tests.HistoryDatabaseTests.setUp + close_databases = history_tests.HistoryDatabaseTests.close_databases + database = history_tests.HistoryDatabaseTests.database + start = history_tests.HistoryDatabaseTests.start + finish = history_tests.HistoryDatabaseTests.finish + count = history_tests.HistoryDatabaseTests.count + + def test_committed_wal_backup_includes_history_and_epoch(self): + db = self.database() + conn = db.model.get_connection() + conn.execute("PRAGMA wal_autocheckpoint=0") + start = self.start(db) + db.write_execution_start(start) + db.write_execution_finish(self.finish(start)) + row = db.get_execution_run(start.run_id) + backup = str(Path(self.temp.name) / "backup.db") + self.assertGreater(Path(self.path + "-wal").stat().st_size, 0) + self.assertTrue(db.model.backup_database(backup)) + restored = self.database(backup) + self.assertEqual(restored.get_execution_epoch(), start.storage_epoch) + self.assertEqual(restored.get_execution_run(start.run_id), row) + + def test_restore_fences_stale_commands_preserves_rows_and_notifies_after_publish( + self, + ): + from database.execution_history import ( + HistoryEpochChanged, + subscribe_history_restore, + ) + + db = self.database() + start = self.start(db) + db.write_execution_start(start) + backup = str(Path(self.temp.name) / "snapshot.db") + self.assertTrue(db.model.backup_database(backup)) + snapshot = Path(backup).read_bytes() + old_row = db.get_execution_run(start.run_id) + notices = [] + unsubscribe = subscribe_history_restore( + str(Path(self.temp.name) / "." / "live.db"), + lambda: notices.append(db.get_execution_epoch()), + ) + self.addCleanup(unsubscribe) + for _ in range(2): + previous = db.get_execution_epoch() + db.model.restore_from_file(backup) + epoch = db.get_execution_epoch() + self.assertNotEqual(epoch, previous) + self.assertNotEqual(epoch, start.storage_epoch) + self.assertEqual(notices[-1], epoch) + self.assertEqual(db.get_execution_run(start.run_id), old_row) + for command, write in ( + (start, db.write_execution_start), + (self.finish(start), db.write_execution_finish), + ): + with self.assertRaises(HistoryEpochChanged) as rejected: + write(command) + self.assertEqual(rejected.exception.storage_epoch, epoch) + self.assertEqual(db.get_execution_run(start.run_id), old_row) + fresh = self.start(db, item=1) + self.assertTrue(db.write_execution_start(fresh)) + self.assertTrue(db.write_execution_finish(self.finish(fresh))) + self.assertEqual(Path(backup).read_bytes(), snapshot) + unsubscribe() + unsubscribe() + db.model.restore_from_file(backup) + self.assertEqual(len(notices), 2) + + def test_old_backup_migrates_with_no_runs_without_changing_source(self): + source = str(Path(self.temp.name) / "old.db") + with closing(sqlite3.connect(source)) as conn, conn: + conn.execute( + "CREATE TABLE prompts (id INTEGER PRIMARY KEY AUTOINCREMENT," + " text TEXT NOT NULL, created_at TIMESTAMP, updated_at TIMESTAMP," + " category TEXT, tags TEXT, rating INTEGER, notes TEXT," + " hash TEXT UNIQUE, workflow_name TEXT)" + ) + conn.execute( + "INSERT INTO prompts(id,text,created_at) " + "VALUES(7,'old prompt','2020-01-01')" + ) + original = Path(source).read_bytes() + self.assertEqual(PromptModel.verify_database_file(source), (True, "ok")) + db = self.database() + previous = db.get_execution_epoch() + db.model.restore_from_file(source) + self.assertNotEqual(db.get_execution_epoch(), previous) + self.assertEqual(self.count(db), 0) + self.assertEqual(Path(source).read_bytes(), original) + self.assertEqual( + tuple( + db.model.get_connection() + .execute("SELECT id,text FROM prompts") + .fetchone() + ), + (7, "old prompt"), + ) + + def test_failed_staging_rotation_keeps_live_epoch_and_never_notifies(self): + from database.execution_history import subscribe_history_restore + + db = self.database() + start = self.start(db) + db.write_execution_start(start) + source = str(Path(self.temp.name) / "source.db") + db.model.backup_database(source) + source_bytes = Path(source).read_bytes() + notices = [] + self.addCleanup( + subscribe_history_restore(self.path, lambda: notices.append(True)) + ) + with patch( + "database.models.rotate_history_epoch", + side_effect=sqlite3.OperationalError("injected staging failure"), + ): + with self.assertRaisesRegex(sqlite3.OperationalError, "injected staging"): + db.model.restore_from_file(source) + self.assertEqual(db.get_execution_epoch(), start.storage_epoch) + self.assertEqual(Path(source).read_bytes(), source_bytes) + self.assertEqual(notices, []) + self.assertTrue(db.write_execution_finish(self.finish(start))) + + def test_malformed_history_rejected_before_replacement(self): + from database.execution_history import subscribe_history_restore + + db = self.database() + epoch = db.get_execution_epoch() + source = str(Path(self.temp.name) / "bad.db") + other = self.database(source) + with other.model.get_connection() as conn: + conn.execute("DROP INDEX idx_execution_runs_identity") + notices = [] + self.addCleanup( + subscribe_history_restore(self.path, lambda: notices.append(True)) + ) + with self.assertRaises(ValueError): + db.model.restore_from_file(source) + self.assertEqual(db.get_execution_epoch(), epoch) + self.assertEqual(notices, []) + + def test_extra_unique_constraints_reject_restore_without_touching_live_or_backup( + self, + ): + from database.execution_history import subscribe_history_restore + + db = self.database() + start = self.start(db) + db.write_execution_start(start) + row = db.get_execution_run(start.run_id) + with db.model.get_connection() as conn: + conn.execute("INSERT INTO prompts(text) VALUES('live library')") + notices = [] + self.addCleanup( + subscribe_history_restore(self.path, lambda: notices.append(True)) + ) + for table_constraint in (False, True): + with self.subTest(table_constraint=table_constraint): + source = str( + Path(self.temp.name) / f"unique-backup-{table_constraint}.db" + ) + self.assertTrue(db.model.backup_database(source)) + with closing(sqlite3.connect(source)) as conn, conn: + if table_constraint: + sql = conn.execute( + "SELECT sql FROM sqlite_master " + "WHERE name='execution_runs'" + ).fetchone()[0] + conn.execute("ALTER TABLE execution_runs RENAME TO old_runs") + conn.execute(sql[:-1] + ", UNIQUE(job_id))") + conn.execute( + "INSERT INTO execution_runs SELECT * FROM old_runs" + ) + conn.execute("DROP TABLE old_runs") + conn.execute( + "CREATE UNIQUE INDEX idx_execution_runs_identity " + "ON execution_runs(session_id, queue_item_id)" + ) + conn.execute( + "CREATE INDEX idx_execution_runs_observed " + "ON execution_runs(observed_dequeue_at, run_id)" + ) + else: + conn.execute( + "CREATE UNIQUE INDEX unexpected_job_uniqueness " + "ON execution_runs(job_id)" + ) + original = Path(source).read_bytes() + with self.assertRaisesRegex(ValueError, "uniqu"): + db.model.restore_from_file(source) + self.assertEqual(db.get_execution_epoch(), start.storage_epoch) + self.assertEqual(db.get_execution_run(start.run_id), row) + self.assertEqual(self.count(db), 1) + self.assertEqual( + db.model.get_connection() + .execute("SELECT text FROM prompts") + .fetchone()[0], + "live library", + ) + self.assertEqual(Path(source).read_bytes(), original) + self.assertEqual(notices, []) + self.assertTrue(db.write_execution_start(self.start(db, item=1))) + self.assertEqual(self.count(db), 2) + + def test_observer_failure_does_not_change_completed_restore(self): + from database.execution_history import subscribe_history_restore + + db = self.database() + source = str(Path(self.temp.name) / "snapshot.db") + db.model.backup_database(source) + original_epoch = db.get_execution_epoch() + + def broken(): + raise RuntimeError("observer failure") + + self.addCleanup(subscribe_history_restore(self.path, broken)) + notices = [] + self.addCleanup( + subscribe_history_restore(self.path, lambda: notices.append(True)) + ) + with self.assertLogs( + "prompt_manager.database.execution_history", level="WARNING" + ): + db.model.restore_from_file(source) + self.assertNotEqual(db.get_execution_epoch(), original_epoch) + self.assertEqual(notices, [True]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_execution_history_writer.py b/tests/test_execution_history_writer.py new file mode 100644 index 0000000..3ad2150 --- /dev/null +++ b/tests/test_execution_history_writer.py @@ -0,0 +1,486 @@ +"""Real-thread, disposable-database coverage for the bounded history writer.""" + +from dataclasses import FrozenInstanceError, replace +import logging +import os +import sqlite3 +import tempfile +import threading +import time +from types import SimpleNamespace +import unittest +from unittest.mock import patch +from uuid import uuid4 + +from database import PromptDatabase +from utils.execution_history_events import ( + LifecycleEvent, + RunFinish, + TerminalEvidence, + make_finish, + make_start, +) +from utils.execution_history_writer import HistoryWriter + + +class ManualClock: + def __init__(self): + self.value = 0.0 + self.lock = threading.Lock() + + def __call__(self): + with self.lock: + return self.value + + def advance(self, seconds): + with self.lock: + self.value += seconds + + +class HistoryWriterTests(unittest.TestCase): + def setUp(self): + logger = logging.getLogger("prompt_manager.execution_history_writer") + handlers, propagate, level = logger.handlers, logger.propagate, logger.level + logger.handlers = [logging.NullHandler()] + logger.propagate = False + logger.setLevel(logging.WARNING) + + def restore_logging(): + logger.handlers, logger.propagate, logger.level = handlers, propagate, level + + self.addCleanup(restore_logging) + self.directory = tempfile.TemporaryDirectory() + self.addCleanup(self.directory.cleanup) + self.db = PromptDatabase(os.path.join(self.directory.name, "history.db")) + self.addCleanup(self.db.model.reset_schema_cache) + self.addCleanup(self.db.model.close_all) + self.clock = ManualClock() + self.session = str(uuid4()) + self.job = str(uuid4()) + self.releases = [] + self.writers = [] + self.addCleanup(self.cleanup_workers) + + def cleanup_workers(self): + for release in self.releases: + release.set() + for writer in self.writers: + self.assertTrue(writer.close(3), "test-created worker leaked") + + def wait_for(self, predicate): + deadline = time.monotonic() + 3 + while time.monotonic() < deadline: + if predicate(): + return + threading.Event().wait(0.005) + self.fail("finite deadline expired waiting for worker behavior") + + def writer(self, factory=None, ready=True): + writer = HistoryWriter(factory or (lambda: self.db), clock=self.clock) + self.writers.append(writer) + writer.start() + if ready: + self.wait_for(lambda: writer.status().accepting) + return writer + + def start_command(self, item=1, epoch=None): + return make_start( + self.session, + ((0, self.job), item), + epoch or self.db.get_execution_epoch(), + 100, + ) + + def finish_command(self, start): + return make_finish( + start.session_id, + start.queue_item_id, + start.storage_epoch, + SimpleNamespace(messages=[("execution_success", {"prompt_id": self.job})]), + 200, + ) + + def block(self): + entered, release = threading.Event(), threading.Event() + self.releases.append(release) + return entered, release + + def test_fifo_single_worker_and_thread_local_connection_cleanup(self): + caller_conn = self.db.model.get_connection() + threads, timeouts = [], [] + original = self.db.write_execution_start + + def write(command): + threads.append(threading.get_ident()) + timeouts.append( + self.db.model.get_connection() + .execute("PRAGMA busy_timeout") + .fetchone()[0] + ) + return original(command) + + with patch.object(self.db, "write_execution_start", write): + writer = self.writer() + writer.start() + start = self.start_command() + self.assertTrue(writer.submit(start)) + self.assertTrue(writer.submit(self.finish_command(start))) + self.wait_for( + lambda: (self.db.get_execution_run(start.run_id) or {}).get("outcome") + == "succeeded" + ) + self.assertTrue(writer.submit(self.start_command(2))) + self.assertTrue(writer.close(2)) + self.assertEqual(len(set(threads)), 1) + self.assertNotEqual(threads[0], threading.get_ident()) + self.assertEqual(timeouts, [250, 250]) + self.assertFalse(writer.status().alive) + self.assertFalse(writer.status().accepting) + self.assertEqual(writer.status().pending, 0) + self.assertEqual( + caller_conn.execute("PRAGMA busy_timeout").fetchone()[0], 30000 + ) + self.assertEqual(len(self.db.model._connections), 1) + with self.assertRaises(FrozenInstanceError): + writer.status().epoch = "changed" + self.assertFalse(writer.submit(start)) + + def test_blocked_initialization_never_accepts_or_blocks_producers(self): + entered, release = self.block() + + def factory(): + entered.set() + release.wait(3) + return self.db + + writer = self.writer(factory, ready=False) + self.assertTrue(entered.wait(1)) + before = time.monotonic() + self.assertFalse(writer.submit(self.start_command())) + self.assertIsNone(writer.status().epoch) + self.assertLess(time.monotonic() - before, 0.1) + self.assertIsNone(self.db.get_execution_run(self.start_command().run_id)) + before = time.monotonic() + self.assertFalse(writer.close(0.03)) + self.assertLess(time.monotonic() - before, 0.2) + release.set() + self.wait_for(lambda: not writer.status().alive) + self.assertFalse(writer.status().accepting) + self.assertEqual(len(self.db.model._connections), 1) + + def test_saturation_bounds_all_pending_commands_and_coalesces_diagnostics(self): + entered, release = self.block() + original = self.db.write_execution_start + + def write(command): + entered.set() + release.wait(3) + return original(command) + + with patch.object(self.db, "write_execution_start", write): + writer = self.writer() + command = self.start_command() + self.assertTrue(writer.submit(command)) + self.assertTrue(entered.wait(1)) + for _ in range(1023): + self.assertTrue(writer.submit(command)) + for _ in range(100): + self.assertFalse(writer.submit(command)) + self.assertEqual(writer.status().pending, 1024) + self.assertEqual(writer.status().recording_gaps, 100) + with self.assertLogs( + "prompt_manager.execution_history_writer", "WARNING" + ) as logs: + release.set() + self.wait_for(lambda: writer.status().pending == 0) + self.assertTrue(writer.close(2)) + self.assertEqual(len(logs.records), 1) + + def test_rejects_oversized_invalid_and_wrong_epoch_commands(self): + writer = self.writer() + start = self.start_command() + oversized = RunFinish( + self.session, + 1, + start.storage_epoch, + 200, + TerminalEvidence( + tuple( + LifecycleEvent( + "execution_error", self.job, i, "n" * 128, None, None + ) + for i in range(20) + ), + False, + False, + False, + ), + ) + for command in ( + oversized, + object(), + replace(start, storage_epoch=str(uuid4())), + ): + self.assertFalse(writer.submit(command)) + self.assertEqual(writer.status().pending, 0) + self.assertEqual(writer.status().recording_gaps, 3) + self.assertIsNone(self.db.get_execution_run(start.run_id)) + + def test_failed_command_once_discards_backlog_and_waits_five_seconds(self): + entered, release = self.block() + probes, writes = [], [] + original_epoch = self.db.get_execution_epoch + original_write = self.db.write_execution_start + + def epoch(): + probes.append(self.clock()) + return original_epoch() + + def write(command): + writes.append(command.queue_item_id) + if command.queue_item_id == 1: + entered.set() + release.wait(3) + raise sqlite3.OperationalError("private failure text") + return original_write(command) + + start, queued, fresh = (self.start_command(i) for i in (1, 2, 3)) + with ( + patch.object(self.db, "get_execution_epoch", epoch), + patch.object(self.db, "write_execution_start", write), + ): + writer = self.writer() + self.assertTrue(writer.submit(start)) + self.assertTrue(entered.wait(1)) + self.assertTrue(writer.submit(queued)) + release.set() + self.wait_for(lambda: not writer.status().accepting) + self.assertFalse(writer.submit(fresh)) + self.clock.advance(4.99) + threading.Event().wait(0.08) + self.assertEqual(probes, [0]) + self.clock.advance(0.01) + self.wait_for(lambda: writer.status().accepting) + self.assertEqual(probes, [0, 5]) + self.assertTrue(writer.submit(fresh)) + self.wait_for(lambda: self.db.get_execution_run(fresh.run_id) is not None) + self.assertEqual(writes, [1, 3]) + self.assertIsNone(self.db.get_execution_run(start.run_id)) + self.assertIsNone(self.db.get_execution_run(queued.run_id)) + + def test_factory_failure_recovery_probes_are_bounded(self): + calls = [] + + def factory(): + calls.append(self.clock()) + if len(calls) < 3: + raise sqlite3.OperationalError("unavailable") + return self.db + + writer = self.writer(factory, ready=False) + self.wait_for(lambda: len(calls) == 1) + self.clock.advance(4.99) + threading.Event().wait(0.08) + self.assertEqual(calls, [0]) + self.clock.advance(0.01) + self.wait_for(lambda: len(calls) == 2) + self.clock.advance(4.99) + threading.Event().wait(0.08) + self.assertEqual(calls, [0, 5]) + self.clock.advance(0.01) + self.wait_for(lambda: writer.status().accepting) + self.assertEqual(calls, [0, 5, 10]) + command = self.start_command() + self.assertTrue(writer.submit(command)) + self.wait_for(lambda: self.db.get_execution_run(command.run_id) is not None) + + def test_failed_epoch_probe_closes_only_worker_connection_and_recovers(self): + command = self.start_command() + original = self.db.get_execution_epoch + failed = threading.Event() + + def epoch(): + if not failed.is_set(): + failed.set() + raise sqlite3.DatabaseError("metadata unavailable") + return original() + + with patch.object(self.db, "get_execution_epoch", epoch): + writer = self.writer(ready=False) + self.wait_for(lambda: writer.status().recording_gaps == 1) + self.assertIsNone(writer.status().epoch) + self.assertFalse(writer.submit(command)) + self.assertEqual(len(self.db.model._connections), 1) + self.clock.advance(5) + self.wait_for(lambda: writer.status().accepting) + self.assertTrue(writer.submit(command)) + self.wait_for(lambda: self.db.get_execution_run(command.run_id) is not None) + + def test_invalidation_during_epoch_probe_cannot_publish_stale_readiness(self): + entered, release = self.block() + original = self.db.get_execution_epoch + stale_epoch = original() + fresh_epoch = str(uuid4()) + + def epoch(): + result = original() + if not entered.is_set(): + entered.set() + release.wait(3) + return result + + with patch.object(self.db, "get_execution_epoch", epoch): + writer = self.writer(ready=False) + self.assertTrue(entered.wait(1)) + with self.db.model.get_connection() as conn: + conn.execute( + "UPDATE execution_history_metadata SET storage_epoch=?", + (fresh_epoch,), + ) + writer.invalidate() + release.set() + threading.Event().wait(0.08) + self.assertFalse(writer.status().accepting) + self.assertIsNone(writer.status().epoch) + self.assertFalse(writer.submit(self.start_command(epoch=stale_epoch))) + self.clock.advance(5) + self.wait_for(lambda: writer.status().epoch == fresh_epoch) + command = self.start_command(epoch=fresh_epoch) + self.assertTrue(writer.submit(command)) + self.wait_for(lambda: self.db.get_execution_run(command.run_id) is not None) + + def test_held_sqlite_lock_does_not_block_submit_and_is_not_retried(self): + writer = self.writer() + command = self.start_command() + lock = sqlite3.connect(self.db.model.db_path) + self.addCleanup(lock.close) + lock.execute("BEGIN IMMEDIATE") + before = time.monotonic() + self.assertTrue(writer.submit(command)) + self.assertLess(time.monotonic() - before, 0.1) + self.wait_for(lambda: not writer.status().accepting) + lock.rollback() + self.clock.advance(5) + self.wait_for(lambda: writer.status().accepting) + self.assertIsNone(self.db.get_execution_run(command.run_id)) + fresh = self.start_command(2) + self.assertTrue(writer.submit(fresh)) + self.wait_for(lambda: self.db.get_execution_run(fresh.run_id) is not None) + + def test_external_epoch_rotation_drops_stale_commands_then_records_fresh(self): + writer = self.writer() + command = self.start_command() + new_epoch = str(uuid4()) + with self.db.model.get_connection() as conn: + conn.execute( + "UPDATE execution_history_metadata SET storage_epoch=?", (new_epoch,) + ) + self.assertTrue(writer.submit(command)) + self.wait_for(lambda: not writer.status().accepting) + self.assertIsNone(writer.status().epoch) + self.clock.advance(5) + self.wait_for(lambda: writer.status().epoch == new_epoch) + self.assertFalse(writer.submit(command)) + fresh = self.start_command(2) + self.assertTrue(writer.submit(fresh)) + self.wait_for(lambda: self.db.get_execution_run(fresh.run_id) is not None) + self.assertIsNone(self.db.get_execution_run(command.run_id)) + + def test_restore_invalidates_queued_and_inflight_commands(self): + backup = os.path.join(self.directory.name, "backup.db") + command = self.start_command() + self.db.write_execution_start(command) + self.assertTrue(self.db.model.backup_database(backup)) + writer = self.writer() + entered, release = self.block() + original = self.db.write_execution_start + + def write(command): + entered.set() + release.wait(3) + return original(command) + + with patch.object(self.db, "write_execution_start", write): + self.assertTrue(writer.submit(command)) + self.assertTrue(entered.wait(1)) + self.assertTrue(writer.submit(self.finish_command(command))) + self.db.model.restore_from_file(backup) + self.assertFalse(writer.status().accepting) + self.assertIsNone(writer.status().epoch) + self.assertFalse(writer.submit(command)) + release.set() + self.wait_for(lambda: writer.status().pending == 0) + self.clock.advance(5) + self.wait_for(lambda: writer.status().accepting) + self.assertNotEqual(writer.status().epoch, command.storage_epoch) + self.assertEqual( + self.db.get_execution_run(command.run_id)["outcome"], "observed" + ) + fresh = self.start_command(2) + timeouts = [] + + def write_fresh(command): + timeouts.append( + self.db.model.get_connection() + .execute("PRAGMA busy_timeout") + .fetchone()[0] + ) + return original(command) + + with patch.object(self.db, "write_execution_start", write_fresh): + self.assertTrue(writer.submit(fresh)) + self.assertTrue(writer.submit(self.finish_command(fresh))) + self.wait_for( + lambda: (self.db.get_execution_run(fresh.run_id) or {}).get("outcome") + == "succeeded" + ) + self.assertEqual(timeouts, [250]) + + def test_bounded_close_discards_backlog_and_disposes_after_blocked_write(self): + entered, release = self.block() + original = self.db.write_execution_start + + def write(command): + entered.set() + release.wait(3) + return original(command) + + with patch.object(self.db, "write_execution_start", write): + writer = self.writer() + first, queued = self.start_command(), self.start_command(2) + self.assertTrue(writer.submit(first)) + self.assertTrue(entered.wait(1)) + self.assertTrue(writer.submit(queued)) + before = time.monotonic() + self.assertFalse(writer.close(0.03)) + self.assertLess(time.monotonic() - before, 0.2) + release.set() + self.wait_for(lambda: not writer.status().alive) + self.assertIsNone(self.db.get_execution_run(queued.run_id)) + self.assertEqual(writer.status().pending, 0) + self.assertEqual(len(self.db.model._connections), 1) + + def test_close_before_start_is_terminal(self): + writer = HistoryWriter(lambda: self.fail("closed writer initialized database")) + self.assertTrue(writer.close()) + writer.start() + self.assertFalse(writer.status().alive) + self.assertFalse(writer.submit(self.start_command())) + + def test_external_capture_gaps_are_nonblocking_and_saturate_all_rejections(self): + writer = HistoryWriter(lambda: self.fail("producer accessed database")) + self.assertTrue(callable(getattr(writer, "record_gap", None))) + writer.record_gap() + self.assertEqual(writer.status().recording_gaps, 1) + with writer._lock: + writer._gaps = 2**63 - 2 + writer.record_gap() + writer.record_gap() + writer.submit(self.start_command()) + writer.submit(object()) + self.assertEqual(writer.status().recording_gaps, 2**63 - 1) + self.assertTrue(writer.close()) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_extension_web_dir.py b/tests/test_extension_web_dir.py index b360f04..018a1cd 100644 --- a/tests/test_extension_web_dir.py +++ b/tests/test_extension_web_dir.py @@ -4,6 +4,7 @@ ComfyUI serves every .js file under WEB_DIRECTORY into the canvas page, so the directory must hold nothing but the node extension. """ +import contextlib import os import sys import tempfile @@ -14,7 +15,7 @@ ROOT = os.path.dirname(TESTS_DIR) sys.path.insert(0, ROOT) sys.path.insert(0, TESTS_DIR) -from test_init_package import close_every_database, load_package # noqa: E402 +from test_init_package import load_package # noqa: E402 class TestExtensionWebDir(unittest.TestCase): @@ -23,21 +24,17 @@ class TestExtensionWebDir(unittest.TestCase): @classmethod def setUpClass(cls): cls.tmp = tempfile.TemporaryDirectory() + cls.addClassCleanup(cls.tmp.cleanup) output_dir = os.path.join(cls.tmp.name, "output") os.makedirs(output_dir) - cls.package, _ = load_package( - cls.package_name, os.path.join(cls.tmp.name, "prompts.db"), output_dir + package_context = contextlib.ExitStack() + cls.addClassCleanup(package_context.close) + cls.package, _ = package_context.enter_context( + load_package( + cls.package_name, os.path.join(cls.tmp.name, "prompts.db"), output_dir + ) ) - @classmethod - def tearDownClass(cls): - monitor = getattr(cls.package, "_global_image_monitor", None) - if monitor is not None: - monitor.stop_monitoring() - sys.modules.pop(cls.package_name, None) - close_every_database() # the package's database must not hold the temp dir open - cls.tmp.cleanup() - def test_web_directory_is_the_canvas_extension_folder(self): self.assertEqual(self.package.WEB_DIRECTORY, "web/comfy") self.assertIn("WEB_DIRECTORY", self.package.__all__) diff --git a/tests/test_init_package.py b/tests/test_init_package.py index 39fc9bf..df95d4e 100644 --- a/tests/test_init_package.py +++ b/tests/test_init_package.py @@ -4,7 +4,11 @@ The package is loaded the way ComfyUI loads it (``__init__`` with relative imports) with ComfyUI's ``server`` and ``folder_paths`` modules stubbed. """ +import contextlib import importlib.util +import io +import logging +import json import os import sys import tempfile @@ -17,6 +21,10 @@ from aiohttp import web ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) sys.path.insert(0, ROOT) +from utils.logging_config import get_logger_manager # noqa: E402 + +get_logger_manager() + try: from tests.open_handles import open_handles_under # noqa: E402 except ImportError: # discovered with tests/ as the top-level directory @@ -37,6 +45,7 @@ def fake_server_module(): routes=web.RouteTableDef(), prompt_queue=FakePromptQueue(), handlers=[], + app=types.SimpleNamespace(middlewares=[]), ) instance.add_on_prompt_handler = instance.handlers.append module = types.ModuleType("server") @@ -53,10 +62,10 @@ def fake_folder_paths_module(output_dir): return module -def load_package(name, db_path, output_dir, server=None): +def _import_package(name, db_path, output_dir, server=None): """Import ROOT/__init__.py as package ``name`` with ComfyUI stubbed. - Returns (package, server instance). The caller stops the image monitor. + Returns (package, server instance) to the scoped fixture. """ saved = {key: sys.modules.get(key) for key in ("server", "folder_paths")} server = server or fake_server_module() @@ -65,7 +74,13 @@ def load_package(name, db_path, output_dir, server=None): try: from py.config import PromptManagerConfig - with mock.patch.object(PromptManagerConfig, "DEFAULT_DB_PATH", db_path): + config_path = os.path.join(output_dir, "promptmanager-test-config.json") + with open(config_path, "w", encoding="utf-8") as config: + json.dump({"database": {"default_path": db_path}}, config) + with ( + mock.patch.object(PromptManagerConfig, "DEFAULT_DB_PATH", db_path), + mock.patch.dict(os.environ, {"PROMPT_MANAGER_CONFIG_PATH": config_path}), + ): spec = importlib.util.spec_from_file_location( name, os.path.join(ROOT, "__init__.py"), @@ -76,7 +91,7 @@ def load_package(name, db_path, output_dir, server=None): try: spec.loader.exec_module(package) except BaseException: - sys.modules.pop(name, None) + cleanup_package(package) raise finally: for key, module in saved.items(): @@ -87,6 +102,73 @@ def load_package(name, db_path, output_dir, server=None): return package, server.PromptServer.instance +def cleanup_package(package): + """Stop fixture-owned workers before releasing its other resources.""" + registration = getattr(package, "_history_registration", None) + if registration is not None: + assert registration.dispose(3), "package history worker leaked" + monitor = getattr(package, "_global_image_monitor", None) + if monitor is not None: + monitor.stop_monitoring() + close_every_database() + sys.modules.pop(package.__name__, None) + + +@contextlib.contextmanager +def load_package(name, db_path, output_dir, server=None): + """Dispose package workers before closing its isolated logging handlers.""" + root_logger = logging.getLogger("prompt_manager") + loggers = [root_logger] + [ + logger + for key, logger in list(logging.Logger.manager.loggerDict.items()) + if key.startswith("prompt_manager.") and isinstance(logger, logging.Logger) + ] + previous = { + logger: ( + logger.handlers[:], + logger.level, + logger.propagate, + logger.disabled, + logger.filters[:], + ) + for logger in loggers + } + prior_handlers = {handler for state in previous.values() for handler in state[0]} + root_logger.handlers = [] + try: + package, instance = _import_package(name, db_path, output_dir, server) + # A previously imported package may reuse its existing logger singleton. + if not root_logger.handlers: + root_logger.handlers = previous[root_logger][0][:] + try: + yield package, instance + finally: + cleanup_package(package) + finally: + current_loggers = [root_logger] + [ + logger + for key, logger in list(logging.Logger.manager.loggerDict.items()) + if key.startswith("prompt_manager.") and isinstance(logger, logging.Logger) + ] + created_handlers = { + handler + for logger in current_loggers + for handler in logger.handlers + if handler not in prior_handlers + } + for logger in current_loggers: + ( + logger.handlers, + level, + logger.propagate, + logger.disabled, + logger.filters, + ) = previous.get(logger, ([], logging.NOTSET, True, False, [])) + logger.setLevel(level) + for handler in created_handlers: + handler.close() + + def close_every_database(): """Close connections of every PromptModel registry in the process. @@ -104,11 +186,17 @@ class PackageTestCase(unittest.TestCase): @classmethod def setUpClass(cls): cls.tmp = tempfile.TemporaryDirectory() + cls.addClassCleanup(cls.tmp.cleanup) + cls.addClassCleanup(cls.assert_no_open_handles) cls.output_dir = os.path.join(cls.tmp.name, "output") os.makedirs(cls.output_dir) cls.db_path = cls.make_db_path(cls.tmp.name) - cls.package, cls.server = load_package( - cls.package_name, cls.db_path, cls.output_dir, server=cls.make_server() + package_context = contextlib.ExitStack() + cls.addClassCleanup(package_context.close) + cls.package, cls.server = package_context.enter_context( + load_package( + cls.package_name, cls.db_path, cls.output_dir, server=cls.make_server() + ) ) @classmethod @@ -120,14 +208,8 @@ class PackageTestCase(unittest.TestCase): return fake_server_module() @classmethod - def tearDownClass(cls): - monitor = getattr(cls.package, "_global_image_monitor", None) - if monitor is not None: - monitor.stop_monitoring() - sys.modules.pop(cls.package_name, None) - close_every_database() + def assert_no_open_handles(cls): hits = open_handles_under(cls.tmp.name) - cls.tmp.cleanup() assert not hits, f"files still open before temp dir cleanup: {hits}" @@ -262,9 +344,6 @@ class TestConsoleBanner(unittest.TestCase): """Loading the node must not depend on the console being able to print emoji.""" def test_package_loads_with_a_cp1252_console(self): - import contextlib - import io - tmp = tempfile.TemporaryDirectory() self.addCleanup(tmp.cleanup) self.addCleanup(close_every_database) @@ -273,8 +352,12 @@ class TestConsoleBanner(unittest.TestCase): name = "prompt_manager_pkg_cp1252" stream = io.TextIOWrapper(io.BytesIO(), encoding="cp1252", errors="strict") + package_context = contextlib.ExitStack() + self.addCleanup(package_context.close) with contextlib.redirect_stdout(stream): - package, _ = load_package(name, os.path.join(tmp.name, "p.db"), output_dir) + package, _ = package_context.enter_context( + load_package(name, os.path.join(tmp.name, "p.db"), output_dir) + ) self.addCleanup(sys.modules.pop, name, None) monitor = getattr(package, "_global_image_monitor", None) if monitor is not None: @@ -283,3 +366,126 @@ class TestConsoleBanner(unittest.TestCase): stream.flush() banner = stream.buffer.getvalue().decode("cp1252") self.assertIn("Loaded:", banner) + + +class TestPackageLoggerIsolation(unittest.TestCase): + def setUp(self): + self.tmp = tempfile.TemporaryDirectory() + self.addCleanup(self.tmp.cleanup) + self.logger = logging.getLogger("prompt_manager") + self.previous = ( + self.logger.handlers[:], + self.logger.level, + self.logger.propagate, + self.logger.disabled, + ) + self.addCleanup(self.restore_logging) + self.path = os.path.join(self.tmp.name, "prior.log") + self.handler = logging.FileHandler(self.path) + self.logger.handlers = [self.handler] + self.logger.setLevel(logging.DEBUG) + self.logger.propagate = True + self.logger.disabled = True + self.child = logging.getLogger("prompt_manager.database") + self.child_previous = self.child.level + self.addCleanup(self.child.setLevel, self.child_previous) + self.child.setLevel(logging.ERROR) + self.saved_modules = { + key: sys.modules.get(key) for key in ("server", "folder_paths") + } + self.saved_config = os.environ.get("PROMPT_MANAGER_CONFIG_PATH") + self.server = fake_server_module() + self.server.PromptServer.instance.prompt_queue.task_done = lambda *a, **k: None + + def restore_logging(self): + for handler in self.logger.handlers: + if handler not in self.previous[0]: + handler.close() + self.handler.close() + ( + self.logger.handlers, + self.logger.level, + self.logger.propagate, + self.logger.disabled, + ) = self.previous + + def assert_logging_restored(self, handlers, registration): + for key, module in self.saved_modules.items(): + self.assertIs(sys.modules.get(key), module) + self.assertEqual( + os.environ.get("PROMPT_MANAGER_CONFIG_PATH"), self.saved_config + ) + self.assertEqual(self.logger.handlers, [self.handler]) + self.assertEqual(self.logger.level, logging.DEBUG) + self.assertTrue(self.logger.propagate) + self.assertTrue(self.logger.disabled) + self.assertEqual(self.child.level, logging.ERROR) + self.assertFalse(self.handler.stream.closed) + self.assertTrue(all(handler.stream is None for handler in handlers)) + self.assertFalse(registration.writer.status().alive) + self.logger.disabled = False + self.logger.warning("prior logging still works") + self.logger.disabled = True + self.handler.flush() + with open(self.path, encoding="utf-8") as log: + self.assertIn("prior logging still works", log.read()) + + def test_package_scope_closes_its_handlers_after_worker_disposal(self): + with contextlib.redirect_stdout(io.StringIO()): + with load_package( + "prompt_manager_pkg_logger_scope", + os.path.join(self.tmp.name, "p.db"), + self.tmp.name, + self.server, + ) as (package, _): + handlers = [ + h + for h in self.logger.handlers + if isinstance(h, logging.FileHandler) + ] + self.assertTrue(handlers) + registration = package._history_registration + original_dispose = registration.dispose + open_at_disposal = [] + + def dispose(timeout=2): + open_at_disposal.append(all(h.stream is not None for h in handlers)) + return original_dispose(timeout) + + registration.dispose = dispose + self.assertEqual(open_at_disposal, [True]) + self.assert_logging_restored(handlers, registration) + + def test_failed_package_import_closes_handlers_and_restores_logging(self): + name = "prompt_manager_pkg_logger_failure" + captured = {} + original_print = print + + def fail_banner(*args, **kwargs): + if any("Total:" in str(arg) for arg in args): + captured["package"] = sys.modules[name] + captured["handlers"] = [ + h + for h in self.logger.handlers + if isinstance(h, logging.FileHandler) + ] + raise RuntimeError("injected banner failure") + return original_print(*args, **kwargs) + + with ( + contextlib.redirect_stdout(io.StringIO()), + mock.patch("builtins.print", side_effect=fail_banner), + ): + with self.assertRaisesRegex(RuntimeError, "injected banner failure"): + with load_package( + name, + os.path.join(self.tmp.name, "p.db"), + self.tmp.name, + self.server, + ): + self.fail("failed import yielded a package") + self.assertNotIn(name, sys.modules) + self.assertTrue(captured["handlers"]) + self.assert_logging_restored( + captured["handlers"], captured["package"]._history_registration + ) diff --git a/tests/test_queue_observers.py b/tests/test_queue_observers.py new file mode 100644 index 0000000..e42b520 --- /dev/null +++ b/tests/test_queue_observers.py @@ -0,0 +1,132 @@ +"""Cooperative queue subscriptions preserve the host and other wrappers.""" + +import queue +import unittest + +try: + from utils.queue_observers import subscribe_dequeue, subscribe_task_done +except ImportError: + subscribe_dequeue = subscribe_task_done = None + + +class FakeQueue: + def __init__(self): + self.result = ((0, "job", {}, {}, [], {}), 0) + self.calls = [] + self.failure = None + self.done_result = object() + + def get(self, *args, **kwargs): + self.calls.append(("get", args, kwargs)) + if self.failure: + raise self.failure + return self.result + + def task_done(self, *args, **kwargs): + self.calls.append(("done", args, kwargs)) + if self.failure: + raise self.failure + return self.done_result + + +class QueueObserversTests(unittest.TestCase): + def require_api(self): + self.assertTrue(callable(subscribe_dequeue), "queue subscriptions are absent") + self.assertTrue(callable(subscribe_task_done), "task subscriptions are absent") + + def test_dequeue_preserves_identity_arguments_empty_and_exception(self): + self.require_api() + host, seen, sentinel = FakeQueue(), [], object() + subscribe_dequeue(host, "capture", seen.append) + original = host.result + self.assertIs(host.get(sentinel, timeout=sentinel), original) + self.assertEqual(host.calls, [("get", (sentinel,), {"timeout": sentinel})]) + self.assertEqual(seen, [original]) + host.result = None + self.assertIsNone(host.get()) + self.assertEqual(seen, [original]) + for error in (queue.Empty(), RuntimeError("host failure")): + host.failure = error + with self.assertRaises(type(error)) as raised: + host.get(timeout=sentinel) + self.assertIs(raised.exception, error) + self.assertEqual(seen, [original]) + + def test_task_done_observes_success_only_and_preserves_host_values(self): + self.require_api() + host, seen, sentinel = FakeQueue(), [], object() + subscribe_task_done(host, "capture", lambda a, k: seen.append((a, k))) + self.assertIs(host.task_done(0, sentinel, status=sentinel), host.done_result) + self.assertEqual(seen, [((0, sentinel), {"status": sentinel})]) + error = ValueError("host failure") + host.failure = error + with self.assertRaises(ValueError) as raised: + host.task_done(1) + self.assertIs(raised.exception, error) + self.assertEqual(len(seen), 1) + + def test_keys_are_idempotent_and_independently_removable(self): + self.require_api() + host, first, second = FakeQueue(), [], [] + unsubscribe = subscribe_dequeue(host, "a", first.append) + subscribe_dequeue(host, "a", lambda value: self.fail("duplicate callback")) + remove_second = subscribe_dequeue(host, "b", second.append) + host.get() + unsubscribe() + unsubscribe() + host.get() + self.assertEqual(len(first), 1) + self.assertEqual(len(second), 2) + remove_second() + subscribe_dequeue(host, "a", first.append) + unsubscribe() # old cleanup must not remove a new subscriber with this key + host.get() + self.assertEqual(len(first), 2) + + def test_other_wrappers_before_after_cleanup_and_resubscription_survive(self): + self.require_api() + host, calls, seen = FakeQueue(), [], [] + original = host.get + + def before(*a, **k): + calls.append("before") + return original(*a, **k) + + host.get = before + unsubscribe = subscribe_dequeue(host, "capture", seen.append) + observed = host.get + + def after(*a, **k): + calls.append("after") + return observed(*a, **k) + + host.get = after + host.get() + unsubscribe() + self.assertIs(host.get, after) + host.get() + subscribe_dequeue(host, "capture", seen.append) + host.get() + self.assertIs(host.get, after) + self.assertEqual(calls, ["after", "before"] * 3) + self.assertEqual(len(seen), 2) + + def test_failing_subscriber_cannot_disable_other_subscribers_or_host(self): + self.require_api() + host, seen = FakeQueue(), [] + + def fail(_): + raise ValueError("secret error") + + subscribe_dequeue(host, "bad", fail) + subscribe_dequeue(host, "good", seen.append) + with self.assertLogs("prompt_manager.queue_observers", level="WARNING") as logs: + self.assertIs(host.get(), host.result) + self.assertEqual(seen, [host.result]) + self.assertNotIn("secret error", "\n".join(logs.output)) + + def test_missing_methods_are_rejected_without_installing_a_fallback(self): + self.require_api() + for subscribe in (subscribe_dequeue, subscribe_task_done): + with self.assertRaises(ValueError): + subscribe(object(), "key", lambda *a: None) diff --git a/tests/test_usage_tracking.py b/tests/test_usage_tracking.py index f7070fb..eeabad5 100644 --- a/tests/test_usage_tracking.py +++ b/tests/test_usage_tracking.py @@ -373,6 +373,22 @@ class TestCountAtDequeue(DbTestCase): self.assertEqual(self._runs(positive), 1) self.assertEqual(calls, []) + def test_usage_subscription_remains_once_after_an_unrelated_wrapper(self): + positive = self._save("a cat") + prompt_queue = self.server.prompt_queue + original = prompt_queue.get + + def unrelated(*args, **kwargs): + return original(*args, **kwargs) + + prompt_queue.get = unrelated + self.assertFalse( + _wrap_queue_get(prompt_queue, lambda _: self.fail("duplicate")) + ) + self.assertIs(prompt_queue.get, unrelated) + self._dequeue(self._submit(workflow())) + self.assertEqual(self._runs(positive), 1) + def test_wrap_rejects_queues_without_a_callable_get(self): self.assertFalse(_wrap_queue_get(object(), lambda prompt_id: None)) self.assertFalse(_wrap_queue_get(None, lambda prompt_id: None)) diff --git a/tests/vendor/comfyui-config.test.mjs b/tests/vendor/comfyui-config.test.mjs index 3895637..48c6a6c 100644 --- a/tests/vendor/comfyui-config.test.mjs +++ b/tests/vendor/comfyui-config.test.mjs @@ -2,6 +2,7 @@ import test from 'node:test'; import assert from 'node:assert/strict'; import { execFileSync } from 'node:child_process'; import { fileURLToPath } from 'node:url'; +import { readFileSync } from 'node:fs'; const root = fileURLToPath(new URL('../../', import.meta.url)); test('real-ComfyUI suite registers Classic and Nodes 2.0 startup projects', () => { @@ -13,3 +14,34 @@ test('real-ComfyUI suite registers Classic and Nodes 2.0 startup projects', () = assert.match(output, /\[classic-start\]/); assert.match(output, /\[nodes2-start\]/); }); + +test('execution script discovers only the new matrix and preserves all 78 stock cases', () => { + const manifest = JSON.parse(readFileSync(new URL('../../tools/vendor/package.json', import.meta.url), 'utf8')); + assert.equal(manifest.scripts['test:executions'], 'playwright test --config playwright.execution.config.mjs'); + const list = config => execFileSync(process.execPath, ['tools/vendor/node_modules/@playwright/test/cli.js', + 'test', '--config', config, '--list'], { cwd: root, encoding: 'utf8', + env: { ...process.env, PM_COMFY_ROOT: '/unused-pinned-checkout' } }); + const stock = list('tools/vendor/comfyui.config.mjs'); + assert.match(stock, /Total: 78 tests in 7 files/); + assert.doesNotMatch(stock, /execution-history/); + const executions = list('tools/vendor/playwright.execution.config.mjs'); + assert.match(executions, /\[classic-start\]/); + assert.match(executions, /\[nodes2-start\]/); + assert.match(executions, /Total: 10 tests in 1 file/); + assert.doesNotMatch(executions, /prompt-manager.spec/); +}); + +test('execution config refuses live-host URLs and binds a managed loopback port', () => { + for (const url of ['http://127.0.0.1:8188/', 'https://127.0.0.1:8196/', 'http://localhost:8196/', 'http://example.com:8196/']) { + assert.throws(() => execFileSync(process.execPath, ['--input-type=module', '-e', + "await import('./tools/vendor/playwright.execution.config.mjs')"], { cwd: root, stdio: 'pipe', + env: { ...process.env, PM_COMFY_ROOT: '/unused', PM_EXECUTION_URL: url } }), /loopback|8196/); + } +}); + +test('CI retains stock, execution and pack suites in order with uploaded execution reports', () => { + const workflow = readFileSync(new URL('../../.github/workflows/test.yml', import.meta.url), 'utf8'); + const commands = [...workflow.matchAll(/^ run: npm run --prefix tools\/vendor (test:[a-z]+)$/gm)].map(match => match[1]); + assert.deepEqual(commands, ['test:comfyui', 'test:executions', 'test:compat']); + assert.match(workflow, /path: tools\/vendor\/test-results\//); +}); diff --git a/tests/vendor/comfyui-host.test.mjs b/tests/vendor/comfyui-host.test.mjs index fb5c9cd..8bbfa19 100644 --- a/tests/vendor/comfyui-host.test.mjs +++ b/tests/vendor/comfyui-host.test.mjs @@ -1,10 +1,11 @@ import test from 'node:test'; import assert from 'node:assert/strict'; import { execFileSync } from 'node:child_process'; -import { mkdtemp, writeFile, mkdir, readFile, realpath, rm, access } from 'node:fs/promises'; +import { mkdtemp, writeFile, mkdir, readFile, realpath, rm, access, readdir } from 'node:fs/promises'; import { tmpdir } from 'node:os'; import { join } from 'node:path'; import { once } from 'node:events'; +import { createServer } from 'node:net'; import { launchComfyHost } from '../../tools/vendor/comfyui-host.mjs'; async function fixture(t, program) { @@ -59,3 +60,71 @@ test('PromptManager adapter isolates database, configuration and monitored direc await access(join(options.repository, '__init__.py')); }); + +test('execution host adds only the tracked fixture pack and removes all owned state', async t => { + const { launchExecutionHost } = await import('../comfyui/execution-host.mjs'); + const options = await fixture(t, idle); + const execution = await launchExecutionHost({ ...options, port: 8196 }); + t.after(() => execution.dispose()); + await once(execution.host.child.stdout, 'data'); + const args = execution.host.child.spawnargs; + assert.deepEqual(args.slice(args.indexOf('--whitelist-custom-nodes') + 1), ['promptmanager', 'pm_execution_fixture']); + assert.equal(execution.databasePath, join(execution.host.stateDirectory, 'prompts.db')); + const installed = join(execution.host.stateDirectory, 'custom_nodes/pm_execution_fixture'); + assert.match(await readFile(join(installed, '__init__.py'), 'utf8'), /PMHistoryCached/); + assert.deepEqual(await readdir(installed), ['__init__.py']); + const source = execution.fixtureDirectory; + assert.notEqual(await realpath(source), await realpath(options.repository)); + await execution.dispose(); + await assert.rejects(access(execution.host.stateDirectory), { code: 'ENOENT' }); + await assert.rejects(access(source), { code: 'ENOENT' }); + await access(join(options.repository, '__init__.py')); +}); + +test('read-only ledger helper distinguishes absent schema from missing, corrupt and oversized data', async t => { + const { readExecutionRuns } = await import('../comfyui/execution-fixture.mjs'); + const directory = await mkdtemp(join(tmpdir(), 'pm-ledger-read-')); + t.after(() => rm(directory, { recursive: true, force: true })); + const database = join(directory, 'prompts.db'); + const python = process.env.PM_TEST_PYTHON || 'python'; + const job = '10000000-0000-4000-8000-000000000001'; + const session = '20000000-0000-4000-8000-000000000001'; + execFileSync(python, ['-c', 'import sqlite3,sys; sqlite3.connect(sys.argv[1]).close()', database]); + assert.deepEqual(await readExecutionRuns(database, job), { schemaAvailable: false, rows: [] }); + execFileSync(python, ['-c', `import sqlite3,sys +c=sqlite3.connect(sys.argv[1]) +c.execute('CREATE TABLE execution_runs (run_id TEXT, job_id TEXT, session_id TEXT, outcome TEXT)') +c.executemany('INSERT INTO execution_runs VALUES (?,?,?,?)', [('first',sys.argv[2],sys.argv[3],'succeeded'),('other','30000000-0000-4000-8000-000000000001',sys.argv[3],'failed')]) +c.commit()`, database, job, session]); + const before = await readFile(database); + assert.deepEqual(await readExecutionRuns(database, job, session), { schemaAvailable: true, rows: [ + { run_id: 'first', job_id: job, session_id: session, outcome: 'succeeded' }, + ] }); + assert.deepEqual((await readExecutionRuns(database, job, '40000000-0000-4000-8000-000000000001')).rows, []); + assert.deepEqual(await readFile(database), before); + await assert.rejects(readExecutionRuns(database, "' OR 1=1 --")); + execFileSync(python, ['-c', `import sqlite3,sys +c=sqlite3.connect(sys.argv[1]) +c.executemany('INSERT INTO execution_runs VALUES (?,?,?,?)', [(str(i),sys.argv[2],sys.argv[3],'succeeded') for i in range(256)]) +c.commit()`, database, job, session]); + await assert.rejects(readExecutionRuns(database, job), /bound|256/); + await assert.rejects(readExecutionRuns(join(directory, 'missing.db'), job)); + await writeFile(database, 'corrupt database'); + await assert.rejects(readExecutionRuns(database, job)); +}); + + +test('execution launcher refuses an occupied port without touching an existing server', async t => { + const { launchExecutionHost } = await import('../comfyui/execution-host.mjs'); + const options = await fixture(t, idle); + const server = createServer(); + await new Promise((resolve, reject) => { server.once('error', reject); server.listen(8196, '127.0.0.1', resolve); }); + let execution; + try { + await assert.rejects(async () => { execution = await launchExecutionHost({ ...options, port: 8196 }); }, /already in use/); + assert.equal(server.listening, true); + } finally { + await execution?.dispose(); + await new Promise(resolve => server.close(resolve)); + } +}); diff --git a/tools/vendor/comfyui-host.mjs b/tools/vendor/comfyui-host.mjs index 33b5168..a545487 100644 --- a/tools/vendor/comfyui-host.mjs +++ b/tools/vendor/comfyui-host.mjs @@ -4,6 +4,8 @@ import { fileURLToPath, pathToFileURL } from 'node:url'; import { launchComfyHost as launchSharedHost } from '@comfyassets/toolkit/testing/host'; import runtime from './comfyui-runtime.json' with { type: 'json' }; +export { launchSharedHost }; + const repositoryRoot = fileURLToPath(new URL('../../', import.meta.url)); export async function preparePromptManagerPlugin({ installed, stateDirectory }) { diff --git a/tools/vendor/comfyui.config.mjs b/tools/vendor/comfyui.config.mjs index c1a0c1a..1d2f206 100644 --- a/tools/vendor/comfyui.config.mjs +++ b/tools/vendor/comfyui.config.mjs @@ -14,6 +14,7 @@ const servers = [{ export default defineConfig({ testDir: '../../tests/comfyui', + testIgnore: 'execution-history.spec.mjs', outputDir: './test-results/comfyui', timeout: 60000, workers: 1, diff --git a/tools/vendor/package.json b/tools/vendor/package.json index 0a4c465..29aa97f 100644 --- a/tools/vendor/package.json +++ b/tools/vendor/package.json @@ -12,6 +12,7 @@ "test:browser": "playwright test --config playwright.config.mjs", "test:workflows": "playwright test --config playwright.workflows.config.mjs", "test:comfyui": "playwright test --config comfyui.config.mjs", + "test:executions": "playwright test --config playwright.execution.config.mjs", "test:coverage": "node frontend-coverage.mjs", "coverage:update": "node frontend-coverage.mjs --update-baseline", "test:compat": "playwright test --config pack-compat.config.mjs" diff --git a/tools/vendor/playwright.execution.config.mjs b/tools/vendor/playwright.execution.config.mjs new file mode 100644 index 0000000..c57349b --- /dev/null +++ b/tools/vendor/playwright.execution.config.mjs @@ -0,0 +1,21 @@ +import { defineConfig } from './node_modules/@playwright/test/index.mjs'; + +if (!process.env.PM_COMFY_ROOT) throw new Error('Set PM_COMFY_ROOT to a pinned disposable ComfyUI checkout'); +const host = new URL(process.env.PM_EXECUTION_URL || 'http://127.0.0.1:8196/'); +if (host.href !== 'http://127.0.0.1:8196/') throw new Error('Execution tests require managed loopback http://127.0.0.1:8196/'); + +export default defineConfig({ + testDir: '../../tests/comfyui', + testMatch: 'execution-history.spec.mjs', + outputDir: './test-results/executions/cases', + timeout: 60000, + workers: 1, + forbidOnly: Boolean(process.env.CI), + retries: 0, + use: { baseURL: host.href, browserName: 'chromium', viewport: { width: 1600, height: 1100 }, + trace: 'retain-on-failure', screenshot: 'only-on-failure' }, + projects: [ + { name: 'classic-start', metadata: { initialRenderer: false } }, + { name: 'nodes2-start', metadata: { initialRenderer: true } }, + ], +}); diff --git a/utils/execution_history.py b/utils/execution_history.py new file mode 100644 index 0000000..c05c07b --- /dev/null +++ b/utils/execution_history.py @@ -0,0 +1,145 @@ +"""Own bounded queue-attempt capture without retaining host workflow payloads.""" + +import atexit +import threading +import time +from typing import Callable +from uuid import uuid4 + +try: + from ..database import PromptDatabase +except ImportError: + from database import PromptDatabase + +from .execution_history_events import make_finish, make_start +from .execution_history_writer import HistoryWriter +from .logging_config import get_logger +from .queue_observers import subscribe_dequeue, subscribe_task_done + +_logger = get_logger("prompt_manager.execution_history") +_STATE_ATTRIBUTE = "_prompt_manager_execution_history" +_registration_lock = threading.Lock() +_MAX_ACTIVE = 128 + + +class HistoryRegistration: + """Own subscriptions, one asynchronous writer and bounded attempt identities.""" + + def __init__( + self, queue: object, session_id: str, db_factory: Callable[[], PromptDatabase] + ) -> None: + self.writer = HistoryWriter(db_factory) + self._session_id = session_id + self._lock = threading.Lock() + self._active: dict[int, str] = {} + self._disposed = False + self._unsubscribers: list[Callable[[], None]] = [] + try: + self._unsubscribers.append( + subscribe_dequeue(queue, "execution_history", self._dequeue) + ) + self._unsubscribers.append( + subscribe_task_done(queue, "execution_history", self._task_done) + ) + self.writer.start() + atexit.register(self.dispose) + except Exception: + self.dispose() + raise + + def _ready_epoch(self) -> str | None: + status = self.writer.status() + epoch = status.epoch if status.accepting else None + if any(value != epoch for value in self._active.values()): + self._active.clear() + return epoch + + def _dequeue(self, result: object) -> None: + try: + with self._lock: + if self._disposed: + return + epoch = self._ready_epoch() + if epoch is None: + self.writer.record_gap() + return + command = make_start( + self._session_id, result, epoch, time.time_ns() // 1_000_000 + ) + if command is None: + self.writer.record_gap() + return + item_id = command.queue_item_id + if item_id in self._active: + return + if len(self._active) >= _MAX_ACTIVE: + self.writer.record_gap() + return + if self.writer.submit(command): + self._active[item_id] = epoch + except Exception: + self.writer.record_gap() + + def _task_done(self, args: tuple, kwargs: dict) -> None: + try: + with self._lock: + if self._disposed: + return + epoch = self._ready_epoch() + if epoch is None: + self.writer.record_gap() + return + item_id = args[0] if args else kwargs.get("item_id") + status = args[2] if len(args) > 2 else kwargs.get("status") + command = make_finish( + self._session_id, + item_id, + epoch, + status, + time.time_ns() // 1_000_000, + ) + if command is None: + self.writer.record_gap() + return + self._active.pop(command.queue_item_id, None) + self.writer.submit(command) + except Exception: + self.writer.record_gap() + + def dispose(self, timeout: float = 2.0) -> bool: + """Stop capture, release subscriptions and drain the owned writer boundedly.""" + with self._lock: + if not self._disposed: + self._disposed = True + self._active.clear() + for unsubscribe in self._unsubscribers: + unsubscribe() + self._unsubscribers.clear() + atexit.unregister(self.dispose) + return self.writer.close(timeout) + + +def register_history_hook( + server: object, db_factory: Callable[[], PromptDatabase] +) -> HistoryRegistration | None: + """Register only on queues supporting actual dequeue and final callbacks.""" + try: + queue = getattr(server, "prompt_queue", None) + if not all( + callable(getattr(queue, method, None)) for method in ("get", "task_done") + ): + return None + with _registration_lock: + state = getattr(queue, _STATE_ATTRIBUTE, None) + if state is None: + state = {"session_id": str(uuid4()), "registration": None} + setattr(queue, _STATE_ATTRIBUTE, state) + previous = state["registration"] + if previous is not None and not previous._disposed: + return previous + registration = HistoryRegistration(queue, state["session_id"], db_factory) + state["registration"] = registration + return registration + except Exception: + _logger.warning("Execution history queue capability unavailable") + return None diff --git a/utils/execution_history_events.py b/utils/execution_history_events.py new file mode 100644 index 0000000..91eb4c1 --- /dev/null +++ b/utils/execution_history_events.py @@ -0,0 +1,360 @@ +"""Pure, bounded execution-lifecycle records for the background ledger writer.""" + +from dataclasses import asdict, dataclass, replace +import json +from uuid import UUID, uuid5 + +MAX_COMMAND_BYTES = 2048 +MAX_MESSAGES = 1024 +MAX_INTEGER = 2**63 - 1 +TERMINAL_KINDS = { + "execution_success": ("succeeded", "terminal_success"), + "execution_error": ("failed", "terminal_error"), + "execution_interrupted": ("interrupted", "terminal_interrupted"), +} +EVENT_KINDS = frozenset((*TERMINAL_KINDS, "execution_start")) +MAX_EVENT_KIND_LENGTH = max(map(len, EVENT_KINDS)) + + +@dataclass(frozen=True) +class RunStart: + run_id: str + session_id: str + queue_item_id: int + job_id: str + storage_epoch: str + observed_dequeue_at: int + + +@dataclass(frozen=True) +class LifecycleEvent: + kind: str + job_id: str + timestamp: int | None + node_id: str | None + node_type: str | None + exception_type: str | None + + +@dataclass(frozen=True) +class TerminalEvidence: + events: tuple[LifecycleEvent, ...] + invalid: bool + limited: bool + diagnostics_omitted: bool + invalid_job_ids: tuple[str, ...] = () + + +@dataclass(frozen=True) +class RunFinish: + session_id: str + queue_item_id: int + storage_epoch: str + observed_callback_at: int + evidence: TerminalEvidence + + +@dataclass(frozen=True) +class Completion: + outcome: str + reason: str + executor_started_at: int | None + terminal_event_at: int | None + node_id: str | None + node_type: str | None + exception_type: str | None + diagnostics_omitted: bool + + +RunCommand = RunStart | RunFinish + + +def _valid_event_kind(value: object) -> bool: + return ( + isinstance(value, str) + and len(value) <= MAX_EVENT_KIND_LENGTH + and value in EVENT_KINDS + ) + + +def _valid_integer(value: object) -> bool: + return type(value) is int and 0 <= value <= MAX_INTEGER + + +def _valid_uuid(value: object) -> bool: + if not isinstance(value, str) or len(value) != 36: + return False + try: + return str(UUID(value)) == value + except ValueError: + return False + + +def _valid_identity(session_id: object, item_id: object, epoch: object) -> bool: + return _valid_uuid(session_id) and _valid_integer(item_id) and _valid_uuid(epoch) + + +def make_start( + session_id: str, queue_result: object, storage_epoch: str, observed_at: int +) -> RunStart | None: + """Extract only the public job ID and queue item ID from a dequeue result.""" + if not isinstance(queue_result, (tuple, list)) or len(queue_result) != 2: + return None + queue_item, item_id = queue_result + if not isinstance(queue_item, (tuple, list)) or len(queue_item) < 2: + return None + job_id = queue_item[1] + if not ( + _valid_identity(session_id, item_id, storage_epoch) + and _valid_uuid(job_id) + and _valid_integer(observed_at) + ): + return None + return RunStart( + str(uuid5(UUID(session_id), str(item_id))), + session_id, + item_id, + job_id, + storage_epoch, + observed_at, + ) + + +def _optional_diagnostic(value: object, limit: int) -> tuple[str | None, bool]: + if value is None: + return None, False + if isinstance(value, str) and len(value) <= limit: + # Lone surrogates cannot be encoded as UTF-8 allowed data. + try: + value.encode("utf-8") + except UnicodeEncodeError: + return None, True + return value, False + return None, True + + +def make_finish( + session_id: str, + queue_item_id: int, + storage_epoch: str, + status: object, + observed_at: int, +) -> RunFinish | None: + """Capture allowed lifecycle fields without copying the host envelope.""" + if not ( + _valid_identity(session_id, queue_item_id, storage_epoch) + and _valid_integer(observed_at) + ): + return None + messages = getattr(status, "messages", None) + events = {} + invalid_job_ids = {} + invalid = not isinstance(messages, (tuple, list)) + limited = False + omitted = False + if not invalid: + limited = len(messages) > MAX_MESSAGES + for index in range(min(len(messages), MAX_MESSAGES)): + message = messages[index] + usable_job_id = None + if ( + isinstance(message, (tuple, list)) + and len(message) >= 2 + and isinstance(message[1], dict) + and _valid_uuid(message[1].get("prompt_id")) + ): + usable_job_id = message[1]["prompt_id"] + malformed = not isinstance(message, (tuple, list)) or len(message) != 2 + if not malformed: + kind, data = message + malformed = not isinstance(kind, str) + if not malformed and not _valid_event_kind(kind): + continue + malformed |= usable_job_id is None + if not malformed: + timestamp = data.get("timestamp") + malformed = timestamp is not None and not _valid_integer(timestamp) + if malformed: + if usable_job_id is None: + invalid = True + else: + invalid_job_ids[usable_job_id] = None + continue + diagnostics = [] + for field, limit in ( + ("node_id", 128), + ("node_type", 128), + ("exception_type", 256), + ): + value, dropped = _optional_diagnostic(data.get(field), limit) + diagnostics.append(value) + omitted |= dropped + event = LifecycleEvent(kind, data["prompt_id"], timestamp, *diagnostics) + events[event] = None + command = RunFinish( + session_id, + queue_item_id, + storage_epoch, + observed_at, + TerminalEvidence( + tuple(events), invalid, limited, omitted, tuple(invalid_job_ids) + ), + ) + if len(_serialize(command)) <= MAX_COMMAND_BYTES: + return command + # Diagnostics yield to all candidate identities, kinds and event times. + stripped = tuple( + dict.fromkeys( + replace(event, node_id=None, node_type=None, exception_type=None) + for event in command.evidence.events + ) + ) + had_diagnostics = any( + event.node_id is not None + or event.node_type is not None + or event.exception_type is not None + for event in command.evidence.events + ) + command = replace( + command, + evidence=TerminalEvidence( + stripped, + invalid, + limited, + omitted or had_diagnostics, + tuple(invalid_job_ids), + ), + ) + if len(_serialize(command)) > MAX_COMMAND_BYTES: + # Required evidence could not fit: no partial terminal may prove success. + command = replace( + command, + evidence=replace( + command.evidence, events=(), limited=True, invalid_job_ids=() + ), + ) + return command + + +def classify_evidence(job_id: str, evidence: TerminalEvidence) -> Completion: + """Classify only candidates matching the job ID persisted at dequeue.""" + events = tuple(event for event in evidence.events if event.job_id == job_id) + start_times = [ + event.timestamp + for event in events + if event.kind == "execution_start" and event.timestamp is not None + ] + started_at = min(start_times, default=None) + terminals = tuple(event for event in events if event.kind in TERMINAL_KINDS) + kinds = {event.kind for event in terminals} + reason = None + if evidence.limited: + reason = "evidence_limit" + elif evidence.invalid or job_id in evidence.invalid_job_ids: + reason = "invalid_evidence" + elif len(kinds) > 1: + reason = "conflicting_terminal" + elif not kinds: + reason = "missing_terminal" + if reason is not None: + return Completion( + "unknown", + reason, + started_at, + None, + None, + None, + None, + evidence.diagnostics_omitted, + ) + outcome, reason = TERMINAL_KINDS[terminals[0].kind] + terminal_times = [ + event.timestamp for event in terminals if event.timestamp is not None + ] + return Completion( + outcome, + reason, + started_at, + max(terminal_times, default=None), + next((event.node_id for event in terminals if event.node_id is not None), None), + next( + (event.node_type for event in terminals if event.node_type is not None), + None, + ), + next( + ( + event.exception_type + for event in terminals + if event.exception_type is not None + ), + None, + ), + evidence.diagnostics_omitted, + ) + + +def _valid_evidence(evidence: object) -> bool: + if type(evidence) is not TerminalEvidence or not isinstance(evidence.events, tuple): + return False + if len(evidence.events) > MAX_MESSAGES or any( + type(flag) is not bool + for flag in (evidence.invalid, evidence.limited, evidence.diagnostics_omitted) + ): + return False + if ( + not isinstance(evidence.invalid_job_ids, tuple) + or len(evidence.invalid_job_ids) > MAX_MESSAGES + or not all(_valid_uuid(job_id) for job_id in evidence.invalid_job_ids) + ): + return False + for event in evidence.events: + if type(event) is not LifecycleEvent: + return False + if ( + not _valid_event_kind(event.kind) + or not _valid_uuid(event.job_id) + or (event.timestamp is not None and not _valid_integer(event.timestamp)) + ): + return False + for value, limit in ( + (event.node_id, 128), + (event.node_type, 128), + (event.exception_type, 256), + ): + if _optional_diagnostic(value, limit)[1]: + return False + return True + + +def _serialize(command: RunCommand) -> bytes: + return json.dumps( + asdict(command), ensure_ascii=False, allow_nan=False, separators=(",", ":") + ).encode("utf-8") + + +def encode_command(command: RunCommand) -> bytes: + """Validate and encode an immutable command within the transport budget.""" + if type(command) not in (RunStart, RunFinish): + raise ValueError("Invalid execution history command") + valid = _valid_identity( + command.session_id, command.queue_item_id, command.storage_epoch + ) + if isinstance(command, RunStart): + valid = ( + valid + and _valid_uuid(command.run_id) + and _valid_uuid(command.job_id) + and _valid_integer(command.observed_dequeue_at) + ) + else: + valid = ( + valid + and _valid_integer(command.observed_callback_at) + and _valid_evidence(command.evidence) + ) + if not valid: + raise ValueError("Invalid execution history command") + encoded = _serialize(command) + if len(encoded) > MAX_COMMAND_BYTES: + raise ValueError("Execution history command exceeds byte limit") + return encoded diff --git a/utils/execution_history_writer.py b/utils/execution_history_writer.py new file mode 100644 index 0000000..3cba705 --- /dev/null +++ b/utils/execution_history_writer.py @@ -0,0 +1,245 @@ +"""A bounded, single-worker transport for best-effort execution history.""" + +from dataclasses import dataclass +import math +import queue +import threading +import time +from typing import Callable + +try: + from ..database import PromptDatabase + from ..database.execution_history import ( + HistoryEpochChanged, + subscribe_history_restore, + ) +except ImportError: + from database import PromptDatabase + from database.execution_history import ( + HistoryEpochChanged, + subscribe_history_restore, + ) + +from .execution_history_events import RunCommand, RunStart, encode_command +from .logging_config import get_logger + +_logger = get_logger("prompt_manager.execution_history_writer") +_MAX_PENDING = 1024 +_PROBE_INTERVAL = 5.0 +_MAX_GAPS = 2**63 - 1 + + +@dataclass(frozen=True) +class HistoryWriterStatus: + epoch: str | None + accepting: bool + pending: int + recording_gaps: int + alive: bool + + +class HistoryWriter: + """Keep database initialization and all writes off producer threads.""" + + def __init__( + self, + db_factory: Callable[[], PromptDatabase], + *, + clock: Callable[[], float] = time.monotonic, + ) -> None: + self._db_factory = db_factory + self._clock = clock + self._lock = threading.Lock() + self._wake = threading.Event() + self._queue: queue.Queue[RunCommand] = queue.Queue(maxsize=_MAX_PENDING) + self._thread: threading.Thread | None = None + self._epoch: str | None = None + self._accepting = False + self._pending = 0 + self._gaps = 0 + self._generation = 0 + self._closing = False + self._close_deadline = 0.0 + self._next_probe = 0.0 + + def start(self) -> None: + with self._lock: + if self._thread is not None or self._closing: + return + self._thread = threading.Thread( + target=self._run, name="prompt-manager-history", daemon=True + ) + self._thread.start() + + def submit(self, command: RunCommand) -> bool: + try: + encode_command(command) + except ValueError: + with self._lock: + self._record_gap_locked() + self._wake.set() + return False + with self._lock: + if ( + not self._accepting + or command.storage_epoch != self._epoch + or self._pending >= _MAX_PENDING + ): + self._record_gap_locked() + accepted = False + else: + self._queue.put_nowait(command) + self._pending += 1 + accepted = True + self._wake.set() + return accepted + + def _record_gap_locked(self) -> None: + self._gaps = min(_MAX_GAPS, self._gaps + 1) + + def record_gap(self) -> None: + """Account for unavailable capture without constructing a fake command.""" + with self._lock: + self._record_gap_locked() + self._wake.set() + + def invalidate(self) -> None: + with self._lock: + self._accepting = False + self._epoch = None + self._generation += 1 + self._next_probe = self._clock() + _PROBE_INTERVAL + self._discard_locked() + self._wake.set() + + def status(self) -> HistoryWriterStatus: + with self._lock: + return HistoryWriterStatus( + self._epoch, + self._accepting, + self._pending, + self._gaps, + self._thread is not None and self._thread.is_alive(), + ) + + def close(self, timeout: float = 2.0) -> bool: + if not math.isfinite(timeout) or timeout < 0: + raise ValueError("Shutdown timeout must be finite and nonnegative") + with self._lock: + deadline = time.monotonic() + timeout + if not self._closing: + self._closing = True + self._close_deadline = deadline + self._accepting = False + thread = self._thread + self._wake.set() + if thread is not None: + thread.join(max(0.0, deadline - time.monotonic())) + return thread is None or not thread.is_alive() + + def _discard_locked(self) -> None: + while True: + try: + self._queue.get_nowait() + except queue.Empty: + return + self._pending -= 1 + self._record_gap_locked() + + def _run(self) -> None: + db: PromptDatabase | None = None + unsubscribe: Callable[[], None] | None = None + reported_gaps = 0 + next_diagnostic = 0.0 + try: + while True: + self._wake.clear() + with self._lock: + if self._closing and ( + self._pending == 0 or time.monotonic() >= self._close_deadline + ): + self._discard_locked() + break + generation = self._generation + probe = ( + self._epoch is None + and not self._closing + and self._clock() >= self._next_probe + ) + if probe: + self._next_probe = self._clock() + _PROBE_INTERVAL + command = None + if self._epoch is not None: + try: + command = self._queue.get_nowait() + except queue.Empty: + pass + if probe: + try: + if db is None: + db = self._db_factory() + with self._lock: + closing = self._closing + if closing: + continue + if unsubscribe is None: + unsubscribe = subscribe_history_restore( + db.model.db_path, self.invalidate + ) + db.model.get_connection().execute("PRAGMA busy_timeout=250") + epoch = db.get_execution_epoch() + with self._lock: + if not self._closing and generation == self._generation: + self._epoch = epoch + self._accepting = True + except Exception: + # The next scheduled probe opens a clean managed connection. + if db is not None: + db.model.close() + with self._lock: + self._record_gap_locked() + if command is not None: + failed = False + try: + # Restore may have reopened this thread's managed connection. + db.model.get_connection().execute("PRAGMA busy_timeout=250") + if isinstance(command, RunStart): + db.write_execution_start(command) + else: + db.write_execution_finish(command) + except HistoryEpochChanged: + # The exception's storage_epoch is only a hint; readiness + # comes from a later probe and never replays this command. + failed = True + except Exception: + failed = True + with self._lock: + self._pending -= 1 + if failed: + self._record_gap_locked() + self._accepting = False + self._epoch = None + self._generation += 1 + self._next_probe = self._clock() + _PROBE_INTERVAL + self._discard_locked() + with self._lock: + gaps = self._gaps + has_work = self._epoch is not None and not self._queue.empty() + now = self._clock() + if gaps != reported_gaps and now >= next_diagnostic: + _logger.warning( + "Execution history recording gaps: %s", gaps - reported_gaps + ) + reported_gaps = gaps + next_diagnostic = now + _PROBE_INTERVAL + if not has_work: + self._wake.wait(0.05) + finally: + with self._lock: + self._accepting = False + self._epoch = None + self._discard_locked() + if unsubscribe is not None: + unsubscribe() + if db is not None: + db.model.close() diff --git a/utils/queue_observers.py b/utils/queue_observers.py new file mode 100644 index 0000000..e07af2a --- /dev/null +++ b/utils/queue_observers.py @@ -0,0 +1,110 @@ +"""Keyed, independently removable observers for successful queue operations.""" + +import functools +import threading +import time +from typing import Callable + +from .logging_config import get_logger + +_logger = get_logger("prompt_manager.queue_observers") +_STATE_ATTRIBUTE = "_prompt_manager_queue_observers" +_registration_lock = threading.Lock() + + +class _Observation: + def __init__(self, queue: object, method: str) -> None: + self.queue = queue + self.method = method + self.lock = threading.Lock() + self.callbacks: dict[str, tuple[object, Callable]] = {} + self.original = None + self.wrapper = None + self.installed = False + self.next_warning = 0.0 + + def install(self) -> None: + original = getattr(self.queue, self.method, None) + if not callable(original): + raise ValueError("Queue capability unavailable") + + @functools.wraps(original) + def observe(*args: object, **kwargs: object) -> object: + result = original(*args, **kwargs) + if self.method == "task_done" or result is not None: + with self.lock: + callbacks = tuple(entry[1] for entry in self.callbacks.values()) + for callback in callbacks: + try: + if self.method == "get": + callback(result) + else: + callback(args, kwargs) + except Exception: + with self.lock: + now = time.monotonic() + warn = now >= self.next_warning + if warn: + self.next_warning = now + 5.0 + if warn: + _logger.warning("Queue observation callback failed") + return result + + setattr(self.queue, self.method, observe) + self.original, self.wrapper, self.installed = original, observe, True + + def subscribe(self, key: str, callback: Callable) -> Callable[[], None]: + with self.lock: + if not self.installed: + self.install() + existing = self.callbacks.get(key) + token = existing[0] if existing is not None else object() + if existing is None: + self.callbacks[key] = (token, callback) + + def unsubscribe() -> None: + with self.lock: + current = self.callbacks.get(key) + if current is None or current[0] is not token: + return + del self.callbacks[key] + if ( + not self.callbacks + and getattr(self.queue, self.method) is self.wrapper + ): + setattr(self.queue, self.method, self.original) + self.installed = False + + return unsubscribe + + +def _subscribe( + queue: object, method: str, key: str, callback: Callable +) -> Callable[[], None]: + with _registration_lock: + states = getattr(queue, _STATE_ATTRIBUTE, None) + if states is None: + states = {} + try: + setattr(queue, _STATE_ATTRIBUTE, states) + except (AttributeError, TypeError) as error: + raise ValueError("Queue does not support observation") from error + state = states.get(method) + if state is None: + state = _Observation(queue, method) + states[method] = state + return state.subscribe(key, callback) + + +def subscribe_dequeue( + queue: object, key: str, callback: Callable[[object], None] +) -> Callable[[], None]: + """Observe each actual get result once per key, after the host returns.""" + return _subscribe(queue, "get", key, callback) + + +def subscribe_task_done( + queue: object, key: str, callback: Callable[[tuple, dict], None] +) -> Callable[[], None]: + """Observe successful task_done calls with their original argument objects.""" + return _subscribe(queue, "task_done", key, callback) diff --git a/utils/usage_tracking.py b/utils/usage_tracking.py index 2acbdf2..4f0b094 100644 --- a/utils/usage_tracking.py +++ b/utils/usage_tracking.py @@ -19,7 +19,6 @@ Text from known pure string nodes is resolved at queue time too. Anything else (PromptSearchList batches, unknown nodes) is counted by the node as it executes. """ -import functools import threading import time import uuid @@ -32,6 +31,7 @@ try: from .prompt_graph import intrinsic_pair_nodes, resolve_text, run_prompt_components from .prompt_pair_resolver import resolve_text_input from .pair_execution import clear_pair_execution + from .queue_observers import subscribe_dequeue except ImportError: from utils.hashing import generate_prompt_hash from utils.logging_config import get_logger @@ -42,6 +42,7 @@ except ImportError: ) from utils.prompt_pair_resolver import resolve_text_input from utils.pair_execution import clear_pair_execution + from utils.queue_observers import subscribe_dequeue logger = get_logger("prompt_manager.usage_tracking") @@ -294,21 +295,20 @@ def _wrap_queue_get(prompt_queue: Any, on_dequeued: Callable[[str], None]) -> bo original = getattr(prompt_queue, "get", None) if not callable(original): return False - if getattr(original, _WRAPPED_MARKER, False): + if getattr(prompt_queue, _WRAPPED_MARKER, False) or getattr( + original, _WRAPPED_MARKER, False + ): return False - @functools.wraps(original) - def get(*args, **kwargs): - result = original(*args, **kwargs) - if result is not None: - try: - on_dequeued(str(result[0][1])) - except Exception as e: - logger.warning(f"Dequeue hook failed: {e}") - return result + def report(result: object) -> None: + try: + on_dequeued(str(result[0][1])) + except Exception as e: + logger.warning(f"Dequeue hook failed: {e}") - setattr(get, _WRAPPED_MARKER, True) - prompt_queue.get = get + subscribe_dequeue(prompt_queue, "usage_tracking", report) + setattr(prompt_queue.get, _WRAPPED_MARKER, True) + setattr(prompt_queue, _WRAPPED_MARKER, True) return True