Compare commits

...
37 changed files with 4417 additions and 88 deletions
+5
View File
@@ -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
+28 -2
View File
@@ -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 `<extension>/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
`<extension>/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
+60 -2
View File
@@ -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.
+28
View File
@@ -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
+194
View File
@@ -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
+225
View File
@@ -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
+14
View File
@@ -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:
+6 -2
View File
@@ -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):
+2 -2
View File
@@ -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
+52 -36
View File
@@ -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.
+47
View File
@@ -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 }],
});
+168
View File
@@ -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'); }
});
+41
View File
@@ -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 }); }
} };
}
+71
View File
@@ -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,
}
+46
View File
@@ -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)
)
)
+5 -1
View File
@@ -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
+32
View File
@@ -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:
+452
View File
@@ -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()
+379
View File
@@ -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)
)
)
+450
View File
@@ -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)
+246
View File
@@ -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()
+486
View File
@@ -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()
+9 -12
View File
@@ -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__)
+223 -17
View File
@@ -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
)
+132
View File
@@ -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)
+16
View File
@@ -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))
+32
View File
@@ -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\//);
});
+70 -1
View File
@@ -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));
}
});
+2
View File
@@ -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 }) {
+1
View File
@@ -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,
+1
View File
@@ -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"
+21
View File
@@ -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 } },
],
});
+145
View File
@@ -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
+360
View File
@@ -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
+245
View File
@@ -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()
+110
View File
@@ -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)
+13 -13
View File
@@ -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