Compare commits
10
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
1c552da7a4 | ||
|
|
c826b593b3 | ||
|
|
978fdfb2cc | ||
|
|
6abcdf5516 | ||
|
|
eb9800a898 | ||
|
|
246496d175 | ||
|
|
4f800c15aa | ||
|
|
c377064145 | ||
|
|
508f03d0a5 | ||
|
|
b731fdfe1b |
@@ -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
|
||||
|
||||
@@ -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
@@ -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
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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:
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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 }],
|
||||
});
|
||||
@@ -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'); }
|
||||
});
|
||||
@@ -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 }); }
|
||||
} };
|
||||
}
|
||||
@@ -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,
|
||||
}
|
||||
@@ -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)
|
||||
)
|
||||
)
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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()
|
||||
@@ -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)
|
||||
)
|
||||
)
|
||||
@@ -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)
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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
@@ -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
|
||||
)
|
||||
|
||||
@@ -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)
|
||||
@@ -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))
|
||||
|
||||
Vendored
+32
@@ -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\//);
|
||||
});
|
||||
|
||||
Vendored
+70
-1
@@ -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));
|
||||
}
|
||||
});
|
||||
|
||||
Vendored
+2
@@ -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 }) {
|
||||
|
||||
Vendored
+1
@@ -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,
|
||||
|
||||
Vendored
+1
@@ -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
@@ -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 } },
|
||||
],
|
||||
});
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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
@@ -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
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user