diff --git a/TESTING.md b/TESTING.md index 2fed74e..94c9296 100644 --- a/TESTING.md +++ b/TESTING.md @@ -200,10 +200,10 @@ 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`, +`npm run --prefix tools/vendor test:executions` runs 26 real execution cases +(13 each in Classic and Nodes 2.0, retaining the original ten ledger cases), +alongside 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 @@ -218,7 +218,11 @@ 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 +interruptible queue gate, deterministic CLIP/model/latent inputs and control routes. +It replaces external sampler inference and trigger injection only inside the +disposable host; real registered PromptManager classes, queue, executor and cache +remain active. Its cache control resets the real executor on its worker to exercise +Classic, LRU, RAM-pressure and disabled caching; 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. @@ -230,8 +234,10 @@ 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. +Sensitive diagnostic and extra-data sentinels must never appear in lifecycle +rows. Exact prompt fields are deliberately persisted only in context snapshots. +The same read-only helper calls the internal `get_execution_contexts` reader for +requested run IDs and inspects their snapshot identities and current usage counts. For a meaningful historical failure, create a separate disposable detached checkout at `fdba6fd0126a8a40cbd6f3e50ed4c5de8991efd4` and pass its path to @@ -253,5 +259,34 @@ 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. +runs. Schema v2 migrates automatically when the configured database opens, adding +immutable prompt-context tables without backfilling old runs. The internal reader +returns `not_captured` for preexisting runs; a new capture can instead be complete, +partial, ambiguous or unavailable. Keep the existing database backup/restore +procedure: restored history changes storage epoch, and old library navigation +hints must never be rebound to unrelated current rows. + +Context cases verify A,A,B as three jobs/two combinations/one positive snapshot +when every positive field matches; all five registered PM variants; exact raw +whitespace and Unicode; known empty negatives; Text negative roles and crossed +pairs; fully cached repeats after trigger configuration changes; successful source +text before downstream failure/interruption; unknown cached dynamic raw; unsupported +transforms and duplicate list invocation ambiguity. Fresh values use +`runtime_return`, cached effective STRING outputs use `host_cache_output`, and +submitted literals use graph inference. SQL NULL means unknown; empty text is +known. Text is never trimmed or normalized by capture. Cache availability and +inferred sampler roles prove neither sampler execution nor consumption, and list +indices do not identify generated images. External trigger configuration changes +must not rewrite the effective text of an existing cached result. The cache matrix +also checks that positive/pair usage increments once per dequeued job while a +negative-only component retains its existing zero standalone-use count. + +Fault injection for restore fencing, inherited async capability revocation, queued +finalization, epoch and generation revocation, provider-local cache entries, eviction/reset, malformed or +oversized payloads and missing host interfaces lives in the focused Python +`test_execution*.py` suites. These use disposable storage and mocked host interfaces; +they are distinct from the pinned real executor/cache checks above. No output-file +associations or chronological timeline are introduced. Deploying the Python changes +requires a ComfyUI restart. After restart, queue three new jobs A,A,B with identical +positive fields and two negative recipes; existing counters should advance as usual, +and internal context inspection should find three attempts and two combinations. diff --git a/database/execution_context_schema.py b/database/execution_context_schema.py new file mode 100644 index 0000000..be06283 --- /dev/null +++ b/database/execution_context_schema.py @@ -0,0 +1,406 @@ +"""Owned execution-context v1 tables within execution-history schema v2. + +Content IDs hash exact scalar values; evidence JSON contains metadata only. +All validation scans use cursors and bounded per-row payloads. +""" + +import json +import sqlite3 + +from .execution_history_schema import _compact, _uuid_check + +try: + from ..utils.execution_contexts import ( + BUNDLE_STATES, + CAPTURE_REASONS, + EVIDENCE_SOURCES, + UNKNOWN_REASONS, + MAX_CONTEXT_BYTES, + ContextBundle, + PromptSnapshot, + RecipeContext, + TextEvidence, + validate_bundle, + content_digest, + text_size, + ) +except ImportError: + from utils.execution_contexts import ( + BUNDLE_STATES, + CAPTURE_REASONS, + EVIDENCE_SOURCES, + UNKNOWN_REASONS, + MAX_CONTEXT_BYTES, + ContextBundle, + PromptSnapshot, + RecipeContext, + TextEvidence, + validate_bundle, + content_digest, + text_size, + ) + + +FIELDS = ("raw", "effective", "prepend", "append") +TABLE_COLUMNS = { + "execution_context_texts": ( + ("text_id", "TEXT", 1, 1), + ("text_value", "TEXT", 1, 0), + ), + "execution_prompt_snapshots": ( + ("snapshot_id", "TEXT", 1, 1), + *((f"{field}_text_id", "TEXT", 0, 0) for field in FIELDS), + ), + "execution_contexts": ( + ("context_id", "TEXT", 1, 1), + ("positive_snapshot_id", "TEXT", 0, 0), + ("negative_snapshot_id", "TEXT", 0, 0), + ), + "execution_run_contexts": ( + ("run_id", "TEXT", 1, 1), + ("anchor_id", "TEXT", 1, 2), + ("context_id", "TEXT", 1, 0), + ("evidence_json", "TEXT", 1, 0), + ("positive_prompt_id", "INTEGER", 0, 0), + ("negative_prompt_id", "INTEGER", 0, 0), + ("pair_id", "INTEGER", 0, 0), + ("reference_epoch", "TEXT", 0, 0), + ("capture_version", "INTEGER", 1, 0), + ), + "execution_context_capture": ( + ("run_id", "TEXT", 1, 1), + ("state", "TEXT", 1, 0), + ("reason", "TEXT", 1, 0), + ("format_version", "INTEGER", 1, 0), + ), +} + + +def _digest_check(column: str, nullable: bool = False) -> str: + check = ( + f"typeof({column}) = 'text' AND length({column}) = 64 " + f"AND length(CAST({column} AS BLOB)) = 64 " + f"AND {column} NOT GLOB '*[^0-9a-f]*'" + ) + return f"{column} IS NULL OR ({check})" if nullable else check + + +def _bounded_uuid_check(column: str) -> str: + return f"{_uuid_check(column)} AND length(CAST({column} AS BLOB)) = 36" + + +def _enum_check(column: str, values: frozenset[str]) -> str: + return ( + f"typeof({column}) = 'text' AND {column} IN (" + + ",".join("'" + value + "'" for value in sorted(values)) + + ")" + ) + + +TABLE_CHECKS = { + "execution_context_texts": ( + _digest_check("text_id"), + "typeof(text_value) = 'text' AND length(CAST(text_value AS BLOB)) <= 65536", + ), + "execution_prompt_snapshots": ( + _digest_check("snapshot_id"), + *(_digest_check(f"{field}_text_id", True) for field in FIELDS), + ), + "execution_contexts": ( + _digest_check("context_id"), + _digest_check("positive_snapshot_id", True), + _digest_check("negative_snapshot_id", True), + "positive_snapshot_id IS NOT NULL OR negative_snapshot_id IS NOT NULL", + ), + "execution_run_contexts": ( + _bounded_uuid_check("run_id"), + "typeof(anchor_id) = 'text' AND length(anchor_id) BETWEEN 1 AND 128 " + "AND length(CAST(anchor_id AS BLOB)) <= 512", + _digest_check("context_id"), + "typeof(evidence_json) = 'text' AND length(CAST(evidence_json AS " + "BLOB)) <= 8192", + *( + f"{column} IS NULL OR (typeof({column}) = 'integer' AND {column} > 0)" + for column in ("positive_prompt_id", "negative_prompt_id", "pair_id") + ), + f"reference_epoch IS NULL OR ({_bounded_uuid_check('reference_epoch')})", + "typeof(capture_version) = 'integer' AND capture_version = 1", + ), + "execution_context_capture": ( + _bounded_uuid_check("run_id"), + _enum_check("state", BUNDLE_STATES), + _enum_check("reason", CAPTURE_REASONS), + "length(CAST(reason AS BLOB)) <= 64", + "typeof(format_version) = 'integer' AND format_version = 1", + ), +} +# (source column, destination table, destination column); actions always NO ACTION. +TABLE_FOREIGN_KEYS = { + "execution_context_texts": (), + "execution_prompt_snapshots": tuple( + (f"{field}_text_id", "execution_context_texts", "text_id") for field in FIELDS + ), + "execution_contexts": tuple( + (f"{role}_snapshot_id", "execution_prompt_snapshots", "snapshot_id") + for role in ("positive", "negative") + ), + "execution_run_contexts": ( + ("run_id", "execution_runs", "run_id"), + ("context_id", "execution_contexts", "context_id"), + ), + "execution_context_capture": (("run_id", "execution_runs", "run_id"),), +} +# Only SQLite-generated primary-key indexes are unique. No secondary indexes required. + + +def install_context_schema(conn: sqlite3.Connection) -> None: + """Install tables inside the caller's transaction, without committing.""" + for table, columns in TABLE_COLUMNS.items(): + definitions = [ + f"{name} {kind}{' NOT NULL' if required else ''}" + for name, kind, required, _ in columns + ] + primary = [name for name, _, _, order in columns if order] + definitions.append(f"PRIMARY KEY ({', '.join(primary)})") + definitions.extend(f"CHECK({check})" for check in TABLE_CHECKS[table]) + definitions.extend( + f"FOREIGN KEY ({column}) REFERENCES {target}({key}) " + "ON UPDATE NO ACTION ON DELETE NO ACTION" + for column, target, key in TABLE_FOREIGN_KEYS[table] + ) + conn.execute(f"CREATE TABLE {table} ({', '.join(definitions)})") + + +def _unique_object(pairs: list[tuple[str, object]]) -> dict: + result = {} + for key, value in pairs: + if key in result: + raise ValueError("Duplicate evidence key") + result[key] = value + return result + + +def _snapshot_values(conn: sqlite3.Connection, snapshot_id: str) -> tuple: + row = conn.execute( + "SELECT raw_text_id,effective_text_id,prepend_text_id,append_text_id " + "FROM execution_prompt_snapshots WHERE snapshot_id=?", + (snapshot_id,), + ).fetchone() + if row is None: + raise ValueError("Missing snapshot") + values = [] + for text_id in row: + if text_id is None: + values.append(None) + else: + text = conn.execute( + "SELECT text_value FROM execution_context_texts WHERE text_id=?", + (text_id,), + ).fetchone() + if text is None: + raise ValueError("Missing context text") + text_size(text[0]) + values.append(text[0]) + return tuple(values) + + +def _validate_evidence( + conn: sqlite3.Connection, anchor_id: str, context_id: str, encoded: str +) -> tuple[RecipeContext, bool]: + text_size(encoded, 8192) + evidence = json.loads(encoded, object_pairs_hook=_unique_object) + if ( + type(evidence) is not dict + or set(evidence) + != {"format_version", "role_provenance", "positive", "negative"} + or type(evidence["format_version"]) is not int + or evidence["format_version"] != 1 + or evidence["role_provenance"] != "inferred" + ): + raise ValueError("Invalid context evidence") + row = conn.execute( + "SELECT positive_snapshot_id,negative_snapshot_id FROM " + "execution_contexts WHERE context_id=?", + (context_id,), + ).fetchone() + if row is None: + raise ValueError("Missing context") + complete = True + snapshots = [] + for role, snapshot_id in zip(("positive", "negative"), row): + metadata = evidence[role] + if snapshot_id is None: + if metadata is not None: + raise ValueError("Unexpected role evidence") + complete = False + snapshots.append(None) + continue + values = _snapshot_values(conn, snapshot_id) + if type(metadata) is not dict or set(metadata) != set(FIELDS): + raise ValueError("Invalid snapshot evidence") + fields = [] + for field, value in zip(FIELDS, values): + item = metadata[field] + if type(item) is not dict or set(item) != { + "provenance", + "source", + "reason", + }: + raise ValueError("Invalid field evidence") + provenance, source, reason = ( + item["provenance"], + item["source"], + item["reason"], + ) + if ( + type(provenance) is not str + or type(source) is not str + or source not in EVIDENCE_SOURCES.get(provenance, ()) + ): + raise ValueError("Invalid evidence source") + if value is None: + complete = False + if ( + provenance != "unknown" + or type(reason) is not str + or reason not in UNKNOWN_REASONS + ): + raise ValueError("Unknown evidence requires reason") + elif provenance == "unknown" or reason is not None: + raise ValueError("Known evidence metadata mismatch") + fields.append(TextEvidence(value, provenance, source, reason)) + snapshots.append(PromptSnapshot(*fields)) + return RecipeContext(anchor_id, *snapshots, evidence["role_provenance"]), complete + + +def context_schema_error(conn: sqlite3.Connection) -> str | None: + """Validate exact owned schema, bounded evidence, fingerprints and links.""" + if conn.execute( + "SELECT 1 FROM sqlite_master WHERE type IN ('trigger','view') LIMIT 1" + ).fetchone(): + return "Unsupported execution context trigger/view" + 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 tables.keys() != TABLE_COLUMNS.keys(): + return "Incomplete execution context 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_xinfo({table})") + ) + if columns != expected: + return f"{table} has invalid execution context 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 context checks" + foreign = tuple( + sorted( + (row[3], row[2], row[4], row[5], row[6], row[7]) + for row in conn.execute(f"PRAGMA foreign_key_list({table})") + ) + ) + expected_foreign = tuple( + sorted( + (*key, "NO ACTION", "NO ACTION", "NONE") + for key in TABLE_FOREIGN_KEYS[table] + ) + ) + if foreign != expected_foreign: + return f"{table} has invalid execution context foreign keys" + for index in conn.execute(f"PRAGMA index_list({table})"): + if not index[2]: + 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 context 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 context records" + if conn.execute(f"PRAGMA foreign_key_check({table})").fetchone(): + return f"{table} contains orphan execution context records" + if conn.execute( + "SELECT 1 FROM execution_run_contexts GROUP BY run_id HAVING " + "COUNT(*) > 16 LIMIT 1" + ).fetchone(): + return "Too many execution contexts per run" + try: + for text_id, text in conn.execute( + "SELECT text_id,text_value FROM execution_context_texts" + ): + text_size(text) + if content_digest("text", (text,)) != text_id: + raise ValueError("Invalid text fingerprint") + for (snapshot_id,) in conn.execute( + "SELECT snapshot_id FROM execution_prompt_snapshots" + ): + if ( + content_digest("snapshot", _snapshot_values(conn, snapshot_id)) + != snapshot_id + ): + raise ValueError("Invalid snapshot fingerprint") + for context_id, positive, negative in conn.execute( + "SELECT context_id,positive_snapshot_id,negative_snapshot_id FROM " + "execution_contexts" + ): + if content_digest("context", (positive, negative)) != context_id: + raise ValueError("Invalid context fingerprint") + payload_size = 0 + for snapshot_id in (positive, negative): + if snapshot_id is not None: + payload_size += sum( + text_size(value) + for value in _snapshot_values(conn, snapshot_id) + if value is not None + ) + if payload_size > MAX_CONTEXT_BYTES: + raise ValueError("Oversized detached context") + previous_run = None + contexts = () + for run_id, anchor_id, context_id, evidence in conn.execute( + "SELECT run_id,anchor_id,context_id,evidence_json FROM " + "execution_run_contexts ORDER BY run_id,anchor_id" + ): + if run_id != previous_run: + contexts = () + previous_run = run_id + context, complete = _validate_evidence( + conn, anchor_id, context_id, evidence + ) + status = conn.execute( + "SELECT state,reason FROM execution_context_capture WHERE run_id=?", + (run_id,), + ).fetchone() + if status is None: + raise ValueError("Context link has no capture header") + state, reason = status + if state == "complete" and not complete: + raise ValueError("Incomplete evidence marked complete") + candidate = (*contexts, context) + validate_bundle(ContextBundle(candidate, state, reason)) + contexts = candidate + if conn.execute( + "SELECT 1 FROM execution_context_capture c WHERE state='complete' " + "AND NOT EXISTS (SELECT 1 FROM execution_run_contexts r WHERE " + "r.run_id=c.run_id) LIMIT 1" + ).fetchone(): + raise ValueError("Complete capture has no contexts") + except (ValueError, TypeError, RecursionError, OverflowError) as exc: + return f"Invalid execution context content: {type(exc).__name__}" + return None diff --git a/database/execution_contexts.py b/database/execution_contexts.py new file mode 100644 index 0000000..4771f04 --- /dev/null +++ b/database/execution_contexts.py @@ -0,0 +1,376 @@ +"""Immutable execution capture and exact, epoch-fenced library navigation hints.""" + +from contextlib import ExitStack +import json +import sqlite3 +from typing import Callable, ContextManager, Iterator + +try: + from ..utils.execution_history_events import RunFinish, encode_command + from ..utils.execution_contexts import ( + CAPTURE_REASONS, + ContextBundle, + content_digest, + validate_bundle, + ) +except ImportError: + from utils.execution_history_events import RunFinish, encode_command + from utils.execution_contexts import ( + CAPTURE_REASONS, + ContextBundle, + content_digest, + validate_bundle, + ) + + +FIELDS = ("raw", "effective", "prepend", "append") +ROLES = ("positive", "negative") + + +def _prepare(bundle: ContextBundle) -> list[dict]: + validate_bundle(bundle) + if bundle.reason not in CAPTURE_REASONS: + raise ValueError("Invalid capture reason") + prepared = [] + complete = bool(bundle.contexts) + for context in bundle.contexts: + evidence = {"format_version": 1, "role_provenance": context.role_provenance} + snapshots = [] + for role in ROLES: + snapshot = getattr(context, role) + if snapshot is None: + snapshots.append(None) + evidence[role] = None + complete = False + continue + values = tuple(getattr(snapshot, field).value for field in FIELDS) + complete = complete and all(value is not None for value in values) + snapshots.append((content_digest("snapshot", values), values)) + evidence[role] = { + field: { + "provenance": getattr(snapshot, field).provenance, + "source": getattr(snapshot, field).source, + "reason": getattr(snapshot, field).reason, + } + for field in FIELDS + } + if all(snapshot is None for snapshot in snapshots): + raise ValueError("Context must contain a snapshot") + context_id = content_digest( + "context", tuple(s[0] if s else None for s in snapshots) + ) + encoded = json.dumps( + evidence, + ensure_ascii=False, + allow_nan=False, + separators=(",", ":"), + sort_keys=True, + ) + if len(encoded.encode("utf-8")) > 8192: + raise ValueError("Oversized context evidence") + prepared.append( + { + "anchor_id": context.anchor_id, + "context_id": context_id, + "snapshots": snapshots, + "evidence_json": encoded, + } + ) + if bundle.state == "complete" and not complete: + raise ValueError("Incomplete scalar capture") + return sorted(prepared, key=lambda item: item["anchor_id"]) + + +def _existing_row( + conn: sqlite3.Connection, table: str, columns: tuple[str, ...], values: tuple +) -> bool: + row = conn.execute( + f"SELECT {','.join(columns[1:])} FROM {table} WHERE {columns[0]}=?", + (values[0],), + ).fetchone() + if row is None: + return False + if tuple(row) != values[1:]: + raise ValueError("Execution context digest collision") + return True + + +def _content_rows(item: dict) -> Iterator[tuple[str, tuple[str, ...], tuple]]: + snapshot_ids = [] + for snapshot in item["snapshots"]: + if snapshot is None: + snapshot_ids.append(None) + continue + snapshot_id, values = snapshot + text_ids = [] + for value in values: + text_id = None if value is None else content_digest("text", (value,)) + if value is not None: + yield "execution_context_texts", ("text_id", "text_value"), ( + text_id, + value, + ) + text_ids.append(text_id) + yield "execution_prompt_snapshots", ( + "snapshot_id", + *(f"{field}_text_id" for field in FIELDS), + ), (snapshot_id, *text_ids) + snapshot_ids.append(snapshot_id) + yield "execution_contexts", ( + "context_id", + "positive_snapshot_id", + "negative_snapshot_id", + ), (item["context_id"], *snapshot_ids) + + +def _insert_exact( + conn: sqlite3.Connection, table: str, columns: tuple[str, ...], values: tuple +) -> None: + if not _existing_row(conn, table, columns, values): + names = ",".join(columns) + placeholders = ",".join("?" for _ in columns) + conn.execute(f"INSERT INTO {table} ({names}) VALUES ({placeholders})", values) + + +def _prompt_hint(conn: sqlite3.Connection, raw: str | None) -> int | None: + if raw is None: + return None + rows = conn.execute( + "SELECT id FROM prompts WHERE text=? COLLATE BINARY LIMIT 2", (raw,) + ).fetchall() + return rows[0][0] if len(rows) == 1 else None + + +def _hints( + conn: sqlite3.Connection, positive: str | None, negative: str | None +) -> tuple: + pos_id, neg_id = _prompt_hint(conn, positive), _prompt_hint(conn, negative) + pair_id = None + negative_ambiguous = ( + negative is not None + and neg_id is None + and conn.execute( + "SELECT 1 FROM prompts WHERE text=? COLLATE BINARY LIMIT 1", (negative,) + ).fetchone() + is not None + ) + if pos_id is not None and negative is not None and not negative_ambiguous: + rows = conn.execute( + "SELECT pp.id FROM prompt_pairs pp LEFT JOIN prompts n ON " + "n.id=pp.negative_prompt_id " + "WHERE pp.positive_prompt_id=? AND ((pp.negative_prompt_id IS NULL AND " + "?='') " + "OR (pp.negative_prompt_id=? AND n.text=? COLLATE BINARY)) LIMIT 2", + (pos_id, negative, neg_id, negative), + ).fetchall() + if len(rows) == 1: + pair_id = rows[0][0] + return pos_id, neg_id, pair_id + + +def _values(conn: sqlite3.Connection, snapshot_id: str | None) -> dict | None: + if snapshot_id is None: + return None + row = conn.execute( + "SELECT r.text_value,e.text_value,p.text_value,a.text_value " + "FROM execution_prompt_snapshots s " + "LEFT JOIN execution_context_texts r ON r.text_id=s.raw_text_id " + "LEFT JOIN execution_context_texts e ON e.text_id=s.effective_text_id " + "LEFT JOIN execution_context_texts p ON p.text_id=s.prepend_text_id " + "LEFT JOIN execution_context_texts a ON a.text_id=s.append_text_id " + "WHERE s.snapshot_id=?", + (snapshot_id,), + ).fetchone() + if row is None: + raise sqlite3.DatabaseError("Missing execution context snapshot") + return dict(zip(FIELDS, row)) + + +def _write_capture( + conn: sqlite3.Connection, + run_id: str, + expected_epoch: str, + bundle: ContextBundle, + prepared: list[dict], +) -> bool: + capture = conn.execute( + "SELECT state,reason,format_version FROM execution_context_capture " + "WHERE run_id=?", + (run_id,), + ).fetchone() + if capture is not None: + for item in prepared: + for table, columns, values in _content_rows(item): + _existing_row(conn, table, columns, values) + links = conn.execute( + "SELECT anchor_id,context_id,evidence_json FROM " + "execution_run_contexts WHERE run_id=? ORDER BY anchor_id", + (run_id,), + ).fetchall() + expected = [ + (item["anchor_id"], item["context_id"], item["evidence_json"]) + for item in prepared + ] + # Compare JSON objects so serialization order never creates a conflict. + original = [(row[0], row[1], json.loads(row[2])) for row in links] + proposed = [ + (anchor, context, json.loads(evidence)) + for anchor, context, evidence in expected + ] + if capture[0] != "conflicting" and ( + tuple(capture) != (bundle.state, bundle.reason, bundle.format_version) + or original != proposed + ): + conn.execute( + "UPDATE execution_context_capture SET " + "state='conflicting',reason='conflicting' WHERE run_id=?", + (run_id,), + ) + return True + for item in prepared: + for table, columns, values in _content_rows(item): + _insert_exact(conn, table, columns, values) + raw = [snapshot[1][0] if snapshot else None for snapshot in item["snapshots"]] + positive_id, negative_id, pair_id = _hints(conn, *raw) + conn.execute( + "INSERT INTO " + "execution_run_contexts(run_id,anchor_id,context_id,evidence_json," + "positive_prompt_id,negative_prompt_id,pair_id,reference_epoch," + "capture_version) VALUES(?,?,?,?,?,?,?,?,1)", + ( + run_id, + item["anchor_id"], + item["context_id"], + item["evidence_json"], + positive_id, + negative_id, + pair_id, + expected_epoch, + ), + ) + conn.execute( + "INSERT INTO " + "execution_context_capture(run_id,state,reason,format_version) " + "VALUES(?,?,?,?)", + (run_id, bundle.state, bundle.reason, bundle.format_version), + ) + return True + + +class ExecutionContextOperations: + """Compose with ExecutionHistoryOperations on PromptDatabase.""" + + def write_execution_contexts( + self, run_id: str, expected_epoch: str, bundle: ContextBundle + ) -> bool: + prepared = _prepare(bundle) + with self._execution_transaction(expected_epoch) as conn: + if ( + conn.execute( + "SELECT 1 FROM execution_runs WHERE run_id=? AND storage_epoch=?", + (run_id, expected_epoch), + ).fetchone() + is None + ): + return False + return _write_capture(conn, run_id, expected_epoch, bundle, prepared) + + def write_execution_contexts_for_finish( + self, + command: RunFinish, + bundle: ContextBundle, + *, + generation_guard: Callable[[], ContextManager[bool]] | None = None, + ) -> bool: + encode_command(command) + if type(command) is not RunFinish: + raise ValueError("Expected execution finish command") + prepared = _prepare(bundle) + # Enter the guard after BEGIN, and keep it until COMMIT/ROLLBACK exits. + with ExitStack() as guards: + with self._execution_transaction(command.storage_epoch) as conn: + if generation_guard is not None and not guards.enter_context( + generation_guard() + ): + return False + row = conn.execute( + "SELECT run_id 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 + return _write_capture( + conn, row[0], command.storage_epoch, bundle, prepared + ) + + def get_execution_contexts(self, run_id: str) -> dict: + """Return first-capture scalars; hints are None unless currently valid. + + reference_epoch remains historical evidence even when hints are invalidated. + Reads never rebind a first-capture reference or alter stored capture data. + """ + conn = self.model.get_connection() + own_transaction = not conn.in_transaction + if own_transaction: + conn.execute("BEGIN") + try: + capture = conn.execute( + "SELECT state,reason,format_version FROM execution_context_capture " + "WHERE run_id=?", + (run_id,), + ).fetchone() + if capture is None: + return { + "state": "not_captured", + "reason": "missing", + "format_version": None, + "contexts": [], + } + epoch = self.get_execution_epoch() + contexts = [] + rows = conn.execute( + "SELECT r.*,c.positive_snapshot_id,c.negative_snapshot_id FROM " + "execution_run_contexts r JOIN execution_contexts c ON " + "c.context_id=r.context_id WHERE r.run_id=? ORDER BY r.anchor_id", + (run_id,), + ) + for row in rows: + positive = _values(conn, row["positive_snapshot_id"]) + negative = _values(conn, row["negative_snapshot_id"]) + hints = (None, None, None) + if row["reference_epoch"] == epoch: + current = _hints( + conn, + positive["raw"] if positive else None, + negative["raw"] if negative else None, + ) + hints = tuple( + row[key] if row[key] == value else None + for key, value in zip( + ("positive_prompt_id", "negative_prompt_id", "pair_id"), + current, + ) + ) + contexts.append( + { + "anchor_id": row["anchor_id"], + "context_id": row["context_id"], + "positive": positive, + "negative": negative, + "evidence": json.loads(row["evidence_json"]), + "positive_prompt_id": hints[0], + "negative_prompt_id": hints[1], + "pair_id": hints[2], + "reference_epoch": row["reference_epoch"], + } + ) + return { + "state": capture[0], + "reason": capture[1], + "format_version": capture[2], + "contexts": contexts, + } + finally: + if own_transaction: + conn.rollback() diff --git a/database/execution_history_schema.py b/database/execution_history_schema.py index 562648f..81e3dc6 100644 --- a/database/execution_history_schema.py +++ b/database/execution_history_schema.py @@ -4,7 +4,7 @@ import re import sqlite3 from uuid import uuid4 -SCHEMA_VERSION = 1 +SCHEMA_VERSION = 2 def _uuid_check(column: str) -> str: @@ -24,7 +24,7 @@ def _integer_check(column: str, *, nullable: bool = False) -> str: _METADATA_CHECKS = ( "singleton_id = 1", - "schema_version = 1", + "schema_version = 2", _uuid_check("storage_epoch"), ) _RUN_CHECKS = ( @@ -44,13 +44,13 @@ _RUN_CHECKS = ( ) ), "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 = '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', " + "(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)", *( @@ -92,6 +92,30 @@ _TABLE_CHECKS = { "execution_history_metadata": _METADATA_CHECKS, "execution_runs": _RUN_CHECKS, } +# Row-only guards preserve the supported v1 DDL fingerprint while bounding reads. +_ROW_BYTE_CHECKS = { + "execution_history_metadata": ("length(CAST(storage_epoch AS BLOB)) = 36",), + "execution_runs": ( + *( + f"length(CAST({column} AS BLOB)) = 36" + for column in ( + "run_id", + "session_id", + "job_id", + "storage_epoch", + ) + ), + *( + f"{column} IS NULL OR length(CAST({column} AS BLOB)) <= {limit}" + for column, limit in ( + ("node_id", 512), + ("node_type", 512), + ("exception_type", 1024), + ) + ), + ), +} + _INDEXES = { "idx_execution_runs_identity": (True, ("session_id", "queue_item_id")), "idx_execution_runs_observed": (False, ("observed_dequeue_at", "run_id")), @@ -114,10 +138,29 @@ def history_schema_error( for row in conn.execute("SELECT name,sql FROM sqlite_master WHERE type='table'") if row[0] in _TABLE_COLUMNS } + from .execution_context_schema import TABLE_COLUMNS, context_schema_error + + context_present = any( + row[0] in TABLE_COLUMNS + for row in conn.execute("SELECT name FROM sqlite_master WHERE type='table'") + ) + if not tables and context_present: + return "Execution contexts require execution history" if not tables: return None if allow_missing else "Missing execution history schema" if tables.keys() != _TABLE_COLUMNS.keys(): return "Incomplete execution history schema" + metadata = conn.execute( + "SELECT CASE WHEN typeof(schema_version) = 'integer' " + "THEN schema_version ELSE NULL END FROM execution_history_metadata LIMIT 2" + ).fetchall() + if len(metadata) != 1 or len(metadata[0]) != 1 or metadata[0][0] not in (1, 2): + return "Unsupported execution history metadata/version" + version = metadata[0][0] + if version == 1 and context_present: + return "Unexpected execution contexts for history v1" + if not allow_missing and version != SCHEMA_VERSION: + return "Execution history migration required" for table, expected in _TABLE_COLUMNS.items(): columns = tuple( (row[1], row[2].upper(), row[3], row[5]) @@ -127,6 +170,11 @@ def history_schema_error( return f"{table} has invalid execution history columns" sql = _compact(tables[table]) checks = _TABLE_CHECKS[table] + if table == "execution_history_metadata" and version == 1: + checks = tuple( + check.replace("schema_version = 2", "schema_version = 1") + for check in checks + ) if sql.count("check(") != len(checks) or any( f"check({_compact(check)})" not in sql for check in checks ): @@ -150,16 +198,18 @@ def history_schema_error( 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) + predicate = " AND ".join( + f"({check})" for check in (*checks, *_ROW_BYTE_CHECKS[table]) + ) 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" + "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: + if len(metadata) != 1 or metadata[0][0] != 1 or metadata[0][1] != 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(): @@ -172,13 +222,24 @@ def history_schema_error( item[3] or item[4] != "BINARY" for item in keys ): return f"Invalid execution history index {name}" + if version == 2: + return context_schema_error(conn) return None +def install_context_schema(conn: sqlite3.Connection) -> None: + from .execution_context_schema import install_context_schema as install + + install(conn) + + def migrate_execution_history(conn: sqlite3.Connection) -> None: - """Install both tables and indexes atomically, without backfilling runs.""" + """Install versioned history and context tables atomically, without backfill.""" conn.execute("SAVEPOINT execution_history_schema") try: + error = history_schema_error(conn) + if error: + raise sqlite3.DatabaseError(error) present = conn.execute( "SELECT 1 FROM sqlite_master WHERE name IN " "('execution_runs','execution_history_metadata') LIMIT 1" @@ -201,6 +262,33 @@ def migrate_execution_history(conn: sqlite3.Connection) -> None: "INSERT INTO execution_history_metadata VALUES(1,?,?)", (SCHEMA_VERSION, str(uuid4())), ) + install_context_schema(conn) + else: + version, epoch = conn.execute( + "SELECT schema_version,storage_epoch FROM execution_history_metadata" + ).fetchone() + if version == 1: + conn.execute( + "ALTER TABLE execution_history_metadata RENAME TO " + "execution_history_metadata_v1" + ) + definitions = [ + f"{name} {kind}{' NOT NULL' if required else ''}" + f"{' PRIMARY KEY' if primary else ''}" + for name, kind, required, primary in _TABLE_COLUMNS[ + "execution_history_metadata" + ] + ] + definitions.extend(f"CHECK({check})" for check in _METADATA_CHECKS) + conn.execute( + "CREATE TABLE execution_history_metadata " + f"({', '.join(definitions)})" + ) + conn.execute( + "INSERT INTO execution_history_metadata VALUES(1,2,?)", (epoch,) + ) + conn.execute("DROP TABLE execution_history_metadata_v1") + install_context_schema(conn) error = history_schema_error(conn, allow_missing=False) if error: raise sqlite3.DatabaseError(error) @@ -216,8 +304,8 @@ def rotate_history_epoch(conn: sqlite3.Connection) -> str: epoch = str(uuid4()) with conn: cursor = conn.execute( - "UPDATE execution_history_metadata SET storage_epoch=? " - "WHERE singleton_id=1", + "UPDATE execution_history_metadata SET storage_epoch=? WHERE " + "singleton_id=1", (epoch,), ) if cursor.rowcount != 1: diff --git a/database/operations.py b/database/operations.py index f927196..92c1ead 100644 --- a/database/operations.py +++ b/database/operations.py @@ -12,6 +12,7 @@ from .models import PromptModel, normalize_image_path, utc_now_iso from .pairs import PairDependencyError, PairOperations from .library import LibraryOperations from .execution_history import ExecutionHistoryOperations +from .execution_contexts import ExecutionContextOperations # Import logging system try: @@ -152,7 +153,12 @@ def _coerce_file_info(metadata: Any) -> tuple: ) -class PromptDatabase(PairOperations, LibraryOperations, ExecutionHistoryOperations): +class PromptDatabase( + PairOperations, + LibraryOperations, + ExecutionHistoryOperations, + ExecutionContextOperations, +): """Database operations class for managing prompts.""" def __init__(self, db_path: Optional[str] = None): diff --git a/prompt_manager.py b/prompt_manager.py index b070e02..3f653e9 100644 --- a/prompt_manager.py +++ b/prompt_manager.py @@ -6,6 +6,11 @@ with persistent prompt storage and search capabilities. import time from typing import Any, Dict, Optional, Tuple +if __package__: + from .utils.execution_context_capture import record_runtime_context +else: + from utils.execution_context_capture import record_runtime_context + try: from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict except ImportError: @@ -256,7 +261,11 @@ class PromptManager(PromptManagerBase, ComfyNodeABC): ) self.logger.info(f"CLIP encoding completed, text: {repr(encoding_text)[:80]}") - return (conditioning, encoding_text) + result = (conditioning, encoding_text) + record_runtime_context( + type(self).__name__, text, prepend_text, append_text, encoding_text + ) + return result @classmethod def IS_CHANGED( diff --git a/prompt_manager_pair.py b/prompt_manager_pair.py index 75ae89b..dcef20d 100644 --- a/prompt_manager_pair.py +++ b/prompt_manager_pair.py @@ -5,6 +5,11 @@ import json import time from typing import Any, Dict, Optional, Tuple +if __package__: + from .utils.execution_context_capture import record_runtime_context +else: + from utils.execution_context_capture import record_runtime_context + try: from .prompt_manager import IO, InputTypeDict, PromptManagerFocused from .utils.hashing import generate_prompt_hash @@ -204,7 +209,19 @@ class PromptManagerPair(PromptManagerFocused): ) except Exception as error: self.logger.warning(f"Could not register paired metadata: {error}") - return positive_conditioning, negative_conditioning, positive, negative_text + result = positive_conditioning, negative_conditioning, positive, negative_text + record_runtime_context( + type(self).__name__, + text, + prepend_text, + append_text, + positive, + role="positive", + ) + record_runtime_context( + type(self).__name__, negative_text, "", "", negative_text, role="negative" + ) + return result @classmethod def IS_CHANGED( diff --git a/prompt_manager_text.py b/prompt_manager_text.py index 7ba3123..92f8133 100644 --- a/prompt_manager_text.py +++ b/prompt_manager_text.py @@ -6,6 +6,11 @@ without CLIP encoding, while maintaining all database and search features. import time from typing import Any, Dict, Optional, Tuple +if __package__: + from .utils.execution_context_capture import record_runtime_context +else: + from utils.execution_context_capture import record_runtime_context + try: from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict except ImportError: @@ -227,7 +232,11 @@ class PromptManagerText(PromptManagerBase, ComfyNodeABC): ) self.logger.debug(f"Text processing completed: {final_text[:100]}...") - return (final_text,) + result = (final_text,) + record_runtime_context( + type(self).__name__, text, prepend_text, append_text, final_text + ) + return result @classmethod def IS_CHANGED( diff --git a/tests/comfyui/execution-history.spec.mjs b/tests/comfyui/execution-history.spec.mjs index f39fe00..6cf84ac 100644 --- a/tests/comfyui/execution-history.spec.mjs +++ b/tests/comfyui/execution-history.spec.mjs @@ -166,3 +166,185 @@ test('rejected and removed-before-dequeue jobs create no persisted attempts', as } } finally { await control(request, 'release'); } }); + +function recipeGraph({ text = 'context positive', negative = '', kind = 'PromptManagerPair', prepend = '', append = '', dynamic = false } = {}) { + const graph = { + '10': { class_type: 'PMHistoryInputs', inputs: {} }, + '20': { class_type: kind, inputs: { text: dynamic ? ['11', 0] : text, + prepend_text: prepend, append_text: append, ...(kind === 'PromptManagerText' ? {} : { clip: ['10', 0] }), + ...(kind === 'PromptManagerPair' ? { negative_text: negative } : {}) } }, + '30': { class_type: 'CLIPTextEncode', inputs: { text: negative, clip: ['10', 0] } }, + '40': { class_type: 'KSampler', inputs: { model: ['10', 1], latent_image: ['10', 2], + positive: ['20', 0], negative: kind === 'PromptManagerPair' ? ['20', 1] : ['30', 0], + seed: 1, steps: 1, cfg: 1, sampler_name: 'euler', scheduler: 'normal', denoise: 1 } }, + '50': { class_type: 'PMHistoryOutput', inputs: { value: ['40', 0] } }, + }; + if (kind === 'PromptManagerText') { + graph['31'] = { class_type: 'CLIPTextEncode', inputs: { text: ['20', 0], clip: ['10', 0] } }; + graph['40'].inputs.positive = ['31', 0]; + } + if (dynamic) graph['11'] = { class_type: 'PMHistoryCached', inputs: { value: text } }; + return graph; +} + +async function capture(execution, id, outcome = 'succeeded') { + const [run] = await rows(execution, id, [outcome]); + let result; + await expect.poll(async () => { + result = await readExecutionRuns(execution.databasePath, id); + return result.captures[run.run_id].state; + }, { timeout: 10000, message: 'Background writer persisted context finalization' }).not.toBe('not_captured'); + return result.captures[run.run_id]; +} + +test('actual PM A,A,B jobs persist two combinations and one identical positive snapshot', async ({ execution }) => { + const captures = []; + const positiveSnapshotIds = []; + for (const negative of ['negative A', 'negative A', 'negative B']) { + const id = await submit(execution.request, recipeGraph({ negative })); + await history(execution.request, id); + captures.push(await capture(execution, id)); + const persisted = await readExecutionRuns(execution.databasePath, id); + positiveSnapshotIds.push(persisted.snapshotIds[persisted.rows[0].run_id][0].positive_snapshot_id); + } + expect(new Set(positiveSnapshotIds).size).toBe(1); + expect(captures.map(value => value.state)).toEqual(['complete', 'complete', 'complete']); + const contexts = captures.map(value => value.contexts[0]); + expect(new Set(contexts.map(value => value.context_id)).size).toBe(2); + expect(new Set(contexts.map(value => JSON.stringify(value.positive))).size).toBe(1); + expect(contexts[0].evidence.positive.effective.source).toBe('runtime_return'); + expect(contexts[1].evidence.positive.effective.source).toBe('host_cache_output'); + expect(contexts[1].evidence.positive.raw.source).toBe('graph_literal'); + expect(contexts[1].evidence.role_provenance).toBe('inferred'); + expect(contexts[2].negative.raw).toBe('negative B'); +}); + +test('fresh registered PM variants preserve exact raw whitespace and known empty negative', async ({ execution }) => { + for (const kind of ['PromptManager', 'PromptManagerFocused', 'PromptManagerSplit', 'PromptManagerPair', 'PromptManagerText']) { + const graph = recipeGraph({ kind, text: ` ${kind} e\u0301\n`, prepend: ' pre \t', append: '\n post ' }); + const id = await submit(execution.request, graph); + await history(execution.request, id); + const value = (await capture(execution, id)).contexts[0]; + expect(value.positive).toEqual({ raw: ` ${kind} e\u0301\n`, prepend: ' pre \t', append: '\n post ', effective: `pre ${kind} e\u0301\n post` }); + expect(value.negative.raw).toBe(''); + expect(value.negative.effective).toBe(''); + expect(value.evidence.positive.raw).toMatchObject({ provenance: 'recorded', source: 'runtime_return' }); + } +}); + +test('fully cached text uses old effective output after external trigger configuration changes', async ({ execution }) => { + const request = execution.request; + const configure = suffix => request.post('/pm_execution_fixture/injection', { data: { suffix } }); + try { + expect((await configure(' trigger-A')).ok()).toBe(true); + const graph = recipeGraph({ text: 'cached trigger recipe', negative: 'trigger negative' }); + const first = await submit(request, graph); + await history(request, first); + const original = (await capture(execution, first)).contexts[0]; + expect(original.positive.effective).toBe('cached trigger recipe trigger-A'); + expect((await configure(' trigger-B')).ok()).toBe(true); + const second = await submit(request, graph); + const item = await history(request, second); + expect(item.status.messages.find(([kind]) => kind === 'execution_cached')[1].nodes).toEqual(expect.arrayContaining(['20', '40', '50'])); + const repeated = (await capture(execution, second)).contexts[0]; + expect(repeated.context_id).toBe(original.context_id); + expect(repeated.positive.effective).toBe(original.positive.effective); + expect(repeated.evidence.positive.effective.source).toBe('host_cache_output'); + } finally { expect((await configure('')).ok()).toBe(true); } +}); + +test('crossed Pair and Text negative roles remain independent declared contexts', async ({ execution }) => { + const graph = recipeGraph({ text: 'pair one positive', negative: 'pair one negative' }); + graph['21'] = { class_type: 'PromptManagerPair', inputs: { text: 'pair two positive', negative_text: 'pair two negative', clip: ['10', 0] } }; + graph['22'] = { class_type: 'PromptManagerText', inputs: { text: 'text negative role' } }; + graph['32'] = { class_type: 'CLIPTextEncode', inputs: { text: ['22', 0], clip: ['10', 0] } }; + graph['40'].inputs.negative = ['21', 1]; + graph['41'] = structuredClone(graph['40']); + graph['41'].inputs.positive = ['21', 0]; + graph['41'].inputs.negative = ['32', 0]; + graph['51'] = { class_type: 'PMHistoryOutput', inputs: { value: ['41', 0] } }; + const id = await submit(execution.request, graph); + await history(execution.request, id); + const result = await capture(execution, id); + expect(result.contexts.map(value => [value.anchor_id, value.positive.raw, value.negative.raw])).toEqual([ + ['40', 'pair one positive', 'pair two negative'], ['41', 'pair two positive', 'text negative role'], + ]); + for (const value of result.contexts) expect(value.evidence.negative.effective.source).toBe('runtime_return'); +}); + +test('successful source capture survives downstream failure and real interruption', async ({ execution }) => { + const request = execution.request; + for (const mode of ['failure', 'gate']) { + const graph = recipeGraph({ text: `source before ${mode}` }); + graph['50'].inputs.mode = mode; + if (mode === 'gate') await control(request, 'hold'); + try { + const id = await submit(request, graph); + if (mode === 'gate') { + await entered(request); + expect((await request.post('/interrupt', { data: {} })).ok()).toBe(true); + } + await history(request, id, mode === 'gate' ? 'execution_interrupted' : 'execution_error'); + const result = await capture(execution, id, mode === 'gate' ? 'interrupted' : 'failed'); + expect(result.contexts[0].positive.effective).toBe(`source before ${mode}`); + expect(result.contexts[0].evidence.positive.effective.source).toBe('runtime_return'); + } finally { await control(request, 'release'); } + } +}); + +test('cached dynamic raw stays unknown while literal affixes and effective cache text survive', async ({ execution }) => { + const graph = recipeGraph({ dynamic: true, text: 'runtime dynamic recipe', prepend: ' pre ', append: ' post ' }); + const first = await submit(execution.request, graph); + await history(execution.request, first); + expect((await capture(execution, first)).contexts[0].positive.raw).toBe('runtime dynamic recipe'); + const second = await submit(execution.request, graph); + await history(execution.request, second); + const value = (await capture(execution, second)).contexts[0]; + expect(value.positive).toEqual({ raw: null, prepend: ' pre ', append: ' post ', effective: 'pre runtime dynamic recipe post' }); + expect(value.evidence.positive.raw.provenance).toBe('unknown'); + expect(value.evidence.positive.effective.source).toBe('host_cache_output'); +}); + +test('unsupported conditioning transforms and duplicate list invocations preserve uncertainty', async ({ execution }) => { + const transformed = recipeGraph({ text: 'unsupported transformed recipe' }); + transformed['35'] = { class_type: 'PMHistoryTransform', inputs: { conditioning: ['20', 0] } }; + transformed['40'].inputs.positive = ['35', 0]; + const unsupported = await submit(execution.request, transformed); + await history(execution.request, unsupported); + const unknown = await capture(execution, unsupported); + expect(unknown.state).toBe('partial'); + expect(unknown.contexts[0].positive.effective).toBeNull(); + const graph = recipeGraph({ text: 'duplicate runtime list' }); + graph['11'] = { class_type: 'PMHistoryList', inputs: { value: 'duplicate runtime list' } }; + graph['20'].inputs.text = ['11', 0]; + const id = await submit(execution.request, graph); + await history(execution.request, id); + const result = await capture(execution, id); + expect(result.state).toBe('ambiguous'); + expect(result.reason).toBe('multiple_samples'); + expect(result.contexts[0].positive.effective).toBeNull(); +}); + +test('real executor Classic LRU RAM and disabled caches preserve honest evidence and usage', async ({ execution }) => { + const request = execution.request; + for (const mode of ['CLASSIC', 'LRU', 'RAM_PRESSURE', 'NONE']) { + expect((await request.post('/pm_execution_fixture/cache', { data: { mode } })).ok()).toBe(true); + const graph = recipeGraph({ text: `cache mode ${mode}`, negative: `cache mode negative ${mode}` }); + const results = []; + for (let index = 0; index < 2; index++) { + const id = await submit(request, graph); + await history(request, id); + results.push(await capture(execution, id)); + const persisted = await readExecutionRuns(execution.databasePath, id); + const [run] = persisted.rows; + expect(persisted.usage[run.run_id]).toEqual([ + { role: 'positive', run_count: index + 1 }, + { role: 'negative', run_count: 0 }, + { role: 'pair', run_count: index + 1 }, + ]); + } + expect(results[0].state).toBe('complete'); + expect(results[1].contexts[0].context_id).toBe(results[0].contexts[0].context_id); + expect(results[1].contexts[0].evidence.positive.effective.source).toBe(mode === 'NONE' ? 'runtime_return' : 'host_cache_output'); + } +}); diff --git a/tests/comfyui/fixtures/ledger_nodes.py b/tests/comfyui/fixtures/ledger_nodes.py index 785fa21..e5d0588 100644 --- a/tests/comfyui/fixtures/ledger_nodes.py +++ b/tests/comfyui/fixtures/ledger_nodes.py @@ -3,6 +3,9 @@ import threading import time +import nodes +import execution + from aiohttp import web import comfy.model_management from server import PromptServer @@ -69,3 +72,129 @@ NODE_CLASS_MAPPINGS = { "PMHistoryFailure": PMHistoryFailure, "PMHistoryGate": PMHistoryGate, } + +# Only external inference is replaced; PM classes and host execution remain real. + + +class _Clip: + def tokenize(self, text: str) -> str: + return text + + def encode_from_tokens_scheduled(self, tokens: str) -> list: + return [[tokens, {}]] + + +class PMHistoryInputs: + @classmethod + def INPUT_TYPES(cls) -> dict: + return {"required": {}} + + RETURN_TYPES = ("CLIP", "MODEL", "LATENT") + FUNCTION = "execute" + CATEGORY = "testing/execution-history" + + def execute(self) -> tuple: + return (_Clip(), object(), {"samples": "owned deterministic latent"}) + + +class PMHistoryOutput: + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": {"value": ("LATENT",)}, + "optional": {"mode": (["normal", "failure", "gate"],)}, + } + + RETURN_TYPES = () + FUNCTION = "execute" + CATEGORY = "testing/execution-history" + OUTPUT_NODE = True + + def execute(self, value: dict, mode: str = "normal") -> tuple: + if mode == "failure": + raise ValueError("owned downstream failure") + if mode == "gate": + PMHistoryGate().execute("owned downstream gate") + return () + + +def _sample(self: object, **kwargs: object) -> tuple: + return (kwargs["latent_image"],) + + +nodes.KSampler.sample = _sample +NODE_CLASS_MAPPINGS.update( + {"PMHistoryInputs": PMHistoryInputs, "PMHistoryOutput": PMHistoryOutput} +) + + +class PMHistoryList(PMHistoryCached): + OUTPUT_IS_LIST = (True,) + + def execute(self, value: str) -> tuple: + return ([value, value],) + + +class PMHistoryTransform: + @classmethod + def INPUT_TYPES(cls) -> dict: + return {"required": {"conditioning": ("CONDITIONING",)}} + + RETURN_TYPES = ("CONDITIONING",) + FUNCTION = "execute" + CATEGORY = "testing/execution-history" + + def execute(self, conditioning: list) -> tuple: + return (conditioning,) + + +NODE_CLASS_MAPPINGS.update( + {"PMHistoryList": PMHistoryList, "PMHistoryTransform": PMHistoryTransform} +) + +_INJECTION_ORIGINALS = {} + + +@PromptServer.instance.routes.post("/pm_execution_fixture/injection") +async def injection(request: web.Request) -> web.Response: + suffix = (await request.json())["suffix"] + if not isinstance(suffix, str) or len(suffix) > 128: + raise web.HTTPBadRequest() + for name in ("PromptManager", "PromptManagerText"): + cls = nodes.NODE_CLASS_MAPPINGS[name] + if name not in _INJECTION_ORIGINALS: + _INJECTION_ORIGINALS[name] = cls._inject_lora_trigger_words + cls._inject_lora_trigger_words = ( + (lambda self, text: text + suffix) if suffix else _INJECTION_ORIGINALS[name] + ) + return web.json_response({"configured": True}) + + +# Apply a requested host cache configuration on its own worker before execution. +# Retain only the mode scalar; never retain the worker/executor/cache objects. + +_CACHE_MODE = None +_EXECUTE = execution.PromptExecutor.execute + + +def _execute_with_cache_mode(self: object, *args: object, **kwargs: object) -> object: + global _CACHE_MODE + if _CACHE_MODE is not None: + self.cache_type = execution.CacheType[_CACHE_MODE] + self.cache_args = {**self.cache_args, "lru": 100, "ram": 0} + self.reset() + _CACHE_MODE = None + return _EXECUTE(self, *args, **kwargs) + + +execution.PromptExecutor.execute = _execute_with_cache_mode + + +@PromptServer.instance.routes.post("/pm_execution_fixture/cache") +async def cache_mode(request: web.Request) -> web.Response: + global _CACHE_MODE + mode = (await request.json())["mode"] + if mode not in ("CLASSIC", "LRU", "RAM_PRESSURE", "NONE"): + raise web.HTTPBadRequest() + _CACHE_MODE = mode + return web.json_response({"configured": True}) diff --git a/tests/comfyui/read_execution_history.py b/tests/comfyui/read_execution_history.py index f2ac22a..9974bf2 100644 --- a/tests/comfyui/read_execution_history.py +++ b/tests/comfyui/read_execution_history.py @@ -4,8 +4,19 @@ import argparse import json from pathlib import Path import sqlite3 +import sys +from types import SimpleNamespace from uuid import UUID +sys.path.insert(0, str(Path(__file__).resolve().parents[2])) +from database.execution_contexts import ExecutionContextOperations # noqa: E402 +from database.execution_history import ExecutionHistoryOperations # noqa: E402 + + +class _ReadOnlyHistory(ExecutionContextOperations, ExecutionHistoryOperations): + def __init__(self, connection: sqlite3.Connection) -> None: + self.model = SimpleNamespace(get_connection=lambda: connection) + def read_runs(database: str, job_id: str, session_id: str | None = None) -> dict: job_id = str(UUID(job_id)) @@ -30,7 +41,45 @@ def read_runs(database: str, job_id: str, session_id: str | None = None) -> dict ).fetchall() if len(rows) > 256: raise ValueError("Execution query exceeds the 256-row bound") - return {"schemaAvailable": True, "rows": [dict(row) for row in rows]} + reader = _ReadOnlyHistory(connection) + captures = { + row["run_id"]: reader.get_execution_contexts(row["run_id"]) for row in rows + } + usage = {} + snapshot_ids = {} + for run_id, capture in captures.items(): + usage[run_id] = [] + snapshot_ids[run_id] = [] + for context in capture["contexts"]: + for role in ("positive", "negative"): + prompt_id = context[f"{role}_prompt_id"] + if prompt_id is not None: + count = connection.execute( + "SELECT run_count FROM prompts WHERE id=?", (prompt_id,) + ).fetchone()[0] + usage[run_id].append({"role": role, "run_count": count}) + pair_id = context["pair_id"] + if pair_id is not None: + count = connection.execute( + "SELECT run_count FROM prompt_pairs WHERE id=?", (pair_id,) + ).fetchone()[0] + usage[run_id].append({"role": "pair", "run_count": count}) + snapshot_ids[run_id].append( + dict( + connection.execute( + "SELECT positive_snapshot_id,negative_snapshot_id " + "FROM execution_contexts WHERE context_id=?", + (context["context_id"],), + ).fetchone() + ) + ) + return { + "schemaAvailable": True, + "rows": [dict(row) for row in rows], + "captures": captures, + "usage": usage, + "snapshotIds": snapshot_ids, + } if __name__ == "__main__": diff --git a/tests/test_execution_context_budget.py b/tests/test_execution_context_budget.py new file mode 100644 index 0000000..55f4449 --- /dev/null +++ b/tests/test_execution_context_budget.py @@ -0,0 +1,115 @@ +"""Ownership and accounting behavior for shared context payloads.""" + +import unittest + + +class ContextBudgetTests(unittest.TestCase): + def test_reservation_transfer_shrink_and_idempotent_release(self): + from utils.execution_context_budget import ContextBudget + + pool = ContextBudget(100) + lease = pool.reserve(100) + self.assertIsNotNone(lease) + self.assertIsNone(pool.reserve(1)) + owner = lease.transfer() + lease.release() + self.assertEqual(pool.charged_bytes, 100) + with self.assertRaises(ValueError): + lease.transfer() + with self.assertRaises(ValueError): + lease.shrink_to(1) + owner.shrink_to(60) + self.assertEqual(owner.byte_count, 60) + self.assertEqual(pool.charged_bytes, 60) + for invalid in (-1, 61, True, 1.0): + with self.assertRaises(ValueError): + owner.shrink_to(invalid) + owner.release() + owner.release() + self.assertEqual(pool.charged_bytes, 0) + + def test_default_pool_and_invalid_reservations(self): + from utils.execution_context_budget import ContextBudget + + pool = ContextBudget() + lease = pool.reserve(8388608) + self.assertIsNotNone(lease) + self.assertIsNone(pool.reserve(1)) + lease.release() + for value in (-1, True, 1.0): + with self.assertRaises(ValueError): + pool.reserve(value) + with self.assertRaises(ValueError): + ContextBudget(value) + + def test_only_pool_issued_leases_can_charge_or_release(self): + from utils.execution_context_budget import ContextBudget, PayloadLease + + pool = ContextBudget(100) + with self.assertRaises(ValueError): + PayloadLease(pool, 100) + self.assertEqual(pool.charged_bytes, 0) + + def test_measured_payload_field_context_and_attempt_boundaries(self): + from dataclasses import replace + from utils.execution_contexts import ( + ContextBundle, + RecipeContext, + PromptSnapshot, + TextEvidence, + bundle_payload_size, + validate_bundle, + ) + + def recipe(anchor, fourth): + values = ( + "x" * 65536, + "x" * 65536, + "x" * 65536, + "x" * fourth, + "", + "", + "", + "", + ) + fields = tuple( + TextEvidence(v, "recorded", "runtime_return", None) for v in values + ) + return RecipeContext( + anchor, + PromptSnapshot(*fields[:4]), + PromptSnapshot(*fields[4:]), + "inferred", + ) + + # 8 * (8 provenance + 14 source) + 8 role + 1 anchor = 185. + context = recipe("1", 65351) + bundle = ContextBundle((context,), "complete", "captured") + self.assertEqual(bundle_payload_size(bundle), 262160) + self.assertIsNone(validate_bundle(bundle)) + with self.assertRaises(ValueError): + bundle_payload_size(replace(bundle, contexts=(recipe("1", 65352),))) + field = replace(context.positive.raw, value="é" * 32769) + bad = replace(context, positive=replace(context.positive, raw=field)) + with self.assertRaises(ValueError): + bundle_payload_size(replace(bundle, contexts=(bad,))) + attempt = replace( + bundle, contexts=tuple(recipe(str(i), 65347) for i in range(4)) + ) + self.assertEqual(bundle_payload_size(attempt), 1048576) + with self.assertRaises(ValueError): + bundle_payload_size( + replace(attempt, contexts=(recipe("0", 65348), *attempt.contexts[1:])) + ) + small = replace(context, positive=None, negative=None) + sixteen = replace( + bundle, contexts=tuple(replace(small, anchor_id=str(i)) for i in range(16)) + ) + validate_bundle(sixteen) + with self.assertRaises(ValueError): + bundle_payload_size( + replace( + sixteen, + contexts=(*sixteen.contexts, replace(small, anchor_id="17")), + ) + ) diff --git a/tests/test_execution_context_capture.py b/tests/test_execution_context_capture.py new file mode 100644 index 0000000..db0214d --- /dev/null +++ b/tests/test_execution_context_capture.py @@ -0,0 +1,573 @@ +"""Attempt fencing and bounded runtime evidence, without a running host.""" + +from contextvars import copy_context +from dataclasses import replace +import gc +import sys +import threading +import types +import unittest +import weakref +from unittest.mock import patch + +from test_execution_context_graph import transformed_input_graph + +from utils.execution_context_budget import ContextBudget +from utils.execution_context_graph import SourceRef +from utils.execution_contexts import PromptSnapshot, TextEvidence +from utils.execution_history_events import RunStart +from utils.execution_history_writer import HistoryWriterStatus + +try: + from utils import execution_context_capture as capture +except ImportError: + capture = None + +SESSION = "11111111-1111-4111-8111-111111111111" +EPOCH = "22222222-2222-4222-8222-222222222222" +JOB = "33333333-3333-4333-8333-333333333333" + + +def graph(): + return { + "p": { + "class_type": "PromptManagerPair", + "inputs": {"text": "raw", "negative_text": "negative"}, + }, + "s": { + "class_type": "KSampler", + "inputs": {"positive": ["p", 0], "negative": ["p", 1]}, + }, + } + + +def snapshot(raw=" raw ", effective="injected", prepend=" pre ", append=" post "): + return PromptSnapshot( + *( + TextEvidence(v, "recorded", "runtime_return", None) + for v in (raw, effective, prepend, append) + ) + ) + + +class CaptureTests(unittest.TestCase): + def setUp(self): + self.assertIsNotNone(capture, "attempt capture API is missing") + self.pool = ContextBudget() + self.status = HistoryWriterStatus(EPOCH, True, 0, 0, True, 3) + self.manager = capture.ContextCapture(10, lambda: self.status, self.pool) + self.addCleanup(self.manager.revoke_all) + self.start = RunStart(SESSION, SESSION, 1, JOB, EPOCH, 0) + + def begin(self, item=1, g=None): + return self.manager.begin( + replace(self.start, queue_item_id=item), + 3, + graph() if g is None else g, + None, + ) + + def seal(self, cap): + result = self.manager.seal(cap) + if result.lease: + self.addCleanup(result.lease.release) + return result + + def test_runtime_keeps_raw_affix_whitespace_and_injected_effective_text(self): + cap = self.begin() + self.manager.observe(cap, SourceRef("p", "PromptManagerPair", 0), 0, snapshot()) + result = self.seal(cap) + side = result.bundle.contexts[0].positive + self.assertEqual( + [ + side.raw.value, + side.effective.value, + side.prepend.value, + side.append.value, + ], + [" raw ", "injected", " pre ", " post "], + ) + self.assertIsNone(result.bundle.contexts[0].negative.effective.value) + self.assertEqual(result.state, "partial") + self.assertLess(self.pool.charged_bytes, 1048576) + + def test_reused_uuid_late_async_and_same_epoch_generation_are_fenced(self): + a = self.begin() + inherited = copy_context() + b = self.begin(2) + inherited.run( + self.manager.observe, + a, + SourceRef("p", "PromptManagerPair", 0), + 0, + snapshot(), + ) + self.assertIsNone(self.seal(a).bundle) + self.status = replace(self.status, capture_generation=4) + self.manager.observe(b, SourceRef("p", "PromptManagerPair", 0), 0, snapshot()) + self.assertIsNone(self.seal(b).bundle) + self.assertEqual(self.pool.charged_bytes, 0) + + def test_unbound_thread_and_forged_metadata_cannot_record(self): + cap = self.begin() + host = types.ModuleType("comfy_execution.utils") + host.get_executing_context = lambda: types.SimpleNamespace( + prompt_id=JOB, node_id="p", list_index=0 + ) + with patch.dict(sys.modules, {"comfy_execution.utils": host}): + thread = threading.Thread( + target=capture.record_runtime_context, + args=("PromptManagerPair", "bad", "", "", "bad"), + kwargs={"role": "positive"}, + ) + thread.start() + thread.join(2) + self.manager.observe( + replace(cap, nonce="forged"), + SourceRef("p", "PromptManagerPair", 0), + 0, + snapshot(), + ) + host.get_executing_context = lambda: types.SimpleNamespace( + prompt_id="forged", node_id="p", list_index=0 + ) + capture.record_runtime_context( + "PromptManagerPair", "bad", "", "", "bad", role="positive" + ) + self.assertIsNone(self.seal(cap).bundle.contexts[0].positive.effective.value) + + def test_new_dequeue_revokes_unfinished_attempt(self): + a = self.begin() + b = self.begin(2) + self.assertEqual(capture.active_capture(), (self.manager, b)) + self.manager.revoke(a) + self.assertEqual(capture.active_capture(), (self.manager, b)) + self.manager.revoke_all() + self.assertIsNone(capture.active_capture()) + self.assertEqual(self.pool.charged_bytes, 0) + + def test_equal_distinct_invocations_are_ambiguous_and_replay_is_idempotent(self): + for count, state in ((1, "partial"), (64, "ambiguous"), (65, "ambiguous")): + cap = self.begin(count) + for i in range(count): + self.manager.observe( + cap, SourceRef("p", "PromptManagerPair", 0), i, snapshot() + ) + self.manager.observe( + cap, SourceRef("p", "PromptManagerPair", 0), i, snapshot() + ) + result = self.seal(cap) + self.assertEqual(result.state, state) + if count > 1: + self.assertIsNone(result.bundle.contexts[0].positive.effective.value) + + def test_projection_never_retains_original_containers_and_budget_denial(self): + class Envelope(dict): + pass + + g = Envelope(graph()) + ref = weakref.ref(g) + cap = self.begin(g=g) + del g + gc.collect() + self.assertIsNone(ref()) + self.manager.revoke(cap) + small = capture.ContextCapture(10, lambda: self.status, ContextBudget(1)) + cap = small.begin(self.start, 3, graph(), None) + self.assertEqual(small.seal(cap).reason, "budget_exhausted") + + def test_crossed_roles_and_multiple_anchors_use_exact_slots(self): + g = graph() + g["t"] = { + "class_type": "KSampler", + "inputs": {"positive": ["p", 1], "negative": ["p", 0]}, + } + cap = self.begin(g=g) + self.manager.observe( + cap, SourceRef("p", "PromptManagerPair", 0), 0, snapshot(effective="POS") + ) + self.manager.observe( + cap, SourceRef("p", "PromptManagerPair", 1), 0, snapshot(effective="NEG") + ) + result = self.seal(cap) + self.assertEqual(result.state, "complete") + self.assertEqual( + [ + (r.positive.effective.value, r.negative.effective.value) + for r in result.bundle.contexts + ], + [("POS", "NEG"), ("NEG", "POS")], + ) + + def test_transformed_sources_cannot_replace_composed_text(self): + g = graph() + g["text"] = { + "class_type": "StringConcatenate", + "inputs": {"string_a": ["p", 2], "string_b": " suffix"}, + } + g["clip"] = {"class_type": "CLIPTextEncode", "inputs": {"text": ["text", 0]}} + g["s"]["inputs"]["positive"] = ["clip", 0] + cap = self.begin(g=g) + self.manager.observe( + cap, + SourceRef("p", "PromptManagerPair", 2), + 0, + snapshot(effective="TRIGGER"), + ) + result = self.seal(cap) + self.assertIsNone(result.bundle.contexts[0].positive.effective.value) + self.assertNotEqual(result.state, "complete") + + def test_upstream_transform_preserves_final_pm_runtime_snapshot(self): + for field in ("text", "prepend_text", "append_text"): + for unresolved in (False, True): + with self.subTest(field=field, unresolved=unresolved): + cap = self.begin(g=transformed_input_graph(field, unresolved)) + self.host_context() + capture.record_runtime_context( + "PromptManagerPair", + " actual raw ", + " actual pre ", + " actual post ", + " actual injected effective ", + role="positive", + ) + side = self.seal(cap).bundle.contexts[0].positive + for name, expected in ( + ("raw", " actual raw "), + ("prepend", " actual pre "), + ("append", " actual post "), + ("effective", " actual injected effective "), + ): + evidence = getattr(side, name) + self.assertEqual(evidence.value, expected) + self.assertEqual(evidence.provenance, "recorded") + self.assertEqual(evidence.source, "runtime_return") + + def test_unsupported_graph_has_no_supported_context(self): + result = self.seal(self.begin(g={})) + self.assertEqual(result.reason, "no_supported_context") + + def host_context(self, kind="PromptManagerPair", index=0, job=JOB, node_id="p"): + host = types.ModuleType("comfy_execution.utils") + host.get_executing_context = lambda: types.SimpleNamespace( + prompt_id=job, node_id=node_id, list_index=index + ) + stack = self.enter_stack() + stack.enter_context(patch.dict(sys.modules, {"comfy_execution.utils": host})) + return host + + def enter_stack(self): + from contextlib import ExitStack + + stack = ExitStack() + self.addCleanup(stack.close) + return stack + + def test_nodes_record_only_successful_returns_and_save_failure_preserves_text(self): + import logging + from unittest.mock import Mock + from prompt_manager import ( + PromptManager, + PromptManagerFocused, + PromptManagerSplit, + ) + from prompt_manager_text import PromptManagerText + from prompt_manager_pair import PromptManagerPair + + for cls in ( + PromptManager, + PromptManagerFocused, + PromptManagerSplit, + PromptManagerText, + PromptManagerPair, + ): + kind = cls.__name__ + g = graph() + g["p"]["class_type"] = kind + if kind != "PromptManagerPair": + g["s"]["inputs"]["negative"] = ["p", 0] + if kind == "PromptManagerText": + g["c"] = {"class_type": "CLIPTextEncode", "inputs": {"text": ["p", 0]}} + g["s"]["inputs"] = {"positive": ["c", 0], "negative": ["c", 0]} + cap = self.begin(g=g) + self.host_context(kind) + node = cls.__new__(cls) + node.db = Mock() + node.db.get_prompt_by_hash.return_value = None + node.db.save_prompt.side_effect = RuntimeError("offline") + node.db.save_prompt_pair.side_effect = RuntimeError("offline") + node.logger = logging.getLogger("test.capture.nodes") + node.prompt_tracker = Mock() + node.comfyui_integration = Mock() + node._inject_lora_trigger_words = lambda text: text + " TRIGGER" + clip = Mock() + clip.tokenize.side_effect = lambda text: text + clip.encode_from_tokens_scheduled.side_effect = lambda text: ( + "encoded", + text, + ) + kwargs = dict(text=" raw ", prepend_text=" pre ", append_text=" post ") + if kind != "PromptManagerText": + kwargs["clip"] = clip + if kind == "PromptManagerPair": + kwargs["negative_text"] = " neg " + output = ( + node.process_text(**kwargs) + if kind == "PromptManagerText" + else node.encode_prompt(**kwargs) + ) + result = self.seal(cap) + side = result.bundle.contexts[0].positive + self.assertEqual(side.effective.provenance, "recorded", kind) + self.assertEqual(side.raw.value, " raw ") + self.assertEqual(side.prepend.value, " pre ") + self.assertEqual(side.append.value, " post ") + self.assertEqual( + side.effective.value, + output[-2] if kind == "PromptManagerPair" else output[-1], + ) + if kind == "PromptManagerPair": + self.assertEqual( + result.bundle.contexts[0].negative.effective.value, " neg " + ) + if kind != "PromptManagerText": + cap = self.begin(g=g) + clip.encode_from_tokens_scheduled.side_effect = RuntimeError( + "encode failed" + ) + with self.assertRaises(RuntimeError): + node.encode_prompt(**kwargs) + self.assertIsNone( + self.seal(cap).bundle.contexts[0].positive.effective.value + ) + + def test_oversized_projection_and_observation_release_all_reservations(self): + g = graph() + g["p"]["inputs"].update( + text="x" * 65536, prepend_text=" " * 65536, append_text=" " * 65536 + ) + cap = self.begin(g=g) + self.manager.revoke(cap) + self.assertEqual(self.pool.charged_bytes, 0) + cap = self.begin() + self.manager.observe( + cap, + SourceRef("p", "PromptManagerPair", 0), + 0, + snapshot(effective="x" * 65537), + ) + self.assertIsNone(self.seal(cap).bundle.contexts[0].positive.effective.value) + + def test_recorder_rejects_wrong_class_slot_and_noninteger_invocation(self): + for node_type, role, index, node_id in ( + ("PromptManager", None, 0, "p"), + ("PromptManagerPair", "invalid", 0, "p"), + ("PromptManagerPair", "positive", True, "p"), + ("PromptManagerPair", "positive", -1, "p"), + ("PromptManagerPair", "positive", 0, "other"), + ): + cap = self.begin() + self.host_context(index=index, node_id=node_id) + capture.record_runtime_context(node_type, "bad", "", "", "bad", role=role) + self.assertIsNone( + self.seal(cap).bundle.contexts[0].positive.effective.value + ) + + def test_async_inherited_capability_after_await_cannot_record_new_attempt(self): + import asyncio + + async def scenario(): + old = self.begin() + entered, release = asyncio.Event(), asyncio.Event() + + async def late(): + entered.set() + await release.wait() + capture.record_runtime_context( + "PromptManagerPair", "late", "", "", "late", role="positive" + ) + + self.host_context() + task = asyncio.create_task(late()) + await entered.wait() + new = self.begin(2) + release.set() + await task + self.assertIsNone(self.seal(old).bundle) + self.assertIsNone( + self.seal(new).bundle.contexts[0].positive.effective.value + ) + + asyncio.run(scenario()) + + def test_writer_restore_or_close_revokes_callback_paused_during_encoding(self): + import asyncio + + from utils.execution_history_writer import HistoryWriter + + for action in ("invalidate", "close"): + with self.subTest(action=action): + pool = ContextBudget() + current = [HistoryWriterStatus(EPOCH, True, 0, 0, True, 3)] + manager = capture.ContextCapture(10, lambda: current[0], pool) + self.addCleanup(manager.revoke_all) + writer = HistoryWriter( + lambda: self.fail("This callback test must not open storage"), + context_budget=pool, + on_invalidation=manager.revoke_all, + ) + self.addCleanup(writer.close) + self.host_context() + + async def scenario(): + old = manager.begin(self.start, 3, graph(), None) + manager.observe( + old, SourceRef("p", "PromptManagerPair", 0), 0, snapshot() + ) + entered, release = asyncio.Event(), asyncio.Event() + + async def encoding_return(): + entered.set() + await release.wait() + capture.record_runtime_context( + "PromptManagerPair", "late", "", "", "late", role="positive" + ) + + task = asyncio.create_task(encoding_return()) + await entered.wait() + self.assertGreater(pool.charged_bytes, 0) + if action == "invalidate": + writer.invalidate() + else: + self.assertTrue(writer.close()) + self.assertEqual(pool.charged_bytes, 0) + self.assertIsNone(capture.active_capture()) + new_epoch = "44444444-4444-4444-8444-444444444444" + current[0] = HistoryWriterStatus(new_epoch, True, 0, 0, True, 4) + new = manager.begin( + replace( + self.start, + run_id=JOB, + queue_item_id=2, + storage_epoch=new_epoch, + ), + 4, + graph(), + None, + ) + release.set() + await task + self.assertIsNone(manager.seal(old).bundle) + sealed = manager.seal(new) + try: + self.assertIsNone( + sealed.bundle.contexts[0].positive.effective.value + ) + finally: + if sealed.lease: + sealed.lease.release() + self.assertEqual(pool.charged_bytes, 0) + + asyncio.run(scenario()) + + def test_oversized_runtime_field_is_unknown_without_erasing_other_fields(self): + cap = self.begin() + self.host_context() + capture.record_runtime_context( + "PromptManagerPair", + " actual ", + " pre ", + " post ", + "x" * 65537, + role="positive", + ) + result = self.seal(cap) + side = result.bundle.contexts[0].positive + self.assertEqual(side.raw.value, " actual ") + self.assertEqual(side.prepend.value, " pre ") + self.assertIsNone(side.effective.value) + self.assertEqual(side.effective.reason, "limit_reached") + self.assertEqual(result.reason, "field_limit") + + def test_mixed_cache_provenance_cannot_promote_inferred_pm_effective(self): + cap = self.begin() + field = TextEvidence("literal", "inferred", "graph_literal", None) + cached = TextEvidence("cached TRIGGER", "recorded", "host_cache_output", None) + self.manager.observe( + cap, + SourceRef("p", "PromptManagerPair", 0), + 0, + PromptSnapshot(field, cached, field, field), + ) + side = self.seal(cap).bundle.contexts[0].positive + self.assertEqual(side.raw.provenance, "inferred") + self.assertEqual(side.effective.value, "cached TRIGGER") + cap = self.begin(2) + self.manager.observe( + cap, + SourceRef("p", "PromptManagerPair", 0), + 0, + PromptSnapshot(field, field, field, field), + ) + self.assertIsNone(self.seal(cap).bundle.contexts[0].positive.effective.value) + + def test_restore_and_close_release_weakly_bound_payload_immediately(self): + cap = self.begin() + binding = copy_context() + self.manager.revoke_all() + self.assertEqual(self.pool.charged_bytes, 0) + self.assertIsNone(binding.run(capture.active_capture)) + self.manager.observe(cap, SourceRef("p", "PromptManagerPair", 0), 0, snapshot()) + self.assertIsNone(self.seal(cap).bundle) + owner = capture.ContextCapture(11, lambda: self.status, self.pool) + owner.begin(self.start, 3, graph(), None) + binding = copy_context() + owner.revoke_all() + ref = weakref.ref(owner) + del owner + gc.collect() + self.assertIsNone(ref()) + self.assertIsNone(binding.run(capture.active_capture)) + + def test_conflicting_replay_stays_ambiguous_after_limited_other_source(self): + cap = self.begin() + source = SourceRef("p", "PromptManagerPair", 0) + self.manager.observe(cap, source, 0, snapshot(effective="first")) + self.manager.observe(cap, source, 0, snapshot(effective="different")) + self.manager.observe( + cap, + SourceRef("p", "PromptManagerPair", 1), + 0, + snapshot(effective="x" * 65537), + ) + sealed = self.seal(cap) + self.assertEqual(sealed.state, "ambiguous") + self.assertIsNone(sealed.bundle.contexts[0].positive.effective.value) + + def test_capture_owner_bounds_unavailable_attempts_to_128(self): + owner = capture.ContextCapture(12, lambda: self.status, ContextBudget(0)) + self.addCleanup(owner.revoke_all) + capabilities = [ + copy_context().run( + owner.begin, replace(self.start, queue_item_id=i), 3, {}, None + ) + for i in range(129) + ] + self.assertEqual(sum(owner.valid(cap) for cap in capabilities), 128) + + def test_revocation_discards_owned_samples_before_releasing_budget(self): + cap = self.begin() + value = snapshot() + ref = weakref.ref(value.raw) + self.manager.observe(cap, SourceRef("p", "PromptManagerPair", 0), 0, value) + del value + lease = self.manager._attempts[cap.nonce].lease + original = lease.release + + def release(): + self.assertIsNone(ref(), "sample was uncharged while still owned") + original() + + with patch.object(lease, "release", release): + self.manager.revoke(cap) diff --git a/tests/test_execution_context_database.py b/tests/test_execution_context_database.py new file mode 100644 index 0000000..60036d8 --- /dev/null +++ b/tests/test_execution_context_database.py @@ -0,0 +1,613 @@ +"""Immutable scalar capture, deduplication, and conservative navigation hints.""" + +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from uuid import uuid4 +from dataclasses import replace +from pathlib import Path +import subprocess +import sys +import textwrap +from unittest.mock import patch +import unittest + +from database.execution_history import HistoryEpochChanged +from database.execution_context_schema import context_schema_error +from utils.execution_contexts import ( + ContextBundle, + PromptSnapshot, + RecipeContext, + TextEvidence, +) + +try: + from tests import test_execution_history_database as history_tests +except ImportError: + import test_execution_history_database as history_tests + + +def known(value): + return TextEvidence(value, "recorded", "runtime_return", None) + + +def snapshot(raw): + return PromptSnapshot(known(raw), known(raw), known(""), known("")) + + +def bundle(negative="negative", positive="positive", anchor="4"): + return ContextBundle( + (RecipeContext(anchor, snapshot(positive), snapshot(negative), "inferred"),), + "complete", + "captured", + ) + + +class ContextDatabaseTests(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 + + def capture(self, db, item=0, capture=None): + start = self.start(db, item) + db.write_execution_start(start) + db.write_execution_finish(self.finish(start)) + self.assertTrue( + db.write_execution_contexts( + start.run_id, start.storage_epoch, capture or bundle() + ) + ) + return start + + def counts(self, db): + conn = db.model.get_connection() + return tuple( + conn.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0] + for table in ( + "execution_runs", + "execution_contexts", + "execution_prompt_snapshots", + "execution_context_capture", + ) + ) + + def test_finish_capture_uses_stored_parent_identity(self): + db = self.database() + start = replace(self.start(db), run_id=str(uuid4())) + db.write_execution_start(start) + finish = self.finish(start) + self.assertTrue(db.write_execution_contexts_for_finish(finish, bundle())) + self.assertEqual(db.get_execution_contexts(start.run_id)["state"], "complete") + missing = replace(finish, queue_item_id=999) + self.assertFalse(db.write_execution_contexts_for_finish(missing, bundle())) + self.assertEqual(self.counts(db), (1, 1, 2, 1)) + + def test_finish_generation_guard_encloses_commit_and_rollback(self): + db = self.database() + start = self.start(db) + db.write_execution_start(start) + finish = self.finish(start) + observed = [] + + @contextmanager + def guard(): + observed.append(db.model.get_connection().in_transaction) + try: + yield True + finally: + observed.append(db.model.get_connection().in_transaction) + + self.assertTrue( + db.write_execution_contexts_for_finish( + finish, bundle(), generation_guard=guard + ) + ) + self.assertEqual(observed, [True, False]) + observed.clear() + other = replace(self.start(db, 1), run_id=str(uuid4())) + db.write_execution_start(other) + with patch( + "database.execution_contexts._insert_exact", + side_effect=ValueError("collision"), + ): + with self.assertRaises(ValueError): + db.write_execution_contexts_for_finish( + self.finish(other), bundle(), generation_guard=guard + ) + self.assertEqual(observed, [True, False]) + self.assertEqual( + db.get_execution_contexts(other.run_id)["state"], "not_captured" + ) + + def test_finish_generation_guard_rejects_without_capture(self): + db = self.database() + start = self.start(db) + db.write_execution_start(start) + + @contextmanager + def stale(): + yield False + + self.assertFalse( + db.write_execution_contexts_for_finish( + self.finish(start), bundle(), generation_guard=stale + ) + ) + self.assertEqual(self.counts(db), (1, 0, 0, 0)) + + def test_a_a_b_cardinality(self): + db = self.database() + self.assertTrue( + callable(getattr(db, "write_execution_contexts", None)), + "PromptDatabase must expose immutable context persistence", + ) + starts = [self.capture(db, i, bundle(n)) for i, n in enumerate(("A", "A", "B"))] + self.assertEqual(self.counts(db), (3, 2, 3, 3)) + reads = [db.get_execution_contexts(s.run_id) for s in starts] + self.assertEqual( + reads[0]["contexts"][0]["context_id"], reads[1]["contexts"][0]["context_id"] + ) + self.assertNotEqual( + reads[0]["contexts"][0]["context_id"], reads[2]["contexts"][0]["context_id"] + ) + self.assertIsNone(context_schema_error(db.model.get_connection())) + + def test_parallel_identical_inserts_deduplicate(self): + db = self.database() + starts = [self.start(db, i) for i in range(8)] + for start in starts: + db.write_execution_start(start) + + def write(start): + return db.write_execution_contexts( + start.run_id, start.storage_epoch, bundle() + ) + + with ThreadPoolExecutor(max_workers=4) as pool: + self.assertEqual(list(pool.map(write, starts)), [True] * 8) + self.assertEqual(self.counts(db), (8, 1, 2, 8)) + + def test_replay_freezes_first_capture(self): + db = self.database() + start = self.capture(db) + original = db.get_execution_contexts(start.run_id) + counts = self.counts(db) + self.assertTrue( + db.write_execution_contexts(start.run_id, start.storage_epoch, bundle()) + ) + db.write_execution_finish(self.finish(start)) + self.assertEqual(db.get_execution_contexts(start.run_id), original) + self.assertTrue( + db.write_execution_contexts( + start.run_id, start.storage_epoch, bundle("changed") + ) + ) + read = db.get_execution_contexts(start.run_id) + self.assertEqual( + (read["state"], read["reason"]), ("conflicting", "conflicting") + ) + self.assertEqual(read["contexts"], original["contexts"]) + self.assertEqual(self.counts(db), counts) + db.write_execution_contexts(start.run_id, start.storage_epoch, bundle()) + self.assertEqual(db.get_execution_contexts(start.run_id), read) + + def test_evidence_replay_and_anchor_changes_conflict_without_inserting(self): + for change in ("evidence", "anchor", "status"): + with self.subTest(change=change): + db = self.database(str(Path(self.temp.name) / f"{change}.db")) + start = self.capture(db) + original = db.get_execution_contexts(start.run_id)["contexts"] + candidate = bundle() + if change == "evidence": + ctx = candidate.contexts[0] + candidate = replace( + candidate, + contexts=( + replace( + ctx, + positive=replace( + ctx.positive, + raw=TextEvidence( + "positive", "inferred", "graph_literal", None + ), + ), + ), + ), + ) + elif change == "anchor": + candidate = bundle(anchor="5") + else: + candidate = replace(candidate, state="partial", reason="partial") + db.write_execution_contexts( + start.run_id, start.storage_epoch, candidate + ) + read = db.get_execution_contexts(start.run_id) + self.assertEqual(read["state"], "conflicting") + self.assertEqual(read["contexts"], original) + self.assertEqual(self.counts(db), (1, 1, 2, 1)) + + def test_hash_collision_rejected_without_partial_insert(self): + db = self.database() + start = self.capture(db) + original = db.get_execution_contexts(start.run_id) + second = self.start(db, 1) + db.write_execution_start(second) + # A collision must compare stored scalar content, then roll back all rows. + from utils.execution_contexts import content_digest + + collided = content_digest("text", ("negative",)) + + def digest(kind, fields): + return ( + collided + if kind == "text" and fields == ("changed",) + else content_digest(kind, fields) + ) + + before = tuple(db.model.get_connection().iterdump()) + with patch("database.execution_contexts.content_digest", side_effect=digest): + with self.assertRaises(ValueError): + db.write_execution_contexts( + second.run_id, + second.storage_epoch, + bundle("changed", "new-positive"), + ) + self.assertEqual(tuple(db.model.get_connection().iterdump()), before) + self.assertEqual(db.get_execution_contexts(start.run_id), original) + + def test_missing_parent_and_epoch_mismatch_insert_nothing(self): + db = self.database() + start = self.start(db) + self.assertEqual( + db.get_execution_contexts(start.run_id)["state"], "not_captured" + ) + self.assertFalse( + db.write_execution_contexts(start.run_id, start.storage_epoch, bundle()) + ) + self.assertEqual(self.counts(db), (0, 0, 0, 0)) + db.write_execution_start(start) + with self.assertRaises(HistoryEpochChanged): + db.write_execution_contexts( + start.run_id, "11111111-1111-4111-8111-111111111111", bundle() + ) + with db.model.get_connection() as conn: + conn.execute( + "UPDATE execution_runs SET storage_epoch=? WHERE run_id=?", + ("11111111-1111-4111-8111-111111111111", start.run_id), + ) + self.assertFalse( + db.write_execution_contexts(start.run_id, start.storage_epoch, bundle()) + ) + self.assertEqual(self.counts(db), (1, 0, 0, 0)) + + def test_invalid_bundle_has_no_partial_rows(self): + db = self.database() + start = self.start(db) + db.write_execution_start(start) + for candidate in ( + replace(bundle(), reason="arbitrary"), + replace(bundle(), contexts=()), + replace(bundle(), contexts=(RecipeContext("4", None, None, "inferred"),)), + ): + with self.subTest(candidate=candidate), self.assertRaises(ValueError): + db.write_execution_contexts( + start.run_id, start.storage_epoch, candidate + ) + self.assertEqual(self.counts(db), (1, 0, 0, 0)) + + def test_unknown_empty_unicode_and_evidence_round_trip(self): + db = self.database() + unknown = TextEvidence(None, "unknown", "unknown", "missing") + ctx = RecipeContext( + "3", + PromptSnapshot(known(" e\u0301 \n"), unknown, known(""), known("é")), + None, + "inferred", + ) + start = self.capture(db, capture=ContextBundle((ctx,), "partial", "missing")) + read = db.get_execution_contexts(start.run_id)["contexts"][0] + self.assertEqual( + read["positive"], + {"raw": " e\u0301 \n", "effective": None, "prepend": "", "append": "é"}, + ) + self.assertIsNone(read["negative"]) + self.assertEqual( + read["evidence"]["positive"]["effective"], + {"provenance": "unknown", "source": "unknown", "reason": "missing"}, + ) + self.assertEqual( + set(read["evidence"]), + {"format_version", "role_provenance", "positive", "negative"}, + ) + self.assertIsNone(context_schema_error(db.model.get_connection())) + + def test_library_mutations_and_id_reuse_do_not_rewrite_history(self): + db = self.database() + conn = db.model.get_connection() + with conn: + conn.execute( + "INSERT INTO prompts(id,text) VALUES(1,'positive'),(2,'negative')" + ) + conn.execute( + "INSERT INTO prompt_pairs(id,positive_prompt_id,negative_prompt_id) " + "VALUES(7,1,2)" + ) + start = self.capture(db) + original = db.get_execution_contexts(start.run_id)["contexts"][0] + self.assertEqual( + ( + original["positive_prompt_id"], + original["negative_prompt_id"], + original["pair_id"], + ), + (1, 2, 7), + ) + db.archive_prompt_pair(7) + self.assertEqual( + db.get_execution_contexts(start.run_id)["contexts"][0], original + ) + db.delete_prompt_pair(7) + db.update_prompt_text(1, "edited") + read = db.get_execution_contexts(start.run_id)["contexts"][0] + self.assertIsNone(read["positive_prompt_id"]) + self.assertIsNone(read["pair_id"]) + db.delete_prompt(2) + with conn: + conn.execute("INSERT INTO prompts(id,text) VALUES(2,'reused')") + self.assertIsNone( + db.get_execution_contexts(start.run_id)["contexts"][0]["negative_prompt_id"] + ) + self.assertEqual(db.prune_orphaned_prompts(), 2) + after = db.get_execution_contexts(start.run_id)["contexts"][0] + for key in ("positive", "negative", "evidence", "context_id"): + self.assertEqual(after[key], original[key]) + self.assertEqual(conn.execute("SELECT COUNT(*) FROM prompts").fetchone()[0], 0) + self.assertIsNone(context_schema_error(conn)) + + def test_hints_require_exact_unique_raw_text_and_never_refresh(self): + db = self.database() + conn = db.model.get_connection() + with conn: + conn.execute( + "INSERT INTO prompts(id,text) VALUES(1,' positive " + "'),(2,'negative'),(3,'negative')" + ) + start = self.capture(db) + read = db.get_execution_contexts(start.run_id)["contexts"][0] + self.assertIsNone(read["positive_prompt_id"]) + self.assertIsNone(read["negative_prompt_id"]) + with conn: + conn.execute("DELETE FROM prompts WHERE id=3") + conn.execute("UPDATE prompts SET text='positive' WHERE id=1") + db.write_execution_contexts(start.run_id, start.storage_epoch, bundle()) + self.assertEqual(db.get_execution_contexts(start.run_id)["contexts"], [read]) + fresh = self.capture(db, 1) + self.assertEqual( + db.get_execution_contexts(fresh.run_id)["contexts"][0][ + "positive_prompt_id" + ], + 1, + ) + with conn: + conn.execute("INSERT INTO prompts(id,text) VALUES(3,'positive')") + self.assertIsNone( + db.get_execution_contexts(fresh.run_id)["contexts"][0]["positive_prompt_id"] + ) + + def test_null_negative_pair_only_matches_known_exact_empty(self): + db = self.database() + with db.model.get_connection() as conn: + conn.execute("INSERT INTO prompts(id,text) VALUES(1,'positive')") + conn.execute("INSERT INTO prompt_pairs(id,positive_prompt_id) VALUES(7,1)") + for item, raw, expected in ((0, "", 7), (1, " ", None), (2, None, None)): + candidate = bundle(negative=raw or "") + if raw is None: + ctx = candidate.contexts[0] + candidate = replace( + candidate, + contexts=( + replace( + ctx, + negative=replace( + ctx.negative, + raw=TextEvidence(None, "unknown", "unknown", "missing"), + ), + ), + ), + state="partial", + reason="partial", + ) + start = self.capture(db, item, candidate) + self.assertEqual( + db.get_execution_contexts(start.run_id)["contexts"][0]["pair_id"], + expected, + ) + + def test_snapshot_and_context_collisions_roll_back_all_new_rows(self): + from utils.execution_contexts import content_digest + + for kind in ("snapshot", "context"): + with self.subTest(kind=kind): + db = self.database(str(Path(self.temp.name) / f"{kind}.db")) + first = self.capture(db) + old_context = db.get_execution_contexts(first.run_id)["contexts"][0][ + "context_id" + ] + old_snapshot = ( + db.model.get_connection() + .execute( + "SELECT negative_snapshot_id FROM execution_contexts WHERE " + "context_id=?", + (old_context,), + ) + .fetchone()[0] + ) + second = self.start(db, 1) + db.write_execution_start(second) + before = tuple(db.model.get_connection().iterdump()) + + def digest(current_kind, fields): + if current_kind == kind: + return old_snapshot if kind == "snapshot" else old_context + return content_digest(current_kind, fields) + + with ( + patch( + "database.execution_contexts.content_digest", side_effect=digest + ), + self.assertRaises(ValueError), + ): + db.write_execution_contexts( + second.run_id, + second.storage_epoch, + bundle("changed", "new-positive"), + ) + self.assertEqual(tuple(db.model.get_connection().iterdump()), before) + + def test_parallel_contradictory_replay_keeps_one_first_capture(self): + db = self.database() + start = self.start(db) + db.write_execution_start(start) + + def write(negative): + return db.write_execution_contexts( + start.run_id, start.storage_epoch, bundle(negative) + ) + + with ThreadPoolExecutor(max_workers=2) as pool: + self.assertEqual(list(pool.map(write, ("A", "B"))), [True, True]) + read = db.get_execution_contexts(start.run_id) + self.assertEqual(read["state"], "conflicting") + self.assertIn(read["contexts"][0]["negative"]["raw"], ("A", "B")) + self.assertEqual(self.counts(db), (1, 1, 2, 1)) + + def test_merge_never_rebinds_a_hint_to_normalized_equal_primary(self): + db = self.database() + with db.model.get_connection() as conn: + conn.execute( + "INSERT INTO prompts(id,text,created_at) VALUES(1,' POSITIVE " + "','2020-01-01'),(2,'positive','2021-01-01')" + ) + start = self.capture(db) + before = db.get_execution_contexts(start.run_id)["contexts"][0] + self.assertEqual(before["positive_prompt_id"], 2) + self.assertEqual(db.cleanup_duplicates(), 1) + after = db.get_execution_contexts(start.run_id)["contexts"][0] + self.assertIsNone(after["positive_prompt_id"]) + self.assertEqual(after["positive"], before["positive"]) + self.assertEqual(after["evidence"], before["evidence"]) + self.assertEqual( + db.model.get_connection().execute("SELECT text FROM prompts").fetchone()[0], + " POSITIVE ", + ) + + def test_ambiguous_empty_negative_omits_pair_hint(self): + db = self.database() + with db.model.get_connection() as conn: + conn.execute( + "INSERT INTO prompts(id,text) VALUES(1,'positive'),(2,''),(3,'')" + ) + conn.execute("INSERT INTO prompt_pairs(id,positive_prompt_id) VALUES(7,1)") + start = self.capture(db, capture=bundle(negative="")) + read = db.get_execution_contexts(start.run_id)["contexts"][0] + self.assertIsNone(read["negative_prompt_id"]) + self.assertIsNone(read["pair_id"]) + + def test_replay_checks_scalar_content_even_if_digests_collide(self): + db = self.database() + start = self.capture(db) + before = tuple(db.model.get_connection().iterdump()) + from utils.execution_contexts import content_digest + + original_snapshot = content_digest("snapshot", ("negative", "negative", "", "")) + + def digest(kind, fields): + if kind == "snapshot" and fields == ("changed", "changed", "", ""): + return original_snapshot + return content_digest(kind, fields) + + with ( + patch("database.execution_contexts.content_digest", side_effect=digest), + self.assertRaises(ValueError), + ): + db.write_execution_contexts( + start.run_id, start.storage_epoch, bundle("changed") + ) + self.assertEqual(tuple(db.model.get_connection().iterdump()), before) + + def test_package_qualified_capture_with_and_without_repository_on_sys_path(self): + # A host imports node packs under a package name, not as top-level utils. + script = textwrap.dedent("""\ + import importlib + import sys + import tempfile + import types + from pathlib import Path + root, expose_root = sys.argv[1:] + if expose_root == "yes": + sys.path.insert(0, root) + else: + assert all(Path(p).resolve() != Path(root) for p in sys.path) + package = types.ModuleType("context_pack") + package.__path__ = [root] + sys.modules[package.__name__] = package + records = importlib.import_module("context_pack.utils.execution_contexts") + operations = importlib.import_module("context_pack.database.operations") + events = importlib.import_module( + "context_pack.utils.execution_history_events" + ) + schema = importlib.import_module( + "context_pack.database.execution_context_schema" + ) + with tempfile.TemporaryDirectory() as directory: + db = operations.PromptDatabase(str(Path(directory) / "capture.db")) + try: + start = events.make_start( + "11111111-1111-4111-8111-111111111111", + ((0, "22222222-2222-4222-8222-222222222222", {}, {}), 0), + db.get_execution_epoch(), 100, + ) + db.write_execution_start(start) + def known(value): + return records.TextEvidence( + value, "recorded", "runtime_return", None + ) + snapshot = records.PromptSnapshot( + known("raw"), known("effective"), known(""), known("") + ) + context = records.RecipeContext( + "4", snapshot, snapshot, "inferred" + ) + bundle = records.ContextBundle( + (context,), "complete", "captured" + ) + assert db.write_execution_contexts( + start.run_id, start.storage_epoch, bundle + ) + read = db.get_execution_contexts(start.run_id) + assert read["state"] == "complete", read + assert read["contexts"][0]["positive"]["raw"] == "raw", read + role = read["contexts"][0]["positive"] + assert role["effective"] == "effective", read + error = schema.context_schema_error(db.model.get_connection()) + assert error is None, error + verified = db.model.verify_database_file(db.model.db_path) + assert verified == (True, "ok"), verified + if expose_root == "no": + assert "utils.execution_contexts" not in sys.modules + assert "database.execution_context_schema" not in sys.modules + finally: + db.close_all() + """) + root = str(Path(__file__).resolve().parents[1]) + for expose_root in ("yes", "no"): + with self.subTest(repository_on_sys_path=expose_root): + result = subprocess.run( + [sys.executable, "-I", "-c", script, root, expose_root], + cwd=self.temp.name, + capture_output=True, + text=True, + timeout=30, + ) + self.assertEqual(result.returncode, 0, result.stdout + result.stderr) diff --git a/tests/test_execution_context_graph.py b/tests/test_execution_context_graph.py new file mode 100644 index 0000000..8babd5b --- /dev/null +++ b/tests/test_execution_context_graph.py @@ -0,0 +1,340 @@ +"""Projection must retain declared roles without claiming sampler consumption.""" + +import unittest +from test_execution_contexts import load_module + + +def node(kind, **inputs): + return {"class_type": kind, "inputs": inputs} + + +def graph(kind="KSampler", source="PromptManagerPair"): + return { + "p": node( + source, + text=" cat ", + negative_text="", + prepend_text="pre", + append_text="post", + ), + "s": node(kind, positive=["p", 0], negative=["p", 1]), + } + + +def transformed_input_graph(field: str, unresolved: bool = False) -> dict: + g = graph() + g["upstream"] = node("PromptManagerText", text="looks literal") + g["dynamic"] = node("RuntimeGenerator") + g["join"] = node( + "StringConcatenate", + string_a=["upstream", 0], + string_b=["dynamic", 0] if unresolved else " suffix", + ) + if unresolved: + g["inner"] = node( + "StringConcatenate", string_a=["upstream", 0], string_b=" inner" + ) + g["join"]["inputs"]["string_a"] = ["inner", 0] + g["p"]["inputs"][field] = ["join", 0] + return g + + +class ProjectionTests(unittest.TestCase): + def setUp(self): + self.m = load_module("execution_context_graph") + + def test_projection_preserves_roles(self): + for kind in ("KSampler", "KSamplerAdvanced", "SamplerCustom"): + g = graph(kind) + result = self.m.project_contexts(g, None) + self.assertEqual(len(result.recipes), 1) + recipe = result.recipes[0] + self.assertEqual(recipe.anchor_id, "s") + self.assertEqual(recipe.positive.inferred.raw.value, " cat ") + self.assertEqual(recipe.positive.inferred.effective.value, "pre cat post") + self.assertEqual(recipe.negative.inferred.effective.value, "") + self.assertEqual(recipe.positive.sources[0].output_slot, 0) + g["s"]["inputs"]["positive"], g["s"]["inputs"]["negative"] = ["p", 1], [ + "p", + 0, + ] + crossed = self.m.project_contexts(g, None).recipes[0] + self.assertEqual(crossed.positive.inferred.raw.value, "") + self.assertEqual(crossed.negative.sources[0].output_slot, 0) + for source in ("PromptManager", "PromptManagerFocused", "PromptManagerSplit"): + g = graph(source=source) + g["clip"] = node("CLIPTextEncode", text=["p", 1]) + g["s"]["inputs"]["negative"] = ["clip", 0] + recipe = self.m.project_contexts(g, None).recipes[0] + self.assertEqual(recipe.negative.inferred.raw.value, " cat ") + self.assertEqual(recipe.negative.sources[0].output_slot, 1) + for source, slot in ( + ("PromptManagerText", 0), + ("PromptManagerPair", 2), + ("PromptManagerPair", 3), + ): + g = graph(source=source) + g["clip"] = node("CLIPTextEncode", text=["p", slot]) + g["s"]["inputs"] = {"positive": ["clip", 0], "negative": ["clip", 0]} + self.assertEqual( + self.m.project_contexts(g, None) + .recipes[0] + .positive.sources[0] + .output_slot, + slot, + ) + g = graph() + g["p"] = node("CLIPTextEncode", text=["literal", 0]) + g["literal"] = node("PrimitiveString", value=" exact ") + g["s"]["inputs"]["negative"] = ["p", 0] + self.assertEqual( + self.m.project_contexts(g, None) + .recipes[0] + .positive.inferred.effective.value, + " exact ", + ) + g["p"] = node("UnknownTransform", text="secret") + role = self.m.project_contexts(g, None).recipes[0].positive + self.assertIsNone(role.inferred) + self.assertTrue(role.reason) + + def test_output_selection_and_scalar_isolation(self): + g = graph() + g["other"] = node("KSampler", positive=["p", 0], negative=["p", 1]) + g["decode"] = node("VAEDecode", samples=["s", 0]) + g["save"] = node("SaveImage", images=["decode", 0]) + result = self.m.project_contexts(g, ["save"]) + self.assertEqual(tuple(r.anchor_id for r in result.recipes), ("s",)) + g["p"]["inputs"]["text"] = "changed" + self.assertEqual(result.recipes[0].positive.inferred.raw.value, " cat ") + self.assertEqual(self.m.project_contexts(g, []).recipes, ()) + for outputs in (["missing"], [object()], "save"): + result = self.m.project_contexts(g, outputs) + self.assertEqual(result.recipes, ()) + self.assertTrue(result.reason) + + def test_selected_ancestry_ignores_unrelated_widget_arrays(self): + for widget in ( + [1, 2, 3], + ("s", 0, "widget"), + [], + ["s"], + ["s", "0"], + ["s", True], + ["s", -1], + [["s", 0], 0], + ): + with self.subTest(widget=widget): + g = graph() + g["other"] = node("KSampler", positive=["p", 0], negative=["p", 1]) + g["adapter"] = node("CustomAdapter", samples=("s", 0), widget=widget) + g["save"] = node("CustomOutput", images=["adapter", 0]) + result = self.m.project_contexts(g, ["save"]) + self.assertEqual(tuple(r.anchor_id for r in result.recipes), ("s",)) + self.assertEqual(result.recipes[0].positive.inferred.raw.value, " cat ") + + def test_two_element_link_shapes_are_followed_conservatively(self): + g = graph() + # A typed two-element value cannot be distinguished from a real link. + g["save"] = node("CustomOutput", widget=["s", 0]) + self.assertEqual( + tuple(r.anchor_id for r in self.m.project_contexts(g, ["save"]).recipes), + ("s",), + ) + for value in (["missing", 0], ["x" * 129, 0]): + g["save"]["inputs"]["widget"] = value + result = self.m.project_contexts(g, ["save"]) + self.assertEqual(result.recipes, ()) + self.assertTrue(result.reason) + + def test_projection_limits(self): + for count, succeeds in ((500, True), (501, False)): + g = graph() + g.update({str(i): node("Unused") for i in range(count - 2)}) + self.assertEqual(bool(self.m.project_contexts(g, None).recipes), succeeds) + for count, succeeds in ((16, True), (17, False)): + g = {"p": graph()["p"]} + g.update({str(i): graph()["s"] for i in range(count)}) + self.assertEqual(bool(self.m.project_contexts(g, None).recipes), succeeds) + for depth, succeeds in ((64, True), (65, False)): + g = graph() + # The CLIP plus reroutes plus literal give exactly depth visited links. + g["p"] = node("CLIPTextEncode", text=["0", 0]) + for i in range(depth - 2): + g[str(i)] = node("Reroute", value=[str(i + 1), 0]) + g[str(depth - 2)] = node("PrimitiveString", value="known") + role = self.m.project_contexts(g, None).recipes[0].positive + self.assertEqual(role.inferred is not None, succeeds) + g = graph() + g["s" * 129] = g.pop("s") + self.assertFalse(self.m.project_contexts(g, None).recipes) + + def test_dynamic_pm_inputs_preserve_observable_sources(self): + for kind, positive_slot, negative_slot in ( + ("PromptManagerText", 0, 0), + ("PromptManagerPair", 2, 3), + ("PromptManager", 1, 1), + ("PromptManagerFocused", 1, 1), + ("PromptManagerSplit", 1, 1), + ("PromptManagerPair", 0, 1), + ): + with self.subTest(kind=kind, slots=(positive_slot, negative_slot)): + g = graph(source=kind) + g["generator"] = node("UnknownGenerator") + g["p"]["inputs"]["text"] = ["generator", 0] + g["p"]["inputs"]["negative_text"] = ["generator", 0] + g["pos"] = node("CLIPTextEncode", text=["p", positive_slot]) + g["neg"] = node("CLIPTextEncode", text=["p", negative_slot]) + g["s"]["inputs"] = {"positive": ["pos", 0], "negative": ["neg", 0]} + if kind == "PromptManagerPair" and positive_slot == 0: + g["s"]["inputs"] = {"positive": ["p", 1], "negative": ["p", 0]} + positive_slot, negative_slot = 1, 0 + result = self.m.project_contexts(g, None) + recipe = result.recipes[0] + for role, slot in ( + (recipe.positive, positive_slot), + (recipe.negative, negative_slot), + ): + self.assertEqual( + tuple(ref.output_slot for ref in role.sources), (slot,) + ) + self.assertIsNone(role.inferred) + self.assertTrue(role.reason) + self.assertTrue(result.sources) + + def test_graph_evidence_distinguishes_literals_from_semantics(self): + result = self.m.project_contexts(graph(), None) + evidence = result.recipes[0].positive.inferred + self.assertEqual(evidence.raw.provenance, "inferred") + self.assertEqual(evidence.raw.source, "graph_literal") + self.assertEqual(evidence.effective.source, "graph_semantics") + + def test_derived_raw_and_composition_preserve_field_sources(self): + g = graph(source="PromptManagerText") + g["join"] = node( + "Text Concatenate", + text_a=" a ", + text_b=" b ", + delimiter="-", + clean_whitespace="true", + ) + g["p"]["inputs"]["text"] = ["join", 0] + g["clip"] = node("CLIPTextEncode", text=["p", 0]) + g["s"]["inputs"] = {"positive": ["clip", 0], "negative": ["clip", 0]} + role = self.m.project_contexts(g, None).recipes[0].positive + self.assertIsNotNone(role.inferred) + self.assertEqual(role.inferred.raw.value, "a-b") + self.assertEqual(role.inferred.raw.source, "graph_semantics") + self.assertEqual(role.inferred.prepend.source, "graph_literal") + + def test_selected_custom_outputs_and_multistage_ancestry(self): + g = graph() + g["earlier"] = node("SamplerCustom", positive=["p", 0], negative=["p", 1]) + g["s"]["inputs"]["latent_image"] = ["earlier", 0] + g["adapter"] = node("UnknownImageAdapter", source=["s", 0]) + g["custom"] = node("CustomOutput", images=["adapter", 0]) + result = self.m.project_contexts(g, ["custom"]) + self.assertEqual(tuple(r.anchor_id for r in result.recipes), ("s", "earlier")) + + def test_total_resolution_visits_are_bounded(self): + # Sixteen recipes each resolving two independent chains of length 16 + # exceed 500 total visits even though the graph has fewer than 500 nodes. + g = {"v": node("PrimitiveString", value="known")} + for i in range(15): + g[str(i)] = node("Reroute", value=[str(i + 1), 0] if i < 14 else ["v", 0]) + g["clip"] = node("CLIPTextEncode", text=["0", 0]) + for i in range(16): + g["s" + str(i)] = node( + "KSampler", positive=["clip", 0], negative=["clip", 0] + ) + result = self.m.project_contexts(g, None) + self.assertTrue(result.reason) + self.assertFalse( + result.recipes + and all( + r.positive.inferred is not None and r.negative.inferred is not None + for r in result.recipes + ) + ) + + def test_known_concatenation_remains_inferred(self): + g = graph(source="PromptManagerText") + g["join"] = node( + "StringConcatenate", string_a=["p", 0], string_b="!", delimiter="" + ) + g["clip"] = node("CLIPTextEncode", text=["join", 0]) + g["s"]["inputs"] = {"positive": ["clip", 0], "negative": ["clip", 0]} + role = self.m.project_contexts(g, None).recipes[0].positive + self.assertIsNotNone(role.inferred) + self.assertEqual(role.inferred.effective.value, "pre cat post!") + self.assertEqual(role.inferred.effective.provenance, "inferred") + self.assertEqual(role.sources[0].node_id, "p") + self.assertEqual(role.reason, "source-transformed") + + def test_transform_provenance_stops_at_selected_pm_capture_boundary(self): + for field in ("text", "prepend_text", "append_text"): + for unresolved in (False, True): + with self.subTest(field=field, unresolved=unresolved): + g = transformed_input_graph(field, unresolved) + g["outer"] = node( + "StringConcatenate", string_a=["p", 2], string_b=" outer" + ) + g["clip"] = node("CLIPTextEncode", text=["outer", 0]) + g["s"]["inputs"]["negative"] = ["clip", 0] + result = self.m.project_contexts(g, None) + positive, negative = ( + result.recipes[0].positive, + result.recipes[0].negative, + ) + self.assertEqual( + [(ref.node_id, ref.output_slot) for ref in positive.sources], + [("p", 0)], + ) + self.assertNotEqual(positive.reason, "source-transformed") + self.assertEqual(negative.reason, "source-transformed") + self.assertEqual(negative.sources[0].node_id, "p") + + def test_malformed_inputs_do_not_stringify(self): + class HostValue: + def __str__(self): + raise AssertionError("host object stringified") + + for g in ( + {HostValue(): node("KSampler")}, + {"s": node(HostValue())}, + {"s": node("KSampler", positive=HostValue())}, + ): + self.m.project_contexts(g, None) + + def test_composed_affixes_and_retained_projection_are_bounded(self): + g = graph() + g["p"]["inputs"].update( + text="x" * 65536, prepend_text="y" * 65536, append_text="z" * 65536 + ) + result = self.m.project_contexts(g, None) + self.assertIsNone(result.recipes[0].positive.inferred) + g = graph() + g["p"]["inputs"]["text"] = "x" * 60000 + for i in range(15): + g[str(i)] = node("KSampler", positive=["p", 0], negative=["p", 0]) + result = self.m.project_contexts(g, None) + self.assertEqual(result.recipes, ()) + + def test_partial_source_snapshots_keep_literal_fields_without_stateful_substitution( + self, + ): + for upstream in ("RuntimeGenerator", "PromptManagerText"): + g = graph() + g["dynamic"] = node(upstream, text="looks literal") + g["p"]["inputs"]["text"] = ["dynamic", 0] + projection = self.m.project_contexts(g, None) + positive = next( + snapshot + for ref, snapshot in projection.source_snapshots + if ref.output_slot == 0 + ) + self.assertIsNone(positive.raw.value) + self.assertIsNone(positive.effective.value) + self.assertEqual(positive.prepend.value, "pre") + self.assertEqual(positive.append.value, "post") + self.assertEqual(positive.append.provenance, "inferred") diff --git a/tests/test_execution_context_host.py b/tests/test_execution_context_host.py new file mode 100644 index 0000000..7e61729 --- /dev/null +++ b/tests/test_execution_context_host.py @@ -0,0 +1,447 @@ +"""Read-only cached STRING evidence and cooperative executor notifications.""" + +from collections import namedtuple +from contextvars import Context +from dataclasses import replace +import gc +import sys +import weakref +import threading +import types +import unittest +from unittest.mock import patch + +from test_execution_context_graph import transformed_input_graph + +from utils.execution_context_budget import ContextBudget +from utils.execution_context_capture import ContextCapture +from utils.execution_history_events import RunStart +from utils.execution_history_writer import HistoryWriterStatus + +try: + from utils import execution_context_host as host +except ImportError: + host = None + +SESSION = "11111111-1111-4111-8111-111111111111" +EPOCH = "22222222-2222-4222-8222-222222222222" +JOB = "33333333-3333-4333-8333-333333333333" +Entry = namedtuple("CacheEntry", "ui outputs") + + +class HostTests(unittest.TestCase): + def setUp(self): + self.assertIsNotNone(host, "read-only host adapter is missing") + self.queue = object() + self.pool = ContextBudget() + self.status = HistoryWriterStatus(EPOCH, True, 0, 0, True, 3) + self.capture = ContextCapture(id(self.queue), lambda: self.status, self.pool) + self.addCleanup(self.capture.revoke_all) + self.calls = [] + calls = self.calls + + class BasicCache: + def __init__(self): + self.entries = {} + + def get_local(self, node_id): + calls.append(("local", node_id)) + return self.entries.get(node_id) + + async def get(self, node_id): + raise AssertionError("async get must never be called") + + class HierarchicalCache(BasicCache): + pass + + class LRUCache(BasicCache): + pass + + class RAMPressureCache(LRUCache): + pass + + class NullCache: + def get_local(self, node_id): + raise AssertionError("null cache is not usable") + + class PromptExecutor: + def add_message(self, event, data, broadcast): + calls.append(("original", event, data, broadcast)) + if event == "raise": + raise RuntimeError("original") + return "original-return" + + self.execution = types.ModuleType("execution") + self.execution.PromptExecutor = PromptExecutor + self.execution.CacheEntry = Entry + self.caching = types.ModuleType("comfy_execution.caching") + for cls in ( + BasicCache, + HierarchicalCache, + LRUCache, + RAMPressureCache, + NullCache, + ): + setattr(self.caching, cls.__name__, cls) + package = types.ModuleType("comfy_execution") + package.caching = self.caching + self.patch = patch.dict( + sys.modules, + { + "comfy_execution": package, + "execution": self.execution, + "comfy_execution.caching": self.caching, + }, + ) + self.patch.start() + self.addCleanup(self.patch.stop) + self.executor = PromptExecutor() + self.executor.server = types.SimpleNamespace(prompt_queue=self.queue) + self.executor.caches = types.SimpleNamespace(outputs=HierarchicalCache()) + + def begin( + self, kind="PromptManager", raw=" raw ", prepend=" pre ", append=" post " + ): + g = { + "p": { + "class_type": kind, + "inputs": { + "text": raw, + "negative_text": "negative", + "prepend_text": prepend, + "append_text": append, + }, + }, + "clip": {"class_type": "CLIPTextEncode", "inputs": {"text": ["p", 0]}}, + "s": { + "class_type": "KSampler", + "inputs": { + "positive": ( + ["clip", 0] if kind == "PromptManagerText" else ["p", 0] + ), + "negative": ( + ["p", 1] if kind == "PromptManagerPair" else ["clip", 0] + ), + }, + }, + } + if type(raw) is list or type(prepend) is list: + g["dynamic"] = {"class_type": "RuntimeGenerator", "inputs": {}} + g["stateful"] = { + "class_type": "PromptManagerText", + "inputs": {"text": "looks literal"}, + } + self.cap = self.capture.begin( + RunStart(SESSION, SESSION, 1, JOB, EPOCH, 0), 3, g, None + ) + return self.cap + + def observe(self, event, executor=None, **data): + host.observe_cached_context( + executor or self.executor, event, {"prompt_id": JOB, **data} + ) + + def seal(self): + sealed = self.capture.seal(self.cap) + if sealed.lease: + self.addCleanup(sealed.lease.release) + return sealed.bundle.contexts[0] + + def test_original_return_exception_and_wrapper_cleanup_preserved(self): + cls = self.execution.PromptExecutor + original = cls.add_message + + def earlier(executor, *args, **kwargs): + self.calls.append("earlier") + return original(executor, *args, **kwargs) + + cls.add_message = earlier + seen = [] + stop_a = host.subscribe_executor_messages("a", lambda *args: seen.append(args)) + stop_b = host.subscribe_executor_messages("b", lambda *args: seen.append(args)) + wrapped = cls.add_message + + def later(executor, *args, **kwargs): + self.calls.append("later") + return wrapped(executor, *args, **kwargs) + + cls.add_message = later + data = {"prompt_id": JOB} + self.assertEqual( + self.executor.add_message("ok", data, broadcast=False), "original-return" + ) + self.assertEqual(len(seen), 2) + self.assertIs(seen[0][0], self.executor) + self.assertIs(seen[0][1][1], data) + with self.assertRaisesRegex(RuntimeError, "original"): + self.executor.add_message("raise", data, False) + self.assertEqual(len(seen), 2) + stop_a() + self.executor.add_message("ok", data, False) + self.assertEqual(len(seen), 3) + stop_b() + self.assertIs(cls.add_message, later) + self.executor.add_message("ok", data, False) + self.assertEqual(len(seen), 3) + stop_c = host.subscribe_executor_messages("c", lambda *args: seen.append(args)) + self.executor.add_message("ok", data, False) + self.assertEqual(len(seen), 4) + stop_c() + + def test_cache_reads_only_local_string_slots(self): + for mode in ("HierarchicalCache", "LRUCache", "RAMPressureCache"): + for kind, values in ( + ("PromptManager", [object(), ["cached"]]), + ("PromptManagerFocused", [object(), ["cached"]]), + ("PromptManagerSplit", [object(), ["cached"]]), + ("PromptManagerText", [["cached"]]), + ("PromptManagerPair", [object(), object(), ["cached"], [""]]), + ): + with self.subTest(mode=mode, kind=kind): + self.begin(kind) + cache = getattr(self.caching, mode)() + entry = Entry({"secret": object()}, values) + cache.entries["p"] = entry + self.executor.caches.outputs = cache + self.observe("execution_start") + self.observe("execution_cached", nodes=["p"]) + result = self.seal() + self.assertEqual(result.positive.effective.value, "cached") + self.assertEqual( + result.positive.effective.source, "host_cache_output" + ) + if kind == "PromptManagerPair": + self.assertEqual(result.negative.effective.value, "") + self.assertIs(cache.entries["p"], entry) + self.assertEqual(set(cache.__dict__), {"entries"}) + self.assertTrue(all(call[0] == "local" for call in self.calls)) + + def test_changed_trigger_configuration_keeps_cached_effective_text(self): + self.begin() + self.executor.caches.outputs.entries["p"] = Entry( + None, [object(), ["old injected text"]] + ) + self.observe("execution_start") + self.observe("execution_cached", nodes=["p"]) + result = self.seal().positive + self.assertEqual(result.effective.value, "old injected text") + self.assertEqual(result.raw.value, " raw ") + self.assertEqual(result.prepend.value, " pre ") + self.assertEqual(result.append.value, " post ") + self.assertEqual(result.raw.provenance, "inferred") + + def test_dynamic_cached_raw_is_unknown(self): + for dynamic in (["dynamic", 0], ["stateful", 0]): + for field in ("raw", "prepend"): + with self.subTest(dynamic=dynamic, field=field): + self.begin(**{field: dynamic}) + self.executor.caches.outputs.entries["p"] = Entry( + None, [object(), ["cached"]] + ) + self.observe("execution_start") + self.observe("execution_cached", nodes=["p"]) + result = self.seal().positive + self.assertEqual(result.effective.value, "cached") + self.assertIsNone(getattr(result, field).value) + self.assertEqual(result.append.value, " post ") + if field == "raw": + self.assertEqual(result.prepend.value, " pre ") + else: + self.assertEqual(result.raw.value, " raw ") + + def test_upstream_transform_preserves_final_pm_cached_effective_only(self): + for field, snapshot_field in ( + ("text", "raw"), + ("prepend_text", "prepend"), + ("append_text", "append"), + ): + for unresolved in (False, True): + with self.subTest(field=field, unresolved=unresolved): + g = transformed_input_graph(field, unresolved) + g["outer"] = { + "class_type": "StringConcatenate", + "inputs": {"string_a": ["p", 2], "string_b": " outer"}, + } + g["clip"] = { + "class_type": "CLIPTextEncode", + "inputs": {"text": ["outer", 0]}, + } + g["s"]["inputs"]["negative"] = ["clip", 0] + self.cap = self.capture.begin( + RunStart(SESSION, SESSION, 1, JOB, EPOCH, 0), + 3, + g, + None, + ) + self.executor.caches.outputs.entries["p"] = Entry( + None, [object(), object(), [" cached injected "], [""]] + ) + self.observe("execution_start") + self.observe("execution_cached", nodes=["upstream", "p"]) + result = self.seal() + self.assertIsNone(result.negative.effective.value) + side = result.positive + self.assertEqual(side.effective.value, " cached injected ") + self.assertEqual(side.effective.provenance, "recorded") + self.assertEqual(side.effective.source, "host_cache_output") + for name, expected in ( + ("raw", " cat "), + ("prepend", "pre"), + ("append", "post"), + ): + evidence = getattr(side, name) + if name == snapshot_field: + self.assertIsNone(evidence.value) + self.assertEqual(evidence.provenance, "unknown") + else: + self.assertEqual(evidence.value, expected) + self.assertEqual(evidence.provenance, "inferred") + + def test_provider_local_reset_eviction_and_unknown_modes_fail_closed(self): + for mode in ( + "NullCache", + "unknown", + "reset", + "evicted", + "malformed", + "provider-local", + ): + with self.subTest(mode=mode): + self.begin() + cache = self.caching.HierarchicalCache() + if mode == "NullCache": + cache = self.caching.NullCache() + elif mode == "unknown": + cache = types.SimpleNamespace( + get_local=lambda _: self.fail("unknown cache read") + ) + elif mode == "malformed": + cache.entries["p"] = Entry(None, [object(), [object()]]) + elif mode in ("reset", "provider-local"): + cache.entries["p"] = Entry( + None, [object(), ["resident provider output"]] + ) + self.executor.caches.outputs = cache + self.observe("execution_start") + if mode == "reset": + self.executor.caches.outputs = self.caching.HierarchicalCache() + self.observe("execution_cached", nodes=["p"]) + result = self.seal().positive + self.assertEqual( + result.effective.value, + "resident provider output" if mode == "provider-local" else None, + ) + + def test_worker_queue_job_executor_and_generation_fences(self): + for failure in ( + "unbound", + "thread", + "queue", + "job", + "executor", + "recursive", + "generation", + "no-start", + ): + with self.subTest(failure=failure): + self.begin() + self.executor.server.prompt_queue = self.queue + self.executor.caches.outputs.entries["p"] = Entry( + None, [object(), ["cached"]] + ) + if failure != "no-start": + self.observe("execution_start") + target = self.executor + if failure == "queue": + self.executor.server.prompt_queue = object() + if failure in ("executor", "recursive"): + target = self.execution.PromptExecutor() + target.server = self.executor.server + target.caches = self.executor.caches + if failure == "recursive": + self.observe("execution_start", target) + target = self.executor + data = { + "prompt_id": "wrong" if failure == "job" else JOB, + "nodes": ["p"], + } + if failure == "generation": + self.status = replace(self.status, capture_generation=4) + + def action(): + host.observe_cached_context(target, "execution_cached", data) + + if failure == "unbound": + Context().run(action) + elif failure == "thread": + thread = threading.Thread(target=action) + thread.start() + thread.join(2) + else: + action() + if failure == "generation": + self.assertIsNone(self.capture.seal(self.cap).bundle) + self.status = replace(self.status, capture_generation=3) + else: + self.assertIsNone(self.seal().positive.effective.value) + + def test_missing_cache_preserves_known_affixes_when_raw_is_dynamic(self): + self.begin(raw=["dynamic", 0]) + self.observe("execution_start") + self.observe("execution_cached", nodes=["p"]) + result = self.seal().positive + self.assertIsNone(result.raw.value) + self.assertIsNone(result.effective.value) + self.assertEqual(result.prepend.value, " pre ") + self.assertEqual(result.append.value, " post ") + + def test_cache_entries_tensors_and_executor_are_not_retained(self): + class Tensor: + pass + + tensor = Tensor() + tensor_ref = weakref.ref(tensor) + self.begin() + self.executor.caches.outputs.entries["p"] = Entry(tensor, [tensor, ["cached"]]) + self.observe("execution_start") + self.observe("execution_cached", nodes=["p"]) + self.executor.caches.outputs.entries.clear() + executor_ref = weakref.ref(self.executor) + del tensor + del self.executor + gc.collect() + self.assertIsNone(tensor_ref()) + self.assertIsNone(executor_ref()) + self.assertEqual(self.seal().positive.effective.value, "cached") + + def test_unknown_and_malformed_cache_payloads_are_bounded(self): + for values in (["x"] * 65, ["x" * 65537], [object()], ["x", "x"]): + with self.subTest(count=len(values)): + self.begin() + self.executor.caches.outputs.entries["p"] = Entry( + None, [object(), values] + ) + self.observe("execution_start") + self.observe("execution_cached", nodes=["p"]) + result = self.seal().positive + self.assertIsNone(result.effective.value) + self.assertEqual( + result.append.value, " post " if len(values) != 2 else None + ) + + def test_disposed_subscriptions_do_not_retain_dynamic_executor_classes(self): + def install_and_remove(): + class TemporaryExecutor: + def add_message(self, *args, **kwargs): + return None + + self.execution.PromptExecutor = TemporaryExecutor + ref = weakref.ref(TemporaryExecutor) + stop = host.subscribe_executor_messages("temporary", lambda *args: None) + stop() + del self.execution.PromptExecutor + return ref + + ref = install_and_remove() + gc.collect() + self.assertIsNone(ref()) diff --git a/tests/test_execution_context_schema.py b/tests/test_execution_context_schema.py new file mode 100644 index 0000000..de31eec --- /dev/null +++ b/tests/test_execution_context_schema.py @@ -0,0 +1,433 @@ +"""Versioned migration and hostile context restore boundaries.""" + +import json +import sqlite3 +import unittest +from unittest.mock import patch + +from database import execution_history_schema as history +from utils.execution_contexts import content_digest + +EPOCH = "11111111-1111-4111-8111-111111111111" +RUN = "22222222-2222-4222-8222-222222222222" + + +def legacy(conn): + for table, columns in history._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 + ] + checks = ( + ( + "singleton_id = 1", + "schema_version = 1", + history._uuid_check("storage_epoch"), + ) + if table == "execution_history_metadata" + else history._RUN_CHECKS + ) + definitions.extend(f"CHECK({check})" for check in checks) + conn.execute(f"CREATE TABLE {table} ({', '.join(definitions)})") + for name, (unique, columns) in history._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,1,?)", (EPOCH,)) + conn.execute( + "INSERT INTO execution_runs " + "VALUES(?,?,?,?,?,1,100,NULL,NULL,NULL,'observed',NULL,0,NULL,NULL" + ",NULL)", + (RUN, EPOCH, 0, RUN, EPOCH), + ) + conn.execute( + "CREATE TABLE prompts(id INTEGER PRIMARY KEY, text TEXT, run_count INTEGER)" + ) + conn.execute("INSERT INTO prompts VALUES(7,'kept',19)") + conn.commit() + + +class ContextSchemaTests(unittest.TestCase): + def connection(self): + conn = sqlite3.connect(":memory:") + self.addCleanup(conn.close) + return conn + + def test_v1_upgrade_preserves_epoch_and_records(self): + conn = self.connection() + legacy(conn) + before = conn.execute("SELECT * FROM execution_runs").fetchall() + self.assertIsNone(history.history_schema_error(conn)) + history.migrate_execution_history(conn) + self.assertEqual( + conn.execute( + "SELECT schema_version,storage_epoch FROM execution_history_metadata" + ).fetchone(), + (2, EPOCH), + ) + self.assertEqual( + conn.execute("SELECT * FROM execution_runs").fetchall(), before + ) + self.assertEqual( + conn.execute("SELECT * FROM prompts").fetchall(), [(7, "kept", 19)] + ) + for table in ( + "execution_context_texts", + "execution_prompt_snapshots", + "execution_contexts", + "execution_run_contexts", + "execution_context_capture", + ): + self.assertEqual( + conn.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0], 0 + ) + + def test_migration_idempotent(self): + conn = self.connection() + history.migrate_execution_history(conn) + first = "\n".join(conn.iterdump()) + history.migrate_execution_history(conn) + self.assertEqual("\n".join(conn.iterdump()), first) + self.assertEqual( + conn.execute( + "SELECT schema_version FROM execution_history_metadata" + ).fetchone()[0], + 2, + ) + + def test_migration_failure_rolls_back_metadata_and_all_tables(self): + conn = self.connection() + legacy(conn) + before = "\n".join(conn.iterdump()) + + def fail(connection): + connection.execute("CREATE TABLE execution_context_texts(x TEXT)") + raise sqlite3.OperationalError("injected context DDL failure") + + with patch( + "database.execution_history_schema.install_context_schema", + fail, + create=True, + ): + with self.assertRaisesRegex(sqlite3.OperationalError, "injected"): + history.migrate_execution_history(conn) + self.assertEqual("\n".join(conn.iterdump()), before) + + def test_partial_context_schema_rejected_even_for_v1_or_absent(self): + for old in (False, True): + conn = self.connection() + if old: + legacy(conn) + conn.execute("CREATE TABLE execution_context_texts(text_id TEXT)") + self.assertIsNotNone(history.history_schema_error(conn)) + with self.assertRaises(sqlite3.DatabaseError): + history.migrate_execution_history(conn) + + def populate(self, conn): + legacy(conn) + history.migrate_execution_history(conn) + text = content_digest("text", ("hello",)) + snapshot = content_digest("snapshot", ("hello", None, None, None)) + context = content_digest("context", (snapshot, None)) + conn.execute("INSERT INTO execution_context_texts VALUES(?,?)", (text, "hello")) + conn.execute( + "INSERT INTO execution_prompt_snapshots VALUES(?,?,NULL,NULL,NULL)", + (snapshot, text), + ) + conn.execute( + "INSERT INTO execution_contexts VALUES(?,?,NULL)", (context, snapshot) + ) + unknown = dict(provenance="unknown", source="unknown", reason="missing") + evidence = dict( + format_version=1, + role_provenance="inferred", + positive=dict( + raw=dict(provenance="recorded", source="runtime_return", reason=None), + effective=unknown, + prepend=unknown, + append=unknown, + ), + negative=None, + ) + conn.execute( + "INSERT INTO execution_run_contexts VALUES(?,?,?,?,NULL,NULL,NULL,NULL,1)", + (RUN, "node", context, json.dumps(evidence)), + ) + conn.execute( + "INSERT INTO execution_context_capture VALUES(?,'partial','missing',1)", + (RUN,), + ) + return evidence + + def test_owned_identifier_constraints_reject_nul_suffixed_bytes(self): + columns = { + "execution_context_texts": ("text_id",), + "execution_prompt_snapshots": ( + "snapshot_id", + "raw_text_id", + "effective_text_id", + "prepend_text_id", + "append_text_id", + ), + "execution_contexts": ( + "context_id", + "positive_snapshot_id", + "negative_snapshot_id", + ), + "execution_run_contexts": ( + "run_id", + "anchor_id", + "context_id", + "reference_epoch", + ), + "execution_context_capture": ("run_id",), + } + for table, names in columns.items(): + for column in names: + with self.subTest(table=table, column=column): + conn = self.connection() + self.populate(conn) + prefix = ( + RUN + if column in ("run_id", "reference_epoch") + else "node" if column == "anchor_id" else "a" * 64 + ) + with self.assertRaises(sqlite3.IntegrityError): + conn.execute( + f"UPDATE {table} SET {column}=?", + (prefix + "\0" + "x" * 1024,), + ) + + def test_context_validation_rejects_oversized_keys_before_fetch(self): + for table, column, prefix in ( + ("execution_context_texts", "text_id", "a" * 64), + ("execution_run_contexts", "anchor_id", "node"), + ): + with self.subTest(table=table, column=column): + conn = self.connection() + self.populate(conn) + malicious = prefix + "\0" + "x" * 1024 + if table == "execution_context_texts": + for dependent in ( + "execution_context_capture", + "execution_run_contexts", + "execution_contexts", + "execution_prompt_snapshots", + ): + conn.execute(f"DELETE FROM {dependent}") + conn.execute("PRAGMA ignore_check_constraints=ON") + conn.execute(f"UPDATE {table} SET {column}=?", (malicious,)) + + def bounded_text(value): + self.assertNotEqual(value, malicious.encode()) + return value.decode() + + conn.text_factory = bounded_text + self.assertIsNotNone(history.history_schema_error(conn)) + + def test_unicode_anchor_boundary_and_prompt_nuls_preserve_exact_text(self): + conn = self.connection() + self.populate(conn) + text = "hello\0world" + text_id = content_digest("text", (text,)) + snapshot_id = content_digest("snapshot", (text, None, None, None)) + context_id = content_digest("context", (snapshot_id, None)) + conn.execute( + "UPDATE execution_context_texts SET text_id=?,text_value=?", + (text_id, text), + ) + conn.execute( + "UPDATE execution_prompt_snapshots SET snapshot_id=?,raw_text_id=?", + (snapshot_id, text_id), + ) + conn.execute( + "UPDATE execution_contexts SET context_id=?,positive_snapshot_id=?", + (context_id, snapshot_id), + ) + conn.execute( + "UPDATE execution_run_contexts SET context_id=?,anchor_id=?", + (context_id, "😀" * 128), + ) + self.assertIsNone(history.history_schema_error(conn)) + self.assertEqual( + conn.execute("SELECT text_value FROM execution_context_texts").fetchone(), + (text,), + ) + + def test_link_without_capture_header_is_malformed_v2(self): + from database.execution_context_schema import context_schema_error + + conn = self.connection() + self.populate(conn) + conn.execute("DELETE FROM execution_context_capture") + before = "\n".join(conn.iterdump()) + self.assertIsNotNone(context_schema_error(conn)) + self.assertIsNotNone(history.history_schema_error(conn, allow_missing=False)) + self.assertEqual("\n".join(conn.iterdump()), before) + + def test_valid_context_and_fingerprints(self): + conn = self.connection() + self.populate(conn) + self.assertIsNone(history.history_schema_error(conn, allow_missing=False)) + + def test_v2_rejects_partial_schema_extra_uniqueness_bad_fingerprint_or_orphan(self): + mutations = ( + "DROP TABLE execution_context_capture", + "CREATE UNIQUE INDEX extra_text ON execution_context_texts(text_value)", + "UPDATE execution_context_texts SET text_value='changed'", + "UPDATE execution_prompt_snapshots SET raw_text_id='" + "0" * 64 + "'", + "UPDATE execution_contexts SET context_id='" + "0" * 64 + "'", + "UPDATE execution_run_contexts SET " + "run_id='33333333-3333-4333-8333-333333333333'", + "UPDATE execution_run_contexts SET evidence_json='{}'", + "UPDATE execution_context_capture SET " + "run_id='33333333-3333-4333-8333-333333333333'", + ) + for mutation in mutations: + with self.subTest(mutation=mutation): + conn = self.connection() + self.populate(conn) + conn.execute(mutation) + self.assertIsNotNone(history.history_schema_error(conn)) + + def test_evidence_rejects_extra_metadata_and_false_unknown(self): + for change in ("extra", "known", "role", "version"): + conn = self.connection() + evidence = self.populate(conn) + if change == "extra": + evidence["graph"] = {} + elif change == "known": + evidence["positive"]["effective"] = dict( + provenance="inferred", source="graph_literal", reason=None + ) + elif change == "role": + evidence["negative"] = evidence["positive"] + else: + evidence["format_version"] = True + conn.execute( + "UPDATE execution_run_contexts SET evidence_json=?", + (json.dumps(evidence),), + ) + self.assertIsNotNone(history.history_schema_error(conn)) + + def test_more_than_sixteen_context_links_rejected(self): + conn = self.connection() + self.populate(conn) + for i in range(16): + conn.execute( + "INSERT INTO execution_run_contexts SELECT " + "run_id,?,context_id,evidence_json,positive_prompt_id,negative_pro" + "mpt_id,pair_id,reference_epoch,capture_version FROM " + "execution_run_contexts LIMIT 1", + (str(i),), + ) + self.assertIsNotNone(history.history_schema_error(conn)) + + def test_owned_schema_rejects_triggers_views_and_generated_columns(self): + for mutation in ( + "CREATE TRIGGER hidden_write AFTER INSERT ON " + "execution_context_texts BEGIN DELETE FROM " + "execution_context_texts; END", + "CREATE VIEW context_view AS SELECT * FROM execution_context_texts", + "ALTER TABLE execution_context_texts ADD COLUMN hidden TEXT " + "GENERATED ALWAYS AS (text_value) VIRTUAL", + ): + with self.subTest(mutation=mutation): + conn = self.connection() + self.populate(conn) + conn.execute(mutation) + self.assertIsNotNone(history.history_schema_error(conn)) + + def test_duplicate_json_keys_and_unbounded_or_false_complete_rows_rejected(self): + for change in ( + "duplicate", + "oversized", + "reason", + "complete", + "empty_complete", + ): + with self.subTest(change=change): + conn = self.connection() + self.populate(conn) + conn.execute("PRAGMA ignore_check_constraints=ON") + if change == "duplicate": + conn.execute( + "UPDATE execution_run_contexts SET " + "evidence_json=replace(evidence_json, '{', " + "'{\"format_version\":1,')" + ) + elif change == "oversized": + conn.execute( + "UPDATE execution_context_texts SET text_value=?", + ("é" * 32769,), + ) + elif change == "reason": + conn.execute( + "UPDATE execution_context_capture SET reason='arbitrary' " + ) + else: + conn.execute( + "UPDATE execution_context_capture SET state='complete'" + ) + if change == "empty_complete": + conn.execute("DELETE FROM execution_run_contexts") + self.assertIsNotNone(history.history_schema_error(conn)) + + def test_restore_enforces_context_and_run_aggregate_payload_bounds(self): + for size, links, valid in ( + (40000, 0, False), + (40000, 1, False), + (25000, 5, True), + (25000, 6, False), + ): + with self.subTest(size=size, links=links): + conn = self.connection() + evidence = self.populate(conn) + text = "x" * size + text_id = content_digest("text", (text,)) + snapshot_id = content_digest("snapshot", (text,) * 4) + context_id = content_digest("context", (snapshot_id, snapshot_id)) + conn.execute( + "UPDATE execution_context_texts SET text_id=?,text_value=?", + (text_id, text), + ) + conn.execute( + "UPDATE execution_prompt_snapshots SET " + "snapshot_id=?,raw_text_id=?,effective_text_id=?,prepend_text_id=?" + ",append_text_id=?", + (snapshot_id,) + (text_id,) * 4, + ) + conn.execute( + "UPDATE execution_contexts SET " + "context_id=?,positive_snapshot_id=?,negative_snapshot_id=?", + (context_id, snapshot_id, snapshot_id), + ) + known = dict( + provenance="recorded", source="runtime_return", reason=None + ) + evidence["positive"] = dict.fromkeys( + ("raw", "effective", "prepend", "append"), known + ) + evidence["negative"] = evidence["positive"] + conn.execute( + "UPDATE execution_run_contexts SET context_id=?,evidence_json=?", + (context_id, json.dumps(evidence)), + ) + for i in range(links - 1): + conn.execute( + "INSERT INTO execution_run_contexts SELECT " + "run_id,?,context_id,evidence_json,positive_prompt_id," + "negative_prompt_id,pair_id,reference_epoch," + "capture_version FROM " + "execution_run_contexts LIMIT 1", + (str(i),), + ) + if not links: + conn.execute("DELETE FROM execution_run_contexts") + error = history.history_schema_error(conn) + if valid: + self.assertIsNone(error) + else: + self.assertIsNotNone(error) diff --git a/tests/test_execution_contexts.py b/tests/test_execution_contexts.py new file mode 100644 index 0000000..ea55022 --- /dev/null +++ b/tests/test_execution_contexts.py @@ -0,0 +1,218 @@ +"""Exact identities and bounded scalar records; normalization would lose history.""" + +import hashlib +import importlib.util +import sys +import types +import unittest +from pathlib import Path + + +def load_module(name): + package = "execution_context_test_utils" + if package not in sys.modules: + module = types.ModuleType(package) + module.__path__ = [str(Path(__file__).parents[1] / "utils")] + sys.modules[package] = module + qualified = package + "." + name + if qualified not in sys.modules: + spec = importlib.util.spec_from_file_location( + qualified, Path(__file__).parents[1] / "utils" / (name + ".py") + ) + module = importlib.util.module_from_spec(spec) + sys.modules[qualified] = module + spec.loader.exec_module(module) + return sys.modules[qualified] + + +class ExactContextsTests(unittest.TestCase): + def setUp(self): + self.m = load_module("execution_contexts") + + def evidence(self, value, provenance="inferred"): + return self.m.TextEvidence( + value, + "unknown" if value is None else provenance, + ( + "unknown" + if value is None + else "runtime_return" if provenance == "recorded" else "graph_literal" + ), + "unavailable" if value is None else None, + ) + + def snapshot(self, value): + return self.m.PromptSnapshot( + *(self.evidence(v) for v in (value, value, "", "")) + ) + + def test_exact_identity(self): + m = self.m + self.assertEqual( + m.canonical_content("snapshot", ("é", None, "", "")), + '[1,"snapshot","é",null,"",""]'.encode(), + ) + values = ("cat", "Cat", " cat", "cat ", "é", "e\u0301", None, "") + digests = {m.content_digest("snapshot", (v, v, "", "")) for v in values} + self.assertEqual(len(digests), len(values)) + self.assertEqual( + m.content_digest("snapshot", ("cat", "cat", "", "")), + hashlib.sha256(b'[1,"snapshot","cat","cat","",""]').hexdigest(), + ) + a = self.snapshot("cat") + b = m.PromptSnapshot( + *(self.evidence(v, "recorded") for v in ("cat", "cat", "", "")) + ) + self.assertEqual(m.snapshot_digest(a), m.snapshot_digest(b)) + self.assertNotEqual( + m.context_digest(m.RecipeContext("1", a, None, "inferred")), + m.context_digest(m.RecipeContext("1", None, a, "inferred")), + ) + + def test_invalid_text_and_bounds(self): + m = self.m + for value in ("\ud800", "x" * 65537, "é" * 32769, object(), True): + with self.subTest(value_type=type(value).__name__): + with self.assertRaises(ValueError): + m.canonical_content("snapshot", (value,)) + m.canonical_content("snapshot", ("é" * 32768,)) + for identifier in ("x" * 129, object()): + with self.assertRaises(ValueError): + m.validate_bundle( + m.ContextBundle( + (m.RecipeContext(identifier, None, None, "inferred"),), + "unavailable", + "missing", + ) + ) + with self.assertRaises(ValueError): + m.validate_bundle( + m.ContextBundle( + tuple( + m.RecipeContext(str(i), None, None, "inferred") + for i in range(17) + ), + "partial", + "limit", + ) + ) + with self.assertRaises(ValueError): + m.validate_bundle( + m.ContextBundle( + (m.RecipeContext("1", self.snapshot("\ud800"), None, "inferred"),), + "complete", + "", + ) + ) + + def test_bundle_preserves_unknown_and_empty(self): + m = self.m + bundle = m.ContextBundle( + (m.RecipeContext("1", self.snapshot(""), self.snapshot(None), "inferred"),), + "partial", + "unknown", + ) + self.assertIsNone(m.validate_bundle(bundle)) + with self.assertRaises(ValueError): + m.validate_bundle( + m.ContextBundle(bundle.contexts, "partial", "", format_version=2) + ) + + def test_evidence_contract_rejects_incompatible_metadata(self): + m = self.m + invalid = ( + ("", "graph-inferred", "graph_literal", None), + ("", "recorded", "graph_literal", None), + ("", "inferred", "runtime_return", None), + (None, "inferred", "graph_literal", None), + ("", "unknown", "unknown", "unavailable"), + (None, "unknown", "unknown", None), + (None, "unknown", "unknown", "an arbitrary exception message"), + ) + for fields in invalid: + with self.subTest(fields=fields): + evidence = m.TextEvidence(*fields) + snapshot = m.PromptSnapshot(evidence, evidence, evidence, evidence) + with self.assertRaises(ValueError): + m.validate_bundle( + m.ContextBundle( + (m.RecipeContext("1", snapshot, None, "inferred"),), + "partial", + "", + ) + ) + for provenance, source in ( + ("recorded", "runtime_return"), + ("recorded", "host_cache_output"), + ("inferred", "graph_literal"), + ("inferred", "graph_semantics"), + ): + evidence = m.TextEvidence("", provenance, source, None) + snapshot = m.PromptSnapshot(evidence, evidence, evidence, evidence) + m.validate_bundle( + m.ContextBundle( + (m.RecipeContext("1", snapshot, None, "inferred"),), "complete", "" + ) + ) + + def test_bundle_state_and_declared_role_contract(self): + m = self.m + for state, role in ( + ("unknown", "inferred"), + ("complete", "recorded"), + ("complete", "declared"), + ): + with self.subTest(state=state, role=role): + with self.assertRaises(ValueError): + m.validate_bundle( + m.ContextBundle( + (m.RecipeContext("1", None, None, role),), state, "" + ) + ) + for state in ("complete", "partial", "ambiguous", "unavailable", "conflicting"): + m.validate_bundle( + m.ContextBundle( + (m.RecipeContext("1", None, None, "inferred"),), state, "" + ) + ) + + def test_identifier_limit_counts_characters(self): + m = self.m + m.validate_bundle( + m.ContextBundle( + (m.RecipeContext("é" * 128, None, None, "inferred"),), "unavailable", "" + ) + ) + with self.assertRaises(ValueError): + m.validate_bundle( + m.ContextBundle( + (m.RecipeContext("é" * 129, None, None, "inferred"),), + "unavailable", + "", + ) + ) + + def test_aggregate_payload_limits(self): + m = self.m + large = self.snapshot("x" * 65536) + with self.assertRaises(ValueError): + m.validate_bundle( + m.ContextBundle( + (m.RecipeContext("1", large, large, "inferred"),), "complete", "" + ) + ) + contexts = tuple( + m.RecipeContext(str(i), large, None, "inferred") for i in range(8) + ) + with self.assertRaises(ValueError): + m.validate_bundle(m.ContextBundle(contexts, "complete", "")) + + def test_equal_repeated_samples_are_ambiguous(self): + # Projection contains references, never runtime samples. Duplicate anchors + # remain separate invalid records even if their prompt values are equal. + m = self.m + context = m.RecipeContext( + "1", self.snapshot("same"), self.snapshot("same"), "inferred" + ) + with self.assertRaises(ValueError): + m.validate_bundle(m.ContextBundle((context, context), "complete", "")) diff --git a/tests/test_execution_history_database.py b/tests/test_execution_history_database.py index 4c9f37e..e06cfcb 100644 --- a/tests/test_execution_history_database.py +++ b/tests/test_execution_history_database.py @@ -95,6 +95,10 @@ class HistoryDatabaseTests(unittest.TestCase): for pre_pair in (False, True): with self.subTest(pre_pair=pre_pair): with conn: + from database.execution_context_schema import TABLE_COLUMNS + + for table in reversed(TABLE_COLUMNS): + conn.execute(f"DROP TABLE {table}") conn.execute("DROP TABLE execution_runs") conn.execute("DROP TABLE execution_history_metadata") if pre_pair: @@ -232,7 +236,7 @@ class HistoryDatabaseTests(unittest.TestCase): else: conn.execute( "INSERT INTO execution_history_metadata " - "VALUES(1,1,?)", + "VALUES(1,2,?)", (epoch,), ) else: @@ -432,7 +436,7 @@ class HistoryDatabaseTests(unittest.TestCase): "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", + "UPDATE execution_history_metadata SET schema_version=3", "PRAGMA writable_schema=ON; UPDATE sqlite_master SET " "sql=replace(sql, \"'observed'\", \"'OBSERVED'\") " "WHERE name='execution_runs'; PRAGMA writable_schema=OFF", diff --git a/tests/test_execution_history_hooks.py b/tests/test_execution_history_hooks.py index f54b56f..840735e 100644 --- a/tests/test_execution_history_hooks.py +++ b/tests/test_execution_history_hooks.py @@ -76,6 +76,8 @@ class HistoryHooksTests(unittest.TestCase): logger.handlers, logger.propagate, logger.level = previous def cleanup_registrations(self): + for registration in self.registrations: + registration.writer.invalidate() for release in self.releases: release.set() for registration in self.registrations: @@ -356,7 +358,7 @@ class HistoryHooksTests(unittest.TestCase): clock = ManualClock() with patch( "utils.execution_history.HistoryWriter", - lambda factory: HistoryWriter(factory, clock=clock), + lambda factory, **kwargs: HistoryWriter(factory, clock=clock, **kwargs), ): registration = self.register() for index in range(128): @@ -448,3 +450,310 @@ class HistoryHooksTests(unittest.TestCase): self.drain(registration) self.assertEqual(seen, [JOB]) self.assertEqual(len(self.rows()), 1) + + def test_fixture_cleanup_discards_saturation_before_release_and_database_close( + self, + ): + from utils.execution_history_events import make_start + + entered, release, rescue = (threading.Event() for _ in range(3)) + self.releases.append(release) + pending_at_release, writes = [], [] + original = self.db.write_execution_start + + def blocked(command): + entered.set() + release.wait(5) + pending = registration.writer.status().pending + pending_at_release.append(pending) + # Model a slow drain without executing 1,024 SQLite transactions. + if pending > 1: + rescue.wait(5) + writes.append(command.run_id) + return original(command) + + with patch.object(self.db, "write_execution_start", blocked): + registration = self.register() + try: + self.dequeue(1) + self.assertTrue(entered.wait(1)) + epoch = registration.writer.status().epoch + for i in range(1023): + command = make_start( + registration._session_id, item(1000 + i), epoch, 0 + ) + self.assertTrue(registration.writer.submit(command)) + self.assertEqual(registration.writer.status().pending, 1024) + before = time.monotonic() + self.cleanup_registrations() + self.assertLess(time.monotonic() - before, 2) + self.assertEqual(pending_at_release, [1]) + self.assertEqual(len(writes), 1) + self.assertFalse(registration.writer.status().alive) + self.assertEqual(registration.writer.status().pending, 0) + self.db.model.close_all() + finally: + registration.writer.invalidate() + rescue.set() + release.set() + self.assertTrue( + registration.dispose(3), "regression rescue leaked worker" + ) + + def test_rejected_new_start_revokes_previous_worker_capability(self): + from utils.execution_context_capture import active_capture + from utils.execution_history_events import make_start + + 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(1) + self.assertTrue(entered.wait(1)) + previous = active_capture() + self.assertIsNotNone(previous) + epoch = registration.writer.status().epoch + for i in range(1023): + command = make_start(registration._session_id, item(1000 + i), epoch, 0) + self.assertTrue(registration.writer.submit(command)) + self.assertEqual(registration.writer.status().pending, 1024) + self.dequeue(2) + self.assertIsNone(active_capture()) + self.assertFalse(previous[0].valid(previous[1])) + registration.writer.invalidate() + release.set() + self.assertTrue(registration.dispose(3), "owned worker leaked") + + def test_late_finish_does_not_clear_new_attempt_binding(self): + from utils.execution_context_capture import active_capture + + registration = self.register() + self.dequeue(1) + self.dequeue(2) + current = active_capture() + self.assertIsNotNone(current) + self.queue.task_done(1, {}, status()) + self.assertEqual(active_capture(), current) + # An identical callback replay keeps the already accepted binding. + registration._dequeue(item(2)) + self.assertEqual(active_capture(), current) + self.queue.task_done(2, {}, status()) + self.assertIsNone(active_capture()) + + def test_invalidation_releases_capture_immediately_outside_writer_locks(self): + from utils.execution_context_capture import active_capture + + registration = self.register() + self.dequeue(1) + current = active_capture() + pool = registration.writer.context_budget + self.assertGreater(pool.charged_bytes, 0) + registration.writer.invalidate() + self.assertEqual(pool.charged_bytes, 0) + self.assertFalse(current[0].valid(current[1])) + + def test_replayed_older_accepted_dequeue_cannot_revoke_current_binding(self): + from utils.execution_context_capture import active_capture + + registration = self.register() + old = self.dequeue(1) + self.dequeue(2) + current = active_capture() + registration._dequeue(old) + self.assertEqual(active_capture(), current) + + def test_broken_finish_releases_active_capture(self): + from utils.execution_context_capture import active_capture + + registration = self.register() + self.dequeue(1) + + class Broken: + @property + def messages(self): + raise RuntimeError("broken status") + + self.queue.task_done(1, {}, Broken()) + self.assertEqual(registration.writer.context_budget.charged_bytes, 0) + self.assertIsNone(active_capture()) + + def test_finished_accepted_replay_preserves_new_same_uuid_attempt(self): + from utils.execution_context_capture import active_capture + + registration = self.register() + old = self.dequeue(1) + self.queue.task_done(1, {}, status()) + self.dequeue(2) + current = active_capture() + registration._dequeue(old) + self.assertEqual(active_capture(), current) + self.queue.task_done(2, {}, status()) + self.drain(registration) + self.assertEqual(len(self.rows()), 2) + + def test_observed_upstream_survives_downstream_failure_in_database(self): + from utils.execution_context_capture import record_runtime_context + + g = { + "p": { + "class_type": "PromptManagerPair", + "inputs": {"text": "raw", "negative_text": ""}, + }, + "s": { + "class_type": "KSampler", + "inputs": {"positive": ["p", 0], "negative": ["p", 1]}, + }, + } + registration = self.register() + self.queue.results.append(((0, JOB, g, {}, ["s"]), 1)) + self.queue.get() + host = types.ModuleType("comfy_execution.utils") + host.get_executing_context = lambda: SimpleNamespace( + prompt_id=JOB, node_id="p", list_index=0 + ) + with patch.dict(sys.modules, {"comfy_execution.utils": host}): + record_runtime_context( + "PromptManagerPair", + " raw ", + " pre ", + "", + "pre raw TRIGGER", + role="positive", + ) + record_runtime_context("PromptManagerPair", "", "", "", "", role="negative") + self.queue.task_done(1, {}, status("execution_error")) + self.drain(registration) + row = self.rows()[0] + self.assertEqual(row["outcome"], "failed") + result = self.db.get_execution_contexts(row["run_id"]) + self.assertEqual(result["state"], "complete") + self.assertEqual( + result["contexts"][0]["positive"], + { + "raw": " raw ", + "effective": "pre raw TRIGGER", + "prepend": " pre ", + "append": "", + }, + ) + self.assertEqual(result["contexts"][0]["negative"]["effective"], "") + + def test_close_releases_capture_while_writer_is_blocked(self): + entered, release = threading.Event(), threading.Event() + self.releases.append(release) + original = self.db.write_execution_start + + def blocked(command): + entered.set() + release.wait(3) + return original(command) + + with patch.object(self.db, "write_execution_start", blocked): + registration = self.register() + self.dequeue(1) + self.assertTrue(entered.wait(1)) + self.assertGreater(registration.writer.context_budget.charged_bytes, 0) + self.assertFalse(registration.writer.close(0)) + self.assertEqual(registration.writer.context_budget.charged_bytes, 0) + release.set() + + def test_reregistration_cannot_rebind_already_observed_queue_item(self): + from utils.execution_context_capture import active_capture + + first = self.register() + self.dequeue(10) + self.queue.task_done(10, {}, status()) + self.drain(first) + self.assertTrue(first.dispose()) + second = self.register() + self.dequeue(10) + self.assertIsNone(active_capture()) + self.queue.task_done(10, {}, status()) + self.drain(second) + self.assertEqual(len(self.rows()), 1) + + def test_cached_executor_subscription_persists_strings_and_disposes_cooperatively( + self, + ): + from utils.execution_context_host import subscribe_executor_messages + + calls = [] + execution = types.ModuleType("execution") + caching = types.ModuleType("comfy_execution.caching") + package = types.ModuleType("comfy_execution") + entry_type = namedtuple("CacheEntry", "ui outputs") + + class PromptExecutor: + def add_message(self, event, data, broadcast): + calls.append(event) + return "host-result" + + class HierarchicalCache: + def get_local(self, node_id): + return entry_type(None, [object(), object(), ["old trigger"], [""]]) + + execution.PromptExecutor, execution.CacheEntry = PromptExecutor, entry_type + caching.HierarchicalCache = HierarchicalCache + package.caching = caching + with patch.dict( + sys.modules, + { + "execution": execution, + "comfy_execution": package, + "comfy_execution.caching": caching, + }, + ): + observer_calls = [] + independent = subscribe_executor_messages( + "independent", lambda *args: observer_calls.append(args) + ) + try: + registration = self.register() + g = { + "p": { + "class_type": "PromptManagerPair", + "inputs": {"text": "raw", "negative_text": ""}, + }, + "s": { + "class_type": "KSampler", + "inputs": {"positive": ["p", 0], "negative": ["p", 1]}, + }, + } + self.queue.results.append(((0, JOB, g, {}, ["s"]), 1)) + self.queue.get() + executor = PromptExecutor() + executor.server = self.server + executor.caches = SimpleNamespace(outputs=HierarchicalCache()) + self.assertEqual( + executor.add_message("execution_start", {"prompt_id": JOB}, False), + "host-result", + ) + executor.add_message( + event="execution_cached", + data={"prompt_id": JOB, "nodes": ["p"]}, + broadcast=False, + ) + self.queue.task_done(1, {}, status()) + self.drain(registration) + result = self.db.get_execution_contexts(self.rows()[0]["run_id"]) + self.assertEqual( + result["contexts"][0]["positive"]["effective"], "old trigger" + ) + self.assertEqual(result["contexts"][0]["negative"]["effective"], "") + self.assertTrue(registration.dispose()) + executor.add_message( + "execution_cached", {"prompt_id": JOB, "nodes": ["p"]}, False + ) + self.assertEqual(len(observer_calls), 3) + finally: + independent() + self.assertEqual( + calls, ["execution_start", "execution_cached", "execution_cached"] + ) diff --git a/tests/test_execution_history_restore.py b/tests/test_execution_history_restore.py index 601287c..5b379e1 100644 --- a/tests/test_execution_history_restore.py +++ b/tests/test_execution_history_restore.py @@ -22,6 +22,91 @@ class HistoryRestoreTests(unittest.TestCase): finish = history_tests.HistoryDatabaseTests.finish count = history_tests.HistoryDatabaseTests.count + def test_legacy_identifier_and_diagnostic_bytes_rejected_before_fetch(self): + try: + from tests.test_execution_context_schema import legacy, RUN + except ImportError: + from test_execution_context_schema import legacy, RUN + from database.execution_history_schema import history_schema_error + + for table, column, prefix in ( + ("execution_history_metadata", "storage_epoch", RUN), + *( + ("execution_runs", column, RUN) + for column in ( + "run_id", + "session_id", + "job_id", + "storage_epoch", + ) + ), + ("execution_runs", "node_id", "x" * 128), + ("execution_runs", "node_type", "x" * 128), + ("execution_runs", "exception_type", "x" * 256), + ): + with self.subTest(table=table, column=column): + with closing(sqlite3.connect(":memory:")) as conn: + legacy(conn) + malicious = prefix + "\0" + "x" * 1024 + conn.execute(f"UPDATE {table} SET {column}=?", (malicious,)) + + def bounded_text(value): + self.assertNotEqual(value, malicious.encode()) + return value.decode() + + conn.text_factory = bounded_text + self.assertIsNotNone(history_schema_error(conn)) + + def test_legacy_metadata_version_payload_is_rejected_before_fetch(self): + try: + from tests.test_execution_context_schema import legacy + except ImportError: + from test_execution_context_schema import legacy + from database.execution_history_schema import history_schema_error + + for malicious in ("x" * 1024, b"x" * 1024): + with self.subTest(kind=type(malicious).__name__): + with closing(sqlite3.connect(":memory:")) as conn: + legacy(conn) + conn.execute("PRAGMA ignore_check_constraints=ON") + conn.execute( + "UPDATE execution_history_metadata SET schema_version=?", + (malicious,), + ) + + def bounded_row(cursor, row): + self.assertFalse(any(value == malicious for value in row)) + return row + + conn.row_factory = bounded_row + self.assertIsNotNone(history_schema_error(conn)) + + def test_legacy_four_byte_diagnostics_validate_and_migrate(self): + try: + from tests.test_execution_context_schema import legacy + except ImportError: + from test_execution_context_schema import legacy + from database.execution_history_schema import ( + history_schema_error, + migrate_execution_history, + ) + + with closing(sqlite3.connect(":memory:")) as conn: + legacy(conn) + conn.execute( + "UPDATE execution_runs SET node_id=?,node_type=?,exception_type=?", + ("😀" * 128, "😀" * 128, "😀" * 256), + ) + self.assertIsNone(history_schema_error(conn)) + migrate_execution_history(conn) + self.assertIsNone(history_schema_error(conn, allow_missing=False)) + self.assertEqual( + conn.execute( + "SELECT node_id,node_type,exception_type FROM " "execution_runs" + ).fetchone(), + ("😀" * 128, "😀" * 128, "😀" * 256), + ) + def test_committed_wal_backup_includes_history_and_epoch(self): db = self.database() conn = db.model.get_connection() @@ -87,14 +172,14 @@ class HistoryRestoreTests(unittest.TestCase): 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)" + "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')" + "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")) @@ -113,6 +198,113 @@ class HistoryRestoreTests(unittest.TestCase): (7, "old prompt"), ) + def test_wal_and_old_backup_staging_leave_source_unchanged(self): + from database.execution_context_schema import TABLE_COLUMNS + + source = str(Path(self.temp.name) / "v1.db") + old = self.database(source) + start = self.start(old) + old.write_execution_start(start) + with old.model.get_connection() as conn: + conn.execute("INSERT INTO prompts(id,text,run_count) VALUES(7,'old',19)") + for table in reversed(TABLE_COLUMNS): + conn.execute(f"DROP TABLE {table}") + sql = ( + conn.execute( + "SELECT sql FROM sqlite_master WHERE " + "name='execution_history_metadata'" + ) + .fetchone()[0] + .replace("schema_version = 2", "schema_version = 1") + ) + conn.execute("DROP TABLE execution_history_metadata") + conn.execute(sql) + conn.execute( + "INSERT INTO execution_history_metadata VALUES(1,1,?)", + (start.storage_epoch,), + ) + conn.execute("PRAGMA wal_autocheckpoint=0") + with conn: + conn.execute("UPDATE prompts SET run_count=31 WHERE id=7") + source_bytes = Path(source).read_bytes() + wal_bytes = Path(source + "-wal").read_bytes() + self.assertEqual(PromptModel.verify_database_file(source), (True, "ok")) + live = self.database() + live.model.restore_from_file(source) + self.assertEqual(Path(source).read_bytes(), source_bytes) + self.assertEqual(Path(source + "-wal").read_bytes(), wal_bytes) + self.assertEqual(live.get_execution_run(start.run_id)["record_version"], 1) + self.assertEqual( + live.model.get_connection() + .execute("SELECT run_count FROM prompts WHERE id=7") + .fetchone()[0], + 31, + ) + self.assertEqual( + live.model.get_connection() + .execute("SELECT COUNT(*) FROM execution_run_contexts") + .fetchone()[0], + 0, + ) + self.assertEqual( + conn.execute( + "SELECT schema_version FROM execution_history_metadata" + ).fetchone()[0], + 1, + ) + + def test_malformed_v2_context_restore_keeps_live_database_unchanged(self): + live = self.database() + start = self.start(live) + live.write_execution_start(start) + with live.model.get_connection() as conn: + conn.execute("INSERT INTO prompts(text,run_count) VALUES('kept',19)") + before = "\n".join(live.model.get_connection().iterdump()) + for mutation in ( + "DROP TABLE execution_context_capture", + "CREATE UNIQUE INDEX bad_context_unique ON " + "execution_context_texts(text_value)", + "INSERT INTO execution_context_texts VALUES('" + + "0" * 64 + + "','wrong fingerprint')", + ): + with self.subTest(mutation=mutation): + source = str(Path(self.temp.name) / "bad-context.db") + self.assertTrue(live.model.backup_database(source)) + with closing(sqlite3.connect(source)) as conn, conn: + conn.execute(mutation) + original = Path(source).read_bytes() + with self.assertRaises(ValueError): + live.model.restore_from_file(source) + self.assertEqual( + "\n".join(live.model.get_connection().iterdump()), before + ) + self.assertEqual(Path(source).read_bytes(), original) + + def test_missing_capture_header_rejects_backup_verification_and_restore(self): + try: + from tests.test_execution_context_database import bundle + except ImportError: + from test_execution_context_database import bundle + + live = self.database() + start = self.start(live) + live.write_execution_start(start) + live.write_execution_contexts(start.run_id, start.storage_epoch, bundle()) + self.assertEqual(live.get_execution_contexts(start.run_id)["state"], "complete") + before = "\n".join(live.model.get_connection().iterdump()) + source = str(Path(self.temp.name) / "missing-capture.db") + self.assertTrue(live.model.backup_database(source)) + with closing(sqlite3.connect(source)) as conn, conn: + conn.execute("DELETE FROM execution_context_capture") + original = Path(source).read_bytes() + valid, _ = PromptModel.verify_database_file(source) + self.assertFalse(valid) + with self.assertRaises(ValueError): + live.model.restore_from_file(source) + self.assertEqual("\n".join(live.model.get_connection().iterdump()), before) + self.assertEqual(Path(source).read_bytes(), original) + def test_failed_staging_rotation_keeps_live_epoch_and_never_notifies(self): from database.execution_history import subscribe_history_restore @@ -189,17 +381,17 @@ class HistoryRestoreTests(unittest.TestCase): ) conn.execute("DROP TABLE old_runs") conn.execute( - "CREATE UNIQUE INDEX idx_execution_runs_identity " - "ON execution_runs(session_id, queue_item_id)" + "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)" + "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)" + "CREATE UNIQUE INDEX unexpected_job_uniqueness ON " + "execution_runs(job_id)" ) original = Path(source).read_bytes() with self.assertRaisesRegex(ValueError, "uniqu"): @@ -241,6 +433,51 @@ class HistoryRestoreTests(unittest.TestCase): self.assertNotEqual(db.get_execution_epoch(), original_epoch) self.assertEqual(notices, [True]) + def test_restore_preserves_context_scalars_but_never_reepochs_hints(self): + try: + from tests.test_execution_context_database import bundle + except ImportError: + from test_execution_context_database import bundle + from database.execution_history import HistoryEpochChanged + + db = self.database() + with db.model.get_connection() as conn: + conn.execute( + "INSERT INTO prompts(id,text) VALUES(1,'positive'),(2,'negative')" + ) + conn.execute( + "INSERT INTO prompt_pairs(id,positive_prompt_id,negative_prompt_id) " + "VALUES(7,1,2)" + ) + start = self.start(db) + db.write_execution_start(start) + db.write_execution_contexts(start.run_id, start.storage_epoch, bundle()) + original = db.get_execution_contexts(start.run_id) + backup = str(Path(self.temp.name) / "contexts.db") + self.assertTrue(db.model.backup_database(backup)) + source_bytes = Path(backup).read_bytes() + for _ in range(2): + db.model.restore_from_file(backup) + read = db.get_execution_contexts(start.run_id) + expected = original["contexts"][0].copy() + expected.update( + positive_prompt_id=None, negative_prompt_id=None, pair_id=None + ) + self.assertEqual(read["contexts"], [expected]) + self.assertEqual(read["state"], "complete") + self.assertEqual( + read["contexts"][0]["reference_epoch"], start.storage_epoch + ) + with self.assertRaises(HistoryEpochChanged): + db.write_execution_contexts(start.run_id, start.storage_epoch, bundle()) + self.assertFalse( + db.write_execution_contexts( + start.run_id, db.get_execution_epoch(), bundle() + ) + ) + self.assertEqual(db.get_execution_contexts(start.run_id), read) + self.assertEqual(Path(backup).read_bytes(), source_bytes) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_execution_history_writer.py b/tests/test_execution_history_writer.py index 3ad2150..7c1dbb7 100644 --- a/tests/test_execution_history_writer.py +++ b/tests/test_execution_history_writer.py @@ -1,6 +1,7 @@ """Real-thread, disposable-database coverage for the bounded history writer.""" from dataclasses import FrozenInstanceError, replace +from contextlib import nullcontext import logging import os import sqlite3 @@ -105,6 +106,491 @@ class HistoryWriterTests(unittest.TestCase): self.releases.append(release) return entered, release + def sealed(self, writer, value="exact text", generation=None): + from utils.execution_context_budget import SealedContext + from utils.execution_contexts import ( + ContextBundle, + RecipeContext, + PromptSnapshot, + TextEvidence, + bundle_payload_size, + ) + + field = TextEvidence(value, "recorded", "runtime_return", None) + snapshot = PromptSnapshot(field, field, field, field) + bundle = ContextBundle( + (RecipeContext("1", snapshot, snapshot, "inferred"),), + "complete", + "captured", + ) + lease = writer.context_budget.reserve(bundle_payload_size(bundle)) + return SealedContext( + bundle, + lease, + "complete", + "captured", + writer.status().capture_generation if generation is None else generation, + ) + + def test_context_success_releases_and_legacy_finish_preserves_capture(self): + writer = self.writer() + start = self.start_command() + sealed = self.sealed(writer) + self.assertTrue(writer.submit(start)) + self.assertTrue(writer.submit_finalization(self.finish_command(start), sealed)) + self.wait_for(lambda: writer.status().pending == 0) + self.assertEqual(writer.context_budget.charged_bytes, 0) + self.assertEqual( + self.db.get_execution_contexts(start.run_id)["state"], "complete" + ) + self.assertTrue(writer.submit(self.finish_command(start))) + self.wait_for(lambda: writer.status().pending == 0) + self.assertEqual( + self.db.get_execution_contexts(start.run_id)["state"], "complete" + ) + + def test_context_saturation_keeps_lifecycle_and_status(self): + from utils.execution_context_budget import SealedContext + from utils.execution_history_events import encode_command + + writer = self.writer() + lease = writer.context_budget.reserve(8388608) + self.assertIsNone(writer.context_budget.reserve(1)) + start = self.start_command() + self.assertTrue(writer.submit(start)) + finish = self.finish_command(start) + self.assertLessEqual(len(encode_command(finish)), 2048) + self.assertTrue( + writer.submit_finalization( + finish, SealedContext(None, None, "unavailable", "budget_exhausted") + ) + ) + self.wait_for(lambda: writer.status().pending == 0) + self.assertEqual( + self.db.get_execution_run(start.run_id)["outcome"], "succeeded" + ) + self.assertEqual( + self.db.get_execution_contexts(start.run_id)["reason"], "budget_exhausted" + ) + lease.release() + self.assertEqual(writer.context_budget.charged_bytes, 0) + + def test_outcome_survives_context_transaction_failure(self): + for error, ready in ( + (ValueError("invalid evidence"), True), + (sqlite3.OperationalError("unavailable"), False), + ): + with self.subTest(error=type(error)): + writer = self.writer() + start = self.start_command(10 if ready else 11) + self.assertTrue(writer.submit(start)) + sealed = self.sealed(writer) + entered, release = self.block() + + def fail_context(*args, **kwargs): + entered.set() + release.wait(3) + raise error + + with patch.object( + self.db, "write_execution_contexts_for_finish", fail_context + ): + self.assertTrue( + writer.submit_finalization(self.finish_command(start), sealed) + ) + self.assertTrue(entered.wait(1)) + self.assertEqual( + self.db.get_execution_run(start.run_id)["outcome"], "succeeded" + ) + self.assertGreater(writer.context_budget.charged_bytes, 0) + release.set() + self.wait_for(lambda: writer.status().pending == 0) + self.assertEqual( + self.db.get_execution_run(start.run_id)["outcome"], "succeeded" + ) + self.assertEqual(writer.status().accepting, ready) + self.assertEqual(writer.context_budget.charged_bytes, 0) + if not ready: + self.clock.advance(5) + self.wait_for(lambda: writer.status().accepting) + self.assertEqual( + self.db.get_execution_contexts(start.run_id)["state"], + "not_captured", + ) + self.assertTrue(writer.close()) + + def test_same_epoch_recovery_changes_capture_generation(self): + writer = self.writer() + start = self.start_command() + old_generation = writer.status().capture_generation + sealed = self.sealed(writer) + writer.invalidate() + self.clock.advance(5) + self.wait_for(lambda: writer.status().accepting) + self.assertEqual(writer.status().epoch, start.storage_epoch) + self.assertGreater(writer.status().capture_generation, old_generation) + self.assertTrue(writer.submit(start)) + self.assertTrue(writer.submit_finalization(self.finish_command(start), sealed)) + self.wait_for(lambda: writer.status().pending == 0) + self.assertEqual( + self.db.get_execution_run(start.run_id)["outcome"], "succeeded" + ) + self.assertEqual(writer.context_budget.charged_bytes, 0) + self.assertEqual(self.db.get_execution_contexts(start.run_id)["contexts"], []) + + def test_restore_between_transactions_rejects_context(self): + writer = self.writer() + start = self.start_command() + self.db.write_execution_start(start) + backup = os.path.join(self.directory.name, "between.db") + self.assertTrue(self.db.model.backup_database(backup)) + original = self.db.write_execution_finish + + def finish(command): + result = original(command) + self.db.model.restore_from_file(backup) + return result + + with patch.object(self.db, "write_execution_finish", finish): + self.assertTrue( + writer.submit_finalization( + self.finish_command(start), self.sealed(writer) + ) + ) + self.wait_for(lambda: writer.status().pending == 0) + self.assertEqual( + self.db.get_execution_contexts(start.run_id)["state"], "not_captured" + ) + self.assertEqual(writer.context_budget.charged_bytes, 0) + self.assertFalse(writer.status().accepting) + + def test_reject_discard_and_bounded_close_release_leases(self): + writer = self.writer() + start = self.start_command() + sealed = self.sealed(writer) + self.assertFalse(writer.submit_finalization(object(), sealed)) + self.assertEqual(writer.context_budget.charged_bytes, 0) + 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(start)) + self.assertTrue(entered.wait(1)) + for _ in range(1023): + self.assertTrue( + writer.submit_finalization( + self.finish_command(start), self.sealed(writer) + ) + ) + self.assertEqual(writer.status().pending, 1024) + self.assertFalse( + writer.submit_finalization( + self.finish_command(start), self.sealed(writer) + ) + ) + self.assertFalse(writer.close(0.01)) + release.set() + self.wait_for(lambda: not writer.status().alive) + self.assertEqual(writer.context_budget.charged_bytes, 0) + self.assertEqual( + self.db.get_execution_contexts(start.run_id)["state"], "not_captured" + ) + + def test_bounded_close_keeps_inflight_charge_and_releases_queued(self): + self.check_bounded_close_payloads() + + def test_early_join_return_still_waits_for_deadline_and_releases_queued(self): + self.check_bounded_close_payloads(early_join=True) + + def check_bounded_close_payloads(self, early_join=False): + from utils.execution_context_budget import ContextBudget + + pool = ContextBudget(1000) + writer = HistoryWriter(lambda: self.db, clock=self.clock, context_budget=pool) + self.writers.append(writer) + writer.start() + self.wait_for(lambda: writer.status().accepting) + start = self.start_command() + self.db.write_execution_start(start) + entered, release = self.block() + original = self.db.write_execution_contexts_for_finish + first = self.sealed(writer) + charge = first.lease.byte_count + + def write(*args, **kwargs): + entered.set() + release.wait(3) + return original(*args, **kwargs) + + with patch.object(self.db, "write_execution_contexts_for_finish", write): + self.assertTrue( + writer.submit_finalization(self.finish_command(start), first) + ) + self.assertTrue(entered.wait(1)) + self.assertTrue( + writer.submit_finalization( + self.finish_command(start), self.sealed(writer) + ) + ) + native_join = writer._thread.join + close_clock = ManualClock() + + def return_early(timeout): + # Keep the real worker blocked while simulating an early native join. + self.assertTrue(writer.status().alive) + self.assertFalse(writer.submit(self.finish_command(start))) + native_join(0) + close_clock.advance(min(timeout, 0.004)) + + def yield_wait(delay): + self.assertGreater(delay, 0) + self.assertLessEqual(delay, 0.001) + close_clock.advance(delay) + + clock = ( + patch( + "utils.execution_history_writer.time", + SimpleNamespace(monotonic=close_clock, sleep=yield_wait), + ) + if early_join + else nullcontext() + ) + join = ( + patch.object(writer._thread, "join", return_early) + if early_join + else nullcontext() + ) + with clock, join: + monotonic = close_clock if early_join else time.monotonic + before = monotonic() + self.assertFalse(writer.close(0.02)) + elapsed = monotonic() - before + self.assertGreaterEqual(elapsed, 0.02) + self.assertLess(elapsed, 0.2) + self.assertEqual(pool.charged_bytes, charge) + replacement = HistoryWriter(lambda: self.db, context_budget=pool) + self.writers.append(replacement) + self.assertIsNone(replacement.context_budget.reserve(1000)) + release.set() + self.wait_for(lambda: not writer.status().alive) + self.assertEqual(pool.charged_bytes, 0) + + def test_close_drain_exits_without_sleeping_after_worker_join(self): + writer = self.writer() + clock = ManualClock() + native_join = writer._thread.join + + def join(timeout): + native_join(3) + self.assertFalse(writer._thread.is_alive()) + clock.advance(0.001) + + def sleep(delay): + self.fail("Close slept after the worker exited") + + with ( + patch( + "utils.execution_history_writer.time", + SimpleNamespace(monotonic=clock, sleep=sleep), + ), + patch.object(writer._thread, "join", join), + ): + self.assertTrue(writer.close(0.02)) + self.assertLess(clock(), 0.02) + self.assertFalse(writer.status().alive) + + def test_same_epoch_invalidation_between_transactions_drops_payload(self): + writer = self.writer() + start = self.start_command() + self.db.write_execution_start(start) + original = self.db.write_execution_finish + + def finish(command): + result = original(command) + writer.invalidate() + return result + + with patch.object(self.db, "write_execution_finish", finish): + self.assertTrue( + writer.submit_finalization( + self.finish_command(start), self.sealed(writer) + ) + ) + self.wait_for(lambda: writer.status().pending == 0) + self.assertEqual( + self.db.get_execution_run(start.run_id)["outcome"], "succeeded" + ) + self.assertEqual( + self.db.get_execution_contexts(start.run_id)["state"], "not_captured" + ) + self.assertEqual(writer.context_budget.charged_bytes, 0) + self.clock.advance(5) + self.wait_for(lambda: writer.status().accepting) + self.assertEqual( + self.db.get_execution_contexts(start.run_id)["state"], "not_captured" + ) + + def test_invalid_seals_release_and_keep_valid_outcome(self): + from utils.execution_context_budget import ContextBudget + + writer = self.writer() + cases = [ + dict(state="private text" * 100), + dict(reason="arbitrary_exception"), + dict(capture_generation=True), + dict(capture_generation=-1), + dict(capture_generation=None), + dict(lease=None), + ] + for item, changes in enumerate(cases, 20): + start = self.start_command(item) + self.db.write_execution_start(start) + sealed = self.sealed(writer) + malformed = replace(sealed, **changes) + if changes == {"lease": None}: + sealed.lease.release() + self.assertTrue( + writer.submit_finalization(self.finish_command(start), malformed) + ) + self.wait_for(lambda: writer.status().pending == 0) + self.assertEqual( + self.db.get_execution_run(start.run_id)["outcome"], "succeeded" + ) + self.assertEqual( + self.db.get_execution_contexts(start.run_id)["reason"], + "invalid_evidence", + ) + self.assertEqual(writer.context_budget.charged_bytes, 0) + foreign = ContextBudget(1000) + sealed = self.sealed(writer) + sealed.lease.release() + sealed = replace(sealed, lease=foreign.reserve(1000)) + start = self.start_command(30) + self.db.write_execution_start(start) + self.assertTrue(writer.submit_finalization(self.finish_command(start), sealed)) + self.wait_for(lambda: writer.status().pending == 0) + self.assertEqual(foreign.charged_bytes, 0) + self.assertEqual(self.db.get_execution_contexts(start.run_id)["contexts"], []) + + def test_custom_run_identity_receives_finalized_context(self): + writer = self.writer() + start = replace(self.start_command(), run_id=str(uuid4())) + self.assertTrue(writer.submit(start)) + self.assertTrue( + writer.submit_finalization(self.finish_command(start), self.sealed(writer)) + ) + self.wait_for(lambda: writer.status().pending == 0) + self.assertEqual( + self.db.get_execution_run(start.run_id)["outcome"], "succeeded" + ) + self.assertEqual( + self.db.get_execution_contexts(start.run_id)["state"], "complete" + ) + self.assertEqual(writer.context_budget.charged_bytes, 0) + + def test_invalidation_during_context_preparation_fences_actual_transaction(self): + import database.execution_contexts as contexts_db + + writer = self.writer() + start = self.start_command() + self.db.write_execution_start(start) + entered, release = self.block() + original = contexts_db._prepare + + def prepare(bundle): + prepared = original(bundle) + entered.set() + release.wait(3) + return prepared + + with patch.object(contexts_db, "_prepare", prepare): + self.assertTrue( + writer.submit_finalization( + self.finish_command(start), self.sealed(writer) + ) + ) + self.assertTrue(entered.wait(1)) + self.assertEqual( + self.db.get_execution_run(start.run_id)["outcome"], "succeeded" + ) + generation = writer.status().capture_generation + writer.invalidate() + self.assertGreater(writer.status().capture_generation, generation) + release.set() + self.wait_for(lambda: writer.status().pending == 0) + self.assertEqual( + self.db.get_execution_contexts(start.run_id)["state"], "not_captured" + ) + self.assertEqual(writer.context_budget.charged_bytes, 0) + self.clock.advance(5) + self.wait_for(lambda: writer.status().accepting) + self.assertEqual(writer.status().epoch, start.storage_epoch) + self.assertEqual( + self.db.get_execution_contexts(start.run_id)["state"], "not_captured" + ) + + def test_context_commit_serializes_invalidation_without_blocking_queue_or_close( + self, + ): + import database.execution_contexts as contexts_db + + writer = self.writer() + start = self.start_command() + self.db.write_execution_start(start) + entered, release = self.block() + invalidating, invalidated = threading.Event(), threading.Event() + original = contexts_db._write_capture + + def capture(*args, **kwargs): + entered.set() + release.wait(3) + return original(*args, **kwargs) + + def invalidate(): + invalidating.set() + writer.invalidate() + invalidated.set() + + with patch.object(contexts_db, "_write_capture", capture): + self.assertTrue( + writer.submit_finalization( + self.finish_command(start), self.sealed(writer) + ) + ) + self.assertTrue(entered.wait(1)) + invalidator = threading.Thread(target=invalidate) + invalidator.start() + try: + self.assertTrue(invalidating.wait(1)) + self.assertFalse(invalidated.wait(0.03)) + before = time.monotonic() + self.assertTrue(writer.status().accepting) + self.assertTrue(writer.submit(self.finish_command(start))) + self.assertLess(time.monotonic() - before, 0.1) + before = time.monotonic() + self.assertFalse(writer.close(0)) + self.assertLess(time.monotonic() - before, 0.1) + self.assertGreater(writer.context_budget.charged_bytes, 0) + self.assertEqual(writer.status().pending, 1) + finally: + release.set() + invalidator.join(3) + self.assertFalse(invalidator.is_alive()) + self.assertTrue(invalidated.is_set()) + self.wait_for(lambda: not writer.status().alive) + self.assertEqual( + self.db.get_execution_run(start.run_id)["outcome"], "succeeded" + ) + self.assertEqual( + self.db.get_execution_contexts(start.run_id)["state"], "complete" + ) + self.assertEqual(writer.context_budget.charged_bytes, 0) + def test_fifo_single_worker_and_thread_local_connection_cleanup(self): caller_conn = self.db.model.get_connection() threads, timeouts = [], [] @@ -481,6 +967,29 @@ class HistoryWriterTests(unittest.TestCase): self.assertEqual(writer.status().recording_gaps, 2**63 - 1) self.assertTrue(writer.close()) + def test_owner_invalidation_callback_runs_outside_both_locks_and_is_contained(self): + observations = [] + writer = None + + def callback(): + acquired = [ + lock.acquire(blocking=False) + for lock in (writer._lock, writer._context_commit_gate) + ] + observations.append(tuple(acquired)) + for lock, owned in zip( + (writer._lock, writer._context_commit_gate), acquired + ): + if owned: + lock.release() + raise RuntimeError("callback failure cannot escape") + + writer = HistoryWriter(lambda: self.db, on_invalidation=callback) + self.writers.append(writer) + writer.invalidate() + writer.close(0) + self.assertEqual(observations, [(True, True), (True, True)]) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_pair_execution.py b/tests/test_pair_execution.py index 5f6bea2..a81399e 100644 --- a/tests/test_pair_execution.py +++ b/tests/test_pair_execution.py @@ -326,3 +326,23 @@ class TestPairExecution(unittest.TestCase): record_pair_execution(extra, "1", **kwargs) with self.assertRaises(ExecutionEvidenceError): read_pair_execution(extra, "1") + + def test_png_evidence_cannot_populate_unobserved_attempt(self): + from utils.execution_context_capture import ContextCapture + from utils.execution_context_budget import ContextBudget + from utils.execution_history_events import RunStart + from utils.execution_history_writer import HistoryWriterStatus + + identity = "11111111-1111-4111-8111-111111111111" + self.execute() + state = HistoryWriterStatus(identity, True, 0, 0, True, 1) + manager = ContextCapture(1, lambda: state, ContextBudget()) + self.addCleanup(manager.revoke_all) + cap = manager.begin( + RunStart(identity, identity, 1, identity, identity, 0), 1, graph(), None + ) + sealed = manager.seal(cap) + if sealed.lease: + self.addCleanup(sealed.lease.release) + self.assertIsNone(sealed.bundle.contexts[0].positive.effective.value) + self.assertIsNone(sealed.bundle.contexts[0].negative.effective.value) diff --git a/tests/test_prompt_pair_node.py b/tests/test_prompt_pair_node.py index 693a6c1..671d4cd 100644 --- a/tests/test_prompt_pair_node.py +++ b/tests/test_prompt_pair_node.py @@ -158,3 +158,45 @@ class TestPromptPairNode(unittest.TestCase): ) result = self.node.encode_prompt(self.clip, "cat", negative_text="blur") self.assertEqual(result[2:], ("cat LORA", "blur")) + + def test_failed_negative_encoding_cannot_publish_positive_runtime_proof(self): + from contextlib import ExitStack + import sys + import types + from utils.execution_context_capture import ContextCapture + from utils.execution_context_budget import ContextBudget + from utils.execution_history_events import RunStart + from utils.execution_history_writer import HistoryWriterStatus + from unittest.mock import patch + + identity = "11111111-1111-4111-8111-111111111111" + state = HistoryWriterStatus(identity, True, 0, 0, True, 1) + budget = ContextBudget() + manager = ContextCapture(1, lambda: state, budget) + self.addCleanup(manager.revoke_all) + cap = manager.begin( + RunStart(identity, identity, 1, identity, identity, 0), + 1, + pair_graph(), + None, + ) + host = types.ModuleType("comfy_execution.utils") + host.get_executing_context = lambda: types.SimpleNamespace( + prompt_id=identity, node_id="1", list_index=0 + ) + self.clip.encode_from_tokens_scheduled.side_effect = [ + ("conditioning", "cat LORA"), + RuntimeError("negative failed"), + ] + with ExitStack() as stack: + stack.enter_context( + patch.dict(sys.modules, {"comfy_execution.utils": host}) + ) + with self.assertRaisesRegex(RuntimeError, "negative failed"): + self.node.encode_prompt(self.clip, "cat", negative_text="blur") + sealed = manager.seal(cap) + if sealed.lease: + self.addCleanup(sealed.lease.release) + self.assertIsNone(sealed.bundle.contexts[0].positive.effective.value) + self.assertIsNone(sealed.bundle.contexts[0].negative.effective.value) + self.assertEqual(self.db.list_prompt_pairs()["total"], 0) diff --git a/tests/vendor/comfyui-config.test.mjs b/tests/vendor/comfyui-config.test.mjs index 48c6a6c..ec2d3c8 100644 --- a/tests/vendor/comfyui-config.test.mjs +++ b/tests/vendor/comfyui-config.test.mjs @@ -27,7 +27,10 @@ test('execution script discovers only the new matrix and preserves all 78 stock 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.match(executions, /Total: 26 tests in 1 file/); + for (const project of ['classic-start', 'nodes2-start']) { + assert.equal(executions.split('\n').filter(line => line.includes(`[${project}]`)).length, 13); + } assert.doesNotMatch(executions, /prompt-manager.spec/); }); diff --git a/tests/vendor/comfyui-host.test.mjs b/tests/vendor/comfyui-host.test.mjs index 8bbfa19..78e11d5 100644 --- a/tests/vendor/comfyui-host.test.mjs +++ b/tests/vendor/comfyui-host.test.mjs @@ -91,23 +91,113 @@ test('read-only ledger helper distinguishes absent schema from missing, corrupt 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 runId = '50000000-0000-4000-8000-000000000001'; + const otherSession = '40000000-0000-4000-8000-000000000001'; + const epoch = execFileSync(python, ['-c', `import sys +from database.operations import PromptDatabase +from utils.execution_history_events import RunStart, make_finish +from utils.execution_contexts import ContextBundle, PromptSnapshot, RecipeContext, TextEvidence +from types import SimpleNamespace + +db = PromptDatabase(sys.argv[1]) +try: + epoch = db.get_execution_epoch() + c = db.model.get_connection() + with c: + c.executemany("INSERT INTO prompts(id,text,run_count) VALUES(?,?,?)", + [(7, ' positive ', 19), (8, 'negative', 23)]) + c.execute("INSERT INTO prompt_pairs(id,positive_prompt_id,negative_prompt_id,run_count) VALUES(2,7,8,31)") + for index, (job, session) in enumerate([ + (sys.argv[2], sys.argv[3]), + ('30000000-0000-4000-8000-000000000001', sys.argv[3]), + (sys.argv[2], sys.argv[4]), + ], 1): + run_id = '50000000-0000-4000-8000-%012d' % index + db.write_execution_start(RunStart(run_id, session, index, job, epoch, 100)) + db.write_execution_finish(make_finish(session, index, epoch, + SimpleNamespace(messages=[('execution_success', {'prompt_id': job})]), 200)) + def snapshot(raw): + def known(value): + return TextEvidence(value, 'recorded', 'runtime_return', None) + return PromptSnapshot(known(raw), known(raw), known(''), known('')) + db.write_execution_contexts(sys.argv[5], epoch, ContextBundle( + (RecipeContext('4', snapshot(' positive '), snapshot('negative'), 'inferred'),), + 'complete', 'captured')) + print(epoch) +finally: + db.close_all() +`, database, job, session, otherSession, runId], { cwd: new URL('../../', import.meta.url), encoding: 'utf8' }).trim(); 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' }, + const result = await readExecutionRuns(database, job, session); + assert.equal(result.schemaAvailable, true); + assert.deepEqual(result.rows, [{ + run_id: runId, job_id: job, session_id: session, queue_item_id: 1, + storage_epoch: epoch, record_version: 1, observed_dequeue_at: 100, + observed_callback_at: 200, executor_started_at: null, terminal_event_at: null, + outcome: 'succeeded', reason: 'terminal_success', diagnostics_omitted: 0, + node_id: null, node_type: null, exception_type: null, + }]); + assert.deepEqual(Object.keys(result.captures), [runId]); + const capture = result.captures[runId]; + assert.equal(capture.contexts.length, 1); + const contextId = capture.contexts[0].context_id; + assert.match(contextId, /^[a-f0-9]{64}$/); + const recorded = { provenance: 'recorded', source: 'runtime_return', reason: null }; + const evidence = { raw: recorded, effective: recorded, prepend: recorded, append: recorded }; + assert.deepEqual(capture, { + state: 'complete', reason: 'captured', format_version: 1, contexts: [{ + anchor_id: '4', context_id: contextId, + positive: { raw: ' positive ', effective: ' positive ', prepend: '', append: '' }, + negative: { raw: 'negative', effective: 'negative', prepend: '', append: '' }, + evidence: { format_version: 1, role_provenance: 'inferred', positive: evidence, negative: evidence }, + positive_prompt_id: 7, negative_prompt_id: 8, pair_id: 2, reference_epoch: epoch, + }], + }); + assert.deepEqual(result.usage, { [runId]: [ + { role: 'positive', run_count: 19 }, { role: 'negative', run_count: 23 }, { role: 'pair', run_count: 31 }, ] }); - assert.deepEqual((await readExecutionRuns(database, job, '40000000-0000-4000-8000-000000000001')).rows, []); + assert.deepEqual(Object.keys(result.snapshotIds), [runId]); + assert.equal(result.snapshotIds[runId].length, 1); + assert.deepEqual(Object.keys(result.snapshotIds[runId][0]).sort(), ['negative_snapshot_id', 'positive_snapshot_id']); + const { positive_snapshot_id: positiveId, negative_snapshot_id: negativeId } = result.snapshotIds[runId][0]; + assert.equal(positiveId, 'd2d716108f66314e127fe87bb81327426d2a148e341a8ca1f7cf6a0b0df2ce8c'); + assert.equal(negativeId, '07db4923fea7f8d00d1a94a902e0b85137471178b5e7bb67ad664ef3370770ed'); + const unfiltered = await readExecutionRuns(database, job); + assert.deepEqual(unfiltered.rows.map(row => row.session_id), [session, otherSession]); + const uncapturedId = '50000000-0000-4000-8000-000000000003'; + assert.deepEqual(unfiltered.captures[uncapturedId], { + state: 'not_captured', reason: 'missing', format_version: null, contexts: [], + }); + assert.deepEqual(unfiltered.usage[uncapturedId], []); + assert.deepEqual(unfiltered.snapshotIds[uncapturedId], []); + assert.deepEqual(await readExecutionRuns(database, job, '60000000-0000-4000-8000-000000000001'), { + schemaAvailable: true, rows: [], captures: {}, usage: {}, snapshotIds: {}, + }); 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, "' OR 1=1 --")); + const addRuns = (start, stop) => execFileSync(python, ['-c', `import sys +from database.operations import PromptDatabase +from utils.execution_history_events import RunStart + +db = PromptDatabase(sys.argv[1]) +try: + epoch = db.get_execution_epoch() + for index in range(int(sys.argv[4]), int(sys.argv[5])): + run_id = '70000000-0000-4000-8000-%012d' % index + db.write_execution_start(RunStart(run_id, sys.argv[3], index, sys.argv[2], epoch, 100)) +finally: + db.close_all() +`, database, job, session, String(start), String(stop)], { cwd: new URL('../../', import.meta.url) }); + addRuns(4, 258); + assert.equal((await readExecutionRuns(database, job)).rows.length, 256); + addRuns(258, 259); await assert.rejects(readExecutionRuns(database, job), /bound|256/); + execFileSync(python, ['-c', `import sqlite3,sys +with sqlite3.connect(sys.argv[1]) as c: + c.execute('DROP TABLE execution_context_capture') +`, database]); + await assert.rejects(readExecutionRuns(database, job, otherSession), /no such table: execution_context_capture/); await assert.rejects(readExecutionRuns(join(directory, 'missing.db'), job)); await writeFile(database, 'corrupt database'); await assert.rejects(readExecutionRuns(database, job)); diff --git a/utils/execution_context_budget.py b/utils/execution_context_budget.py new file mode 100644 index 0000000..0b319fa --- /dev/null +++ b/utils/execution_context_budget.py @@ -0,0 +1,111 @@ +"""Shared payload reservations with explicit, single-owner transport leases.""" + +from dataclasses import dataclass +import json +import threading + +from .execution_contexts import BUNDLE_STATES, CAPTURE_REASONS, ContextBundle + + +def _byte_count(value: int) -> None: + if type(value) is not int or value < 0: + raise ValueError("Byte count must be a nonnegative integer") + + +class ContextBudget: + def __init__(self, limit_bytes: int = 8388608) -> None: + _byte_count(limit_bytes) + self._limit = limit_bytes + self._charged = 0 + self._lock = threading.Lock() + + @property + def charged_bytes(self) -> int: + with self._lock: + return self._charged + + def reserve(self, byte_count: int) -> "PayloadLease | None": + _byte_count(byte_count) + with self._lock: + if byte_count > self._limit - self._charged: + return None + self._charged += byte_count + return PayloadLease._new(self, byte_count) + + +class PayloadLease: + def __init__(self, pool: ContextBudget, byte_count: int) -> None: + raise ValueError("Reserve payload leases through ContextBudget") + + @classmethod + def _new(cls, pool: ContextBudget, byte_count: int) -> "PayloadLease": + owner = object.__new__(cls) + owner._pool = pool + owner._bytes = byte_count + owner._active = True + return owner + + @property + def byte_count(self) -> int: + with self._pool._lock: + self._require_owner() + return self._bytes + + def _require_owner(self) -> None: + if not self._active: + raise ValueError("Inactive payload lease") + + def transfer(self) -> "PayloadLease": + with self._pool._lock: + self._require_owner() + owner = PayloadLease._new(self._pool, self._bytes) + self._active = False + return owner + + def shrink_to(self, byte_count: int) -> None: + _byte_count(byte_count) + with self._pool._lock: + self._require_owner() + if byte_count > self._bytes: + raise ValueError("Cannot grow payload lease") + self._pool._charged -= self._bytes - byte_count + self._bytes = byte_count + + def release(self) -> None: + with self._pool._lock: + if self._active: + self._pool._charged -= self._bytes + self._active = False + + +SHARED_CONTEXT_BUDGET = ContextBudget() + + +@dataclass(frozen=True) +class SealedContext: + bundle: ContextBundle | None + lease: PayloadLease | None + state: str + reason: str + capture_generation: int | None = None + + def validate_header(self) -> None: + if ( + type(self.state) is not str + or self.state not in BUNDLE_STATES + or type(self.reason) is not str + or self.reason not in CAPTURE_REASONS + or ( + self.capture_generation is not None + and ( + type(self.capture_generation) is not int + or not 0 <= self.capture_generation <= 2**63 - 1 + ) + ) + ): + raise ValueError("Invalid capture status") + header = json.dumps( + [self.state, self.reason, self.capture_generation], separators=(",", ":") + ).encode("utf-8") + if len(header) > 256: + raise ValueError("Oversized capture status") diff --git a/utils/execution_context_capture.py b/utils/execution_context_capture.py new file mode 100644 index 0000000..05cfaee --- /dev/null +++ b/utils/execution_context_capture.py @@ -0,0 +1,451 @@ +"""Revocable, bounded runtime evidence owned by an actual dequeue attempt.""" + +from __future__ import annotations + +from contextvars import ContextVar +from dataclasses import dataclass, replace +import threading +from typing import Callable, TYPE_CHECKING +from uuid import uuid4 +import weakref + +from .execution_context_budget import ContextBudget, PayloadLease, SealedContext +from .execution_context_graph import ( + GraphProjection, + GraphRole, + SourceRef, + project_contexts, +) +from .execution_contexts import ( + ContextBundle, + PromptSnapshot, + RecipeContext, + TextEvidence, + bundle_payload_size, + text_size, + MAX_ATTEMPT_BYTES, +) +from .execution_history_events import RunStart, encode_command + +if TYPE_CHECKING: + from .execution_history_writer import HistoryWriterStatus + + +@dataclass(frozen=True) +class AttemptCapability: + nonce: str + run_id: str + session_id: str + queue_item_id: int + job_id: str + storage_epoch: str + capture_generation: int + sources: tuple[SourceRef, ...] + + +@dataclass +class _Attempt: + capability: AttemptCapability + projection: GraphProjection + lease: PayloadLease | None + samples: dict + retained_bytes: int + reason: str = "" + executor_identity: int | None = None + executor_rejected: bool = False + + +_binding: ContextVar[tuple | None] = ContextVar("prompt_manager_attempt", default=None) + + +def active_capture() -> tuple[ContextCapture, AttemptCapability] | None: + binding = _binding.get() + if binding is not None: + manager = binding[0]() + if manager is not None and manager.valid(binding[1]): + return manager, binding[1] + return None + + +def unbind_current() -> None: + binding = _binding.get() + _binding.set(None) + if binding is not None: + manager = binding[0]() + if manager is not None: + manager.revoke(binding[1]) + + +def _unknown(reason: str = "missing") -> PromptSnapshot: + field = TextEvidence(None, "unknown", "unknown", reason) + return PromptSnapshot(field, field, field, field) + + +def _snapshot_bytes(snapshot: PromptSnapshot) -> int: + # Reuse the scalar validator, including provenance and per-field bounds. + bundle = ContextBundle( + (RecipeContext("x", snapshot, None, "inferred"),), "partial", "partial" + ) + return bundle_payload_size(bundle) + + +def _bounded_snapshot(snapshot: PromptSnapshot) -> tuple[PromptSnapshot, bool]: + if type(snapshot) is not PromptSnapshot: + raise ValueError("invalid_snapshot") + fields, limited = [], False + for field in (snapshot.raw, snapshot.effective, snapshot.prepend, snapshot.append): + if type(field) is not TextEvidence: + raise ValueError("invalid_field") + if field.value is not None: + try: + text_size(field.value) + except ValueError: + field = _unknown("limit_reached").raw + limited = True + fields.append(field) + if fields[1].value is not None and fields[1].provenance != "recorded": + fields[1] = _unknown("missing").effective + bounded = PromptSnapshot(*fields) + _snapshot_bytes(bounded) + return bounded, limited + + +class ContextCapture: + def __init__( + self, + queue_identity: int, + status_reader: Callable[[], HistoryWriterStatus], + budget: ContextBudget, + ) -> None: + self.queue_identity = queue_identity + self._status_reader = status_reader + self._budget = budget + self._lock = threading.RLock() + self._attempts: dict[str, _Attempt] = {} + + def _current(self, capability: AttemptCapability) -> _Attempt | None: + attempt = self._attempts.get(capability.nonce) + if attempt is None or attempt.capability is not capability: + return None + status = self._status_reader() + if ( + not status.accepting + or status.epoch != capability.storage_epoch + or status.capture_generation != capability.capture_generation + ): + self._drop(capability) + return None + return attempt + + def valid(self, capability: AttemptCapability) -> bool: + with self._lock: + return self._current(capability) is not None + + def begin( + self, start: RunStart, generation: int, graph: object, outputs: object + ) -> AttemptCapability: + unbind_current() + encode_command(start) + with self._lock: + cap = AttemptCapability( + str(uuid4()), + start.run_id, + start.session_id, + start.queue_item_id, + start.job_id, + start.storage_epoch, + generation, + (), + ) + if len(self._attempts) >= 128: + return cap + lease = self._budget.reserve(MAX_ATTEMPT_BYTES) + projection = GraphProjection((), (), "") + reason = "budget_exhausted" if lease is None else "" + retained = 0 + if lease is not None: + try: + projection = project_contexts(graph, outputs) + retained = ( + sum( + _snapshot_bytes(role.inferred) if role.inferred else 0 + for recipe in projection.recipes + for role in (recipe.positive, recipe.negative) + ) + + sum( + _snapshot_bytes(snapshot) + for _, snapshot in projection.source_snapshots + ) + + sum( + text_size(s.node_id, 512) + text_size(s.node_type, 128) + 32 + for s in projection.sources + ) + ) + if retained > MAX_ATTEMPT_BYTES: + raise ValueError("projection_limit") + except ValueError: + projection = GraphProjection((), (), "") + reason = "graph_limit" + retained = 0 + cap = replace(cap, sources=projection.sources) + self._attempts[cap.nonce] = _Attempt( + cap, projection, lease, {}, retained, reason + ) + _binding.set((weakref.ref(self), cap)) + self._current(cap) + return cap + + def bind_executor(self, capability: AttemptCapability, identity: int) -> None: + """Bind one scalar host identity; recursive starts invalidate cache proof.""" + with self._lock: + attempt = self._current(capability) + if attempt is not None: + if attempt.executor_identity is not None: + attempt.executor_rejected = True + else: + attempt.executor_identity = identity + + def matches_executor(self, capability: AttemptCapability, identity: int) -> bool: + with self._lock: + attempt = self._current(capability) + return bool( + attempt is not None + and not attempt.executor_rejected + and attempt.executor_identity == identity + ) + + def source_snapshot( + self, capability: AttemptCapability, source: SourceRef + ) -> PromptSnapshot | None: + with self._lock: + attempt = self._current(capability) + if attempt is not None: + for ref, snapshot in attempt.projection.source_snapshots: + if ref == source: + return snapshot + return None + + def observe( + self, + capability: AttemptCapability, + source: SourceRef, + invocation_index: int, + snapshot: PromptSnapshot, + ) -> None: + try: + with self._lock: + attempt = self._current(capability) + if ( + attempt is None + or attempt.lease is None + or source not in capability.sources + or type(source) is not SourceRef + or type(invocation_index) is not int + or not 0 <= invocation_index <= 2**63 - 1 + ): + return + snapshot, limited = _bounded_snapshot(snapshot) + size = _snapshot_bytes(snapshot) + if ( + limited + or any( + field.reason == "limit_reached" + for field in ( + snapshot.raw, + snapshot.effective, + snapshot.prepend, + snapshot.append, + ) + ) + ) and attempt.reason not in ("multiple_samples", "limit_reached"): + attempt.reason = "field_limit" + key = source, invocation_index + if key in attempt.samples: + if attempt.samples[key] != snapshot: + attempt.reason = "multiple_samples" + return + if len(attempt.samples) >= 64: + attempt.reason = "multiple_samples" + return + if attempt.retained_bytes + size > MAX_ATTEMPT_BYTES: + attempt.reason = "limit_reached" + return + attempt.samples[key] = snapshot + attempt.retained_bytes += size + except (ValueError, AttributeError, TypeError): + return + + @staticmethod + def _discard_payload(attempt: _Attempt) -> None: + attempt.projection = GraphProjection((), (), "") + attempt.samples.clear() + attempt.retained_bytes = 0 + + def _drop(self, capability: AttemptCapability) -> _Attempt | None: + attempt = self._attempts.get(capability.nonce) + if attempt is not None and attempt.capability is capability: + self._attempts.pop(capability.nonce) + self._discard_payload(attempt) + if attempt.lease is not None: + attempt.lease.release() + return attempt + return None + + def revoke(self, capability: AttemptCapability) -> None: + with self._lock: + self._drop(capability) + + def revoke_all(self) -> None: + with self._lock: + for attempt in self._attempts.values(): + self._discard_payload(attempt) + if attempt.lease is not None: + attempt.lease.release() + self._attempts.clear() + + def seal(self, capability: AttemptCapability) -> SealedContext: + with self._lock: + attempt = self._current(capability) + if attempt is None: + return SealedContext(None, None, "unavailable", "unavailable") + self._attempts.pop(capability.nonce) + lease = attempt.lease + if lease is None: + return SealedContext(None, None, "unavailable", attempt.reason) + try: + ambiguous = attempt.reason == "multiple_samples" + + def role_snapshot(role: GraphRole) -> PromptSnapshot: + nonlocal ambiguous + if role.reason == "source-transformed": + return _unknown("unsupported") + if not role.sources: + return role.inferred or _unknown("unsupported") + values = [ + snapshot + for (source, _), snapshot in attempt.samples.items() + if source in role.sources + ] + if len(values) > 1 or attempt.reason == "multiple_samples": + ambiguous = True + return _unknown("ambiguous") + if len(role.sources) == 1 and len(values) == 1: + return values[0] + inferred = role.inferred + if len(role.sources) == 1: + snapshots = attempt.projection.source_snapshots + inferred = next( + ( + snapshot + for source, snapshot in snapshots + if source == role.sources[0] + ), + inferred, + ) + if inferred is not None: + s = inferred + return PromptSnapshot( + s.raw, _unknown().effective, s.prepend, s.append + ) + return _unknown() + + contexts = tuple( + RecipeContext( + r.anchor_id, + role_snapshot(r.positive), + role_snapshot(r.negative), + "inferred", + ) + for r in attempt.projection.recipes + ) + supported = any( + role.sources or role.inferred is not None + for r in attempt.projection.recipes + for role in (r.positive, r.negative) + ) + complete = bool(contexts) and all( + field.value is not None + for r in contexts + for s in (r.positive, r.negative) + for field in (s.raw, s.effective, s.prepend, s.append) + ) + state = ( + "ambiguous" if ambiguous else "complete" if complete else "partial" + ) + reason = ( + "multiple_samples" + if ambiguous + else "captured" if complete else "partial" + ) + if attempt.reason == "field_limit" and not ambiguous: + state, reason = "partial", "field_limit" + if not supported: + contexts, state, reason = (), "unavailable", "no_supported_context" + if attempt.reason in ("limit_reached", "graph_limit"): + contexts, state, reason = (), "unavailable", attempt.reason + bundle = ContextBundle(contexts, state, reason) + size = bundle_payload_size(bundle) + if size + attempt.retained_bytes > MAX_ATTEMPT_BYTES: + raise ValueError("assembly_limit") + self._discard_payload(attempt) + lease.shrink_to(size) + return SealedContext( + bundle, lease, state, reason, capability.capture_generation + ) + except ValueError: + self._discard_payload(attempt) + lease.release() + return SealedContext(None, None, "unavailable", "limit_reached") + + +def record_runtime_context( + node_type: str, + raw: str, + prepend: str, + append: str, + effective: str, + *, + role: str | None = None, +) -> None: + """Record only a successful node return in a trusted, bound host invocation.""" + try: + binding = active_capture() + if binding is None: + return + # Optional host import is independent of our package namespace. + from comfy_execution.utils import get_executing_context + + host = get_executing_context() + manager, cap = binding + if ( + host is None + or type(host.prompt_id) is not str + or host.prompt_id != cap.job_id + ): + return + if type(host.node_id) is not str or type(host.list_index) is not int: + return + if node_type == "PromptManagerPair": + slots = ( + (0, 2) if role == "positive" else (1, 3) if role == "negative" else () + ) + else: + slots = ( + (0,) + if node_type == "PromptManagerText" + else (0, 1) if role is None else () + ) + for slot in slots: + source = SourceRef(host.node_id, node_type, slot) + if source not in cap.sources: + continue + snapshot = PromptSnapshot( + *( + TextEvidence(value, "recorded", "runtime_return", None) + for value in (raw, effective, prepend, append) + ) + ) + manager.observe(cap, source, host.list_index, snapshot) + except Exception: + # Capture must never change the host result or expose input/exception text. + return diff --git a/utils/execution_context_graph.py b/utils/execution_context_graph.py new file mode 100644 index 0000000..740cfe4 --- /dev/null +++ b/utils/execution_context_graph.py @@ -0,0 +1,426 @@ +"""Bounded projection of declared sampler roles, independent of the host.""" + +import re +from dataclasses import dataclass +from .execution_contexts import PromptSnapshot, TextEvidence, text_size +from .prompt_pair_resolver import _DEFAULT_ADAPTERS + + +@dataclass(frozen=True) +class SourceRef: + node_id: str + node_type: str + output_slot: int + + +@dataclass(frozen=True) +class GraphRole: + sources: tuple[SourceRef, ...] + inferred: PromptSnapshot | None + reason: str + + +@dataclass(frozen=True) +class DeclaredRecipe: + anchor_id: str + positive: GraphRole + negative: GraphRole + + +@dataclass(frozen=True) +class GraphProjection: + recipes: tuple[DeclaredRecipe, ...] + sources: tuple[SourceRef, ...] + reason: str + source_snapshots: tuple[tuple[SourceRef, PromptSnapshot], ...] = () + + +_SAMPLERS = {"KSampler", "KSamplerAdvanced", "SamplerCustom"} +_PM = { + "PromptManager", + "PromptManagerFocused", + "PromptManagerSplit", + "PromptManagerText", + "PromptManagerPair", +} + + +class _Unknown(ValueError): + pass + + +class _Limit(_Unknown): + pass + + +def _identifier(value: object) -> str: + if type(value) is int: + if abs(value) > 10**127: + raise _Unknown("invalid_node_identifier") + value = str(value) + try: + text_size(value, 512) + if len(value) > 128: + raise ValueError("oversized_node_identifier") + except ValueError: + raise _Unknown("invalid_node_identifier") from None + if not value: + raise _Unknown("invalid_node_identifier") + return value + + +def _link(value: object) -> tuple[str, int]: + if ( + type(value) not in (tuple, list) + or len(value) != 2 + or type(value[1]) is not int + or value[1] < 0 + ): + raise _Unknown("expected_typed_graph_link") + return _identifier(value[0]), value[1] + + +class _Projector: + def __init__(self, graph: dict) -> None: + self.graph = graph + self.visits = 0 + self.retained = 0 + self.transformed = False + self.selected_seen = set() + + def node(self, node_id: str, depth: int) -> tuple[str, dict]: + self.visits += 1 + if self.visits > 500 or depth > 64: + raise _Unknown("graph_resolution_limit_reached") + node = self.graph.get(node_id) + if type(node) is not dict or type(node.get("inputs")) is not dict: + raise _Unknown("missing_or_malformed_node") + if len(node["inputs"]) > 500: + raise _Unknown("node_input_limit_reached") + kind = node.get("class_type") + try: + text_size(kind, 128) + except ValueError: + raise _Unknown("invalid_node_type") from None + return kind, node["inputs"] + + def text( + self, value: object, kind: str, depth: int = 1, active: frozenset = frozenset() + ) -> tuple[PromptSnapshot, tuple]: + if type(value) is str and kind == "string": + text_size(value) + return self.snapshot(value, value, "", ""), () + node_id, slot = _link(value) + key = (node_id, slot) + if key in active: + raise _Unknown("cyclic_prompt_graph") + name, inputs = self.node(node_id, depth) + active = active | {key} + if name == "Reroute" and slot == 0: + keys = [key for key in ("value", "input", "in") if key in inputs] + if len(keys) != 1: + raise _Unknown("ambiguous_reroute") + return self.text(inputs[keys[0]], kind, depth + 1, active) + if ( + name in ("StringConcatenate", "Text Concatenate") + and slot == 0 + and kind == "string" + ): + keys = ( + ["string_a", "string_b"] + if name == "StringConcatenate" + else sorted( + key + for key in inputs + if type(key) is str and re.fullmatch(r"text_[a-z]", key) + ) + ) + delimiter = inputs.get( + "delimiter", "" if name == "StringConcatenate" else " " + ) + text_size(delimiter) + clean_value = inputs.get("clean_whitespace", "true") + if name == "Text Concatenate" and type(clean_value) not in (str, bool): + raise _Unknown("invalid_whitespace_mode") + clean = name == "Text Concatenate" and ( + clean_value is True + or type(clean_value) is str + and clean_value.lower() == "true" + ) + if not keys: + raise _Unknown("missing_concatenation_inputs") + parts, sources, size = [], [], 0 + for input_name in keys: + child, refs = self.text( + inputs.get(input_name, ""), "string", depth + 1, active + ) + part = child.effective.value.strip() if clean else child.effective.value + if part or name == "StringConcatenate": + size += text_size(part) + (text_size(delimiter) if parts else 0) + if size > 65536: + raise _Unknown("text_field_limit_reached") + parts.append(part) + sources.extend(refs) + if len(sources) > 64: + raise _Unknown("source_reference_limit_reached") + text = delimiter.join(parts) + return self.snapshot( + text, text, "", "", composed=True, raw_composed=True + ), tuple(sources) + adapter = _DEFAULT_ADAPTERS.get((name, slot)) + if adapter is None or adapter.kind != kind: + raise _Unknown("unsupported_node_semantics_or_output_slot") + child, sources = self.text( + inputs.get(adapter.input_name), "string", depth + 1, active + ) + if name in _PM: + sources = (SourceRef(node_id, name, slot),) + if adapter.compose: + before, before_sources = self.text( + inputs.get("prepend_text", ""), "string", depth + 1, active + ) + after, after_sources = self.text( + inputs.get("append_text", ""), "string", depth + 1, active + ) + raw = child.effective.value + prepend, append = before.effective.value, after.effective.value + parts = [part for part in (prepend.strip(), raw, append.strip()) if part] + if sum(text_size(part) for part in parts) + max(0, len(parts) - 1) > 65536: + raise _Unknown("text_field_limit_reached") + effective = " ".join(parts) + text_size(effective) + snapshot = self.retain( + PromptSnapshot( + child.effective, + TextEvidence(effective, "inferred", "graph_semantics", None), + before.effective, + after.effective, + ) + ) + if name not in _PM: + sources += before_sources + after_sources + elif name in _PM: + snapshot = self.retain( + PromptSnapshot( + child.effective, + child.effective, + TextEvidence("", "inferred", "graph_literal", None), + TextEvidence("", "inferred", "graph_literal", None), + ) + ) + else: + snapshot = child + if len(sources) > 64: + raise _Unknown("source_reference_limit_reached") + return snapshot, sources + + def source_snapshot(self, source: SourceRef) -> PromptSnapshot: + name, inputs = self.node(source.node_id, 1) + adapter = _DEFAULT_ADAPTERS[(name, source.output_slot)] + unknown = TextEvidence(None, "unknown", "unknown", "missing") + + def literal(value: object) -> TextEvidence: + try: + snapshot, sources = self.text(value, "string", 2) + # Stateful PM outputs cannot prove their injected result from inputs. + return unknown if sources else snapshot.effective + except _Limit: + raise + except ValueError: + return unknown + + raw = literal(inputs.get(adapter.input_name)) + before = literal(inputs.get("prepend_text", "") if adapter.compose else "") + after = literal(inputs.get("append_text", "") if adapter.compose else "") + return self.retain(PromptSnapshot(raw, unknown, before, after)) + + def discover_sources( + self, value: object, kind: str, depth: int = 1, active: frozenset = frozenset() + ) -> tuple[SourceRef, ...]: + if type(value) is str and kind == "string": + return () + node_id, slot = _link(value) + key = (node_id, slot) + if key in active: + raise _Unknown("cyclic_prompt_graph") + name, inputs = self.node(node_id, depth) + active = active | {key} + adapter = _DEFAULT_ADAPTERS.get((name, slot)) + if adapter is not None and adapter.kind == kind: + if name in _PM: + return (SourceRef(node_id, name, slot),) + return self.discover_sources( + inputs.get(adapter.input_name), "string", depth + 1, active + ) + if name == "Reroute" and slot == 0: + keys = [key for key in ("value", "input", "in") if key in inputs] + if len(keys) != 1: + raise _Unknown("ambiguous_reroute") + return self.discover_sources(inputs[keys[0]], kind, depth + 1, active) + if ( + name in ("StringConcatenate", "Text Concatenate") + and slot == 0 + and kind == "string" + ): + keys = ( + ["string_a", "string_b"] + if name == "StringConcatenate" + else sorted( + key + for key in inputs + if type(key) is str and re.fullmatch(r"text_[a-z]", key) + ) + ) + sources = [] + for key in keys: + sources.extend( + self.discover_sources( + inputs.get(key, ""), "string", depth + 1, active + ) + ) + if len(sources) > 64: + raise _Unknown("source_limit_reached") + self.transformed = self.transformed or bool(sources) + return tuple(sources) + return () + + def retain(self, snapshot: PromptSnapshot) -> PromptSnapshot: + size = sum( + (text_size(field.value) if field.value is not None else 0) + 64 + for field in ( + snapshot.raw, + snapshot.effective, + snapshot.prepend, + snapshot.append, + ) + ) + if self.retained + size > 1048576: + raise _Limit("projection_payload_limit_reached") + self.retained += size + return snapshot + + def snapshot( + self, + raw: str, + effective: str, + prepend: str, + append: str, + composed: bool = False, + raw_composed: bool = False, + ) -> PromptSnapshot: + return self.retain( + PromptSnapshot( + *( + TextEvidence( + value, + "inferred", + ( + "graph_semantics" + if (composed and index == 1 or raw_composed and index == 0) + else "graph_literal" + ), + None, + ) + for index, value in enumerate((raw, effective, prepend, append)) + ) + ) + ) + + def role(self, value: object) -> GraphRole: + sources = () + try: + self.transformed = False + # Only the capture-source-to-consumer path determines transforms. + sources = self.discover_sources(value, "conditioning") + snapshot, _ = self.text(value, "conditioning") + return GraphRole( + sources, snapshot, "source-transformed" if self.transformed else "" + ) + except _Limit: + raise + except ValueError as error: + reason = str(error) if isinstance(error, _Unknown) else "invalid_text_value" + return GraphRole( + sources, + None, + "source-transformed" if self.transformed and sources else reason, + ) + + def anchors(self, node_id: str, depth: int = 1) -> tuple[str, ...]: + if node_id in self.selected_seen: + return () + self.selected_seen.add(node_id) + name, inputs = self.node(node_id, depth) + anchors = [node_id] if name in _SAMPLERS else [] + for value in inputs.values(): + if ( + type(value) not in (tuple, list) + or len(value) != 2 + or type(value[0]) not in (str, int) + or type(value[1]) is not int + or value[1] < 0 + ): + continue + parent, _ = _link(value) + anchors.extend(self.anchors(parent, depth + 1)) + if len(anchors) > 16: + raise _Unknown("recipe_context_limit_reached") + return tuple(anchors) + + +def project_contexts(graph: object, outputs: object) -> GraphProjection: + """Keep declared roles and graph inference, never evidence of consumption. + + Sources are PM capture points. An inferred snapshot may compose text; its + presence never promotes a source observation to recorded effective text. + """ + try: + if type(graph) is not dict or len(graph) > 500: + raise _Unknown("invalid_or_oversized_graph") + normalized = {_identifier(key): value for key, value in graph.items()} + if len(normalized) != len(graph): + raise _Unknown("duplicate_graph_identifiers") + resolver = _Projector(normalized) + if outputs is None: + anchors = [ + key + for key, value in normalized.items() + if type(value) is dict + and type(value.get("class_type")) is str + and value["class_type"] in _SAMPLERS + ] + elif type(outputs) in (list, tuple) and len(outputs) <= 500: + anchors = [] + for output in outputs: + anchors.extend(resolver.anchors(_identifier(output))) + anchors = list(dict.fromkeys(anchors)) + else: + raise _Unknown("invalid_output_selection") + if len(anchors) > 16: + raise _Unknown("recipe_context_limit_reached") + recipes, sources = [], [] + for anchor in anchors: + _, inputs = resolver.node(anchor, 1) + positive = resolver.role(inputs.get("positive")) + negative = resolver.role(inputs.get("negative")) + recipes.append(DeclaredRecipe(anchor, positive, negative)) + sources.extend(positive.sources + negative.sources) + sources = tuple(dict.fromkeys(sources)) + if len(sources) > 64: + raise _Unknown("source_reference_limit_reached") + reason = "" if recipes else "no_selected_sampler_anchors" + if any( + role.reason + for recipe in recipes + for role in (recipe.positive, recipe.negative) + ): + reason = "some_roles_require_inferred_or_unresolved_text" + return GraphProjection( + tuple(recipes), + sources, + reason, + tuple((source, resolver.source_snapshot(source)) for source in sources), + ) + except ValueError as error: + reason = str(error) if isinstance(error, _Unknown) else "invalid_text_value" + return GraphProjection((), (), reason) diff --git a/utils/execution_context_host.py b/utils/execution_context_host.py new file mode 100644 index 0000000..c3d8ba3 --- /dev/null +++ b/utils/execution_context_host.py @@ -0,0 +1,155 @@ +"""Pinned-host v1 adapter for local cached STRING evidence only.""" + +import functools +import threading +from typing import Callable + +from .execution_context_capture import active_capture +from .execution_contexts import PromptSnapshot, TextEvidence + +_STATE_ATTRIBUTE = "_prompt_manager_executor_messages_v1" +_registration_lock = threading.Lock() + + +class _Messages: + def __init__(self, executor_class: type) -> None: + self.executor_class = executor_class + self.lock = threading.Lock() + self.callbacks = {} + self.wrapper = None + self.original = None + self.installed = False + + def subscribe(self, key: str, callback: Callable) -> Callable[[], None]: + with self.lock: + if not self.installed: + original = self.executor_class.add_message + if not callable(original): + raise ValueError("Executor message capability unavailable") + + @functools.wraps(original) + def observe(executor, *args, **kwargs): + result = original(executor, *args, **kwargs) + with self.lock: + callbacks = tuple(entry[1] for entry in self.callbacks.values()) + for notify in callbacks: + try: + notify(executor, args, kwargs) + except Exception: + # Best-effort observers cannot change host results. + pass + return result + + self.executor_class.add_message = observe + self.original, self.wrapper, self.installed = original, observe, True + 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 self.executor_class.add_message is self.wrapper + ): + self.executor_class.add_message = self.original + self.installed = False + + return unsubscribe + + +def subscribe_executor_messages(key: str, callback: Callable) -> Callable[[], None]: + """Observe successful original calls; preserve earlier and later wrappers.""" + from execution import PromptExecutor + + with _registration_lock: + state = vars(PromptExecutor).get(_STATE_ATTRIBUTE) + if state is None: + state = _Messages(PromptExecutor) + setattr(PromptExecutor, _STATE_ATTRIBUTE, state) + return state.subscribe(key, callback) + + +def observe_cached_context( + executor: object, message_type: object, message_data: object +) -> None: + """Copy bounded local strings without cache/provider execution or retention.""" + try: + binding = active_capture() + if binding is None or type(message_type) is not str: + return + manager, cap = binding + if ( + type(message_data) is not dict + or type(message_data.get("prompt_id")) is not str + or message_data["prompt_id"] != cap.job_id + or id(executor.server.prompt_queue) != manager.queue_identity + ): + return + if message_type == "execution_start": + manager.bind_executor(cap, id(executor)) + return + if message_type != "execution_cached" or not manager.matches_executor( + cap, id(executor) + ): + return + nodes = message_data.get("nodes") + if type(nodes) not in (list, tuple) or len(nodes) > 500: + return + from comfy_execution import caching + from execution import CacheEntry + + cache = executor.caches.outputs + supported = tuple( + cls + for name in ("HierarchicalCache", "LRUCache", "RAMPressureCache") + if isinstance(cls := getattr(caching, name, None), type) + ) + if type(cache) not in supported: + return + local = getattr(cache, "get_local", None) + if not callable(local): + return + # Only source IDs retained by the bounded submitted graph are queried. + for source in cap.sources: + if not any(type(node) is str and node == source.node_id for node in nodes): + continue + slot = ( + (3 if source.output_slot in (1, 3) else 2) + if source.node_type == "PromptManagerPair" + else 0 if source.node_type == "PromptManagerText" else 1 + ) + entry = local(source.node_id) + if type(entry) is not CacheEntry: + continue + outputs = entry.outputs + if type(outputs) not in (list, tuple) or not slot < len(outputs) <= 4: + continue + values = outputs[slot] + if type(values) not in (list, tuple) or len(values) > 64: + continue + inferred = manager.source_snapshot(cap, source) + if inferred is None: + continue + for index, value in enumerate(values): + if type(value) is not str: + continue + manager.observe( + cap, + source, + index, + PromptSnapshot( + inferred.raw, + TextEvidence(value, "recorded", "host_cache_output", None), + inferred.prepend, + inferred.append, + ), + ) + except Exception: + # Unknown host versions and malformed entries leave evidence unavailable. + return diff --git a/utils/execution_contexts.py b/utils/execution_contexts.py new file mode 100644 index 0000000..64c537f --- /dev/null +++ b/utils/execution_contexts.py @@ -0,0 +1,201 @@ +"""Scalar contracts for bounded execution context capture.""" + +import hashlib +import json +import re +from dataclasses import dataclass + + +@dataclass(frozen=True) +class TextEvidence: + value: str | None + provenance: str + source: str + reason: str | None + + +@dataclass(frozen=True) +class PromptSnapshot: + raw: TextEvidence + effective: TextEvidence + prepend: TextEvidence + append: TextEvidence + + +@dataclass(frozen=True) +class RecipeContext: + anchor_id: str + positive: PromptSnapshot | None + negative: PromptSnapshot | None + role_provenance: str + + +@dataclass(frozen=True) +class ContextBundle: + contexts: tuple[RecipeContext, ...] + state: str + reason: str + format_version: int = 1 + + +MAX_FIELD_BYTES = 64 * 1024 +MAX_CONTEXT_BYTES = 256 * 1024 +MAX_ATTEMPT_BYTES = 1024 * 1024 + + +def text_size(value: str, limit: int = MAX_FIELD_BYTES) -> int: + """Reject oversize strings before allocating encoded copies.""" + if type(value) is not str or len(value) > limit: + raise ValueError("Invalid or oversized text") + size = 0 + for char in value: + code = ord(char) + if 0xD800 <= code <= 0xDFFF: + raise ValueError("Invalid Unicode") + size += 1 if code < 128 else 2 if code < 2048 else 3 if code < 65536 else 4 + if size > limit: + raise ValueError("Oversized text") + return size + + +def canonical_content(kind: str, fields: tuple[str | None, ...]) -> bytes: + text_size(kind, 128) + if type(fields) is not tuple: + raise ValueError("Invalid canonical fields") + for value in fields: + if value is not None: + text_size(value) + return json.dumps( + [1, kind, *fields], ensure_ascii=False, allow_nan=False, separators=(",", ":") + ).encode("utf-8") + + +def content_digest(kind: str, fields: tuple[str | None, ...]) -> str: + return hashlib.sha256(canonical_content(kind, fields)).hexdigest() + + +def snapshot_digest(snapshot: PromptSnapshot) -> str: + return content_digest( + "snapshot", + tuple( + getattr(snapshot, field).value + for field in ("raw", "effective", "prepend", "append") + ), + ) + + +def context_digest(context: RecipeContext) -> str: + return content_digest( + "context", + tuple( + snapshot_digest(side) if side is not None else None + for side in (context.positive, context.negative) + ), + ) + + +BUNDLE_STATES = frozenset( + {"complete", "partial", "ambiguous", "unavailable", "conflicting"} +) + +CAPTURE_REASONS = frozenset( + { + "", + "captured", + "partial", + "unavailable", + "unsupported", + "ambiguous", + "limit_reached", + "missing", + "conflicting", + "no_supported_context", + "budget_exhausted", + "invalid_evidence", + "multiple_samples", + "source_transformed", + "field_limit", + "graph_limit", + } +) +UNKNOWN_REASONS = frozenset( + {"unavailable", "unsupported", "ambiguous", "limit_reached", "missing"} +) +EVIDENCE_SOURCES = { + "recorded": frozenset({"runtime_return", "host_cache_output"}), + "inferred": frozenset({"graph_literal", "graph_semantics"}), + "unknown": frozenset({"unknown"}), +} + + +def validate_bundle(bundle: ContextBundle) -> None: + bundle_payload_size(bundle) + + +def bundle_payload_size(bundle: ContextBundle) -> int: + """Validate scalar contracts and measure the charged UTF-8 payload.""" + if ( + type(bundle) is not ContextBundle + or type(bundle.format_version) is not int + or bundle.format_version != 1 + or type(bundle.contexts) is not tuple + or len(bundle.contexts) > 16 + ): + raise ValueError("Invalid context bundle") + total = text_size(bundle.state, 128) + text_size(bundle.reason, 128) + if bundle.state not in BUNDLE_STATES: + raise ValueError("Invalid bundle state") + if bundle.reason and re.fullmatch(r"[a-z][a-z0-9_-]*", bundle.reason) is None: + raise ValueError("Invalid bundle reason code") + anchors = set() + for context in bundle.contexts: + if type(context) is not RecipeContext: + raise ValueError("Invalid recipe context") + size = text_size(context.anchor_id, 512) + if len(context.anchor_id) > 128: + raise ValueError("Oversized recipe anchor") + if not context.anchor_id or context.anchor_id in anchors: + raise ValueError("Duplicate or empty recipe anchor") + anchors.add(context.anchor_id) + size += text_size(context.role_provenance, 128) + if context.role_provenance != "inferred": + raise ValueError("Invalid declared role provenance") + for snapshot in (context.positive, context.negative): + if snapshot is None: + continue + if type(snapshot) is not PromptSnapshot: + raise ValueError("Invalid prompt snapshot") + for field in ( + snapshot.raw, + snapshot.effective, + snapshot.prepend, + snapshot.append, + ): + if type(field) is not TextEvidence: + raise ValueError("Invalid text evidence") + if ( + type(field.provenance) is not str + or type(field.source) is not str + or field.source not in EVIDENCE_SOURCES.get(field.provenance, ()) + ): + raise ValueError("Invalid evidence provenance or source") + if field.value is None: + if ( + field.provenance != "unknown" + or type(field.reason) is not str + or field.reason not in UNKNOWN_REASONS + ): + raise ValueError("Unknown evidence requires a reason code") + elif field.provenance == "unknown" or field.reason is not None: + raise ValueError("Known evidence requires a known provenance") + if field.value is not None: + size += text_size(field.value) + size += text_size(field.provenance, 128) + text_size(field.source, 128) + if field.reason is not None: + size += text_size(field.reason, 128) + if size > MAX_CONTEXT_BYTES: + raise ValueError("Context payload limit reached") + total += size + if total > MAX_ATTEMPT_BYTES: + raise ValueError("Attempt payload limit reached") + return total diff --git a/utils/execution_history.py b/utils/execution_history.py index c05c07b..086626a 100644 --- a/utils/execution_history.py +++ b/utils/execution_history.py @@ -1,5 +1,6 @@ """Own bounded queue-attempt capture without retaining host workflow payloads.""" +from collections import OrderedDict import atexit import threading import time @@ -12,6 +13,8 @@ except ImportError: from database import PromptDatabase from .execution_history_events import make_finish, make_start +from .execution_context_capture import ContextCapture, active_capture, unbind_current +from .execution_context_host import observe_cached_context, subscribe_executor_messages from .execution_history_writer import HistoryWriter from .logging_config import get_logger from .queue_observers import subscribe_dequeue, subscribe_task_done @@ -26,13 +29,27 @@ class HistoryRegistration: """Own subscriptions, one asynchronous writer and bounded attempt identities.""" def __init__( - self, queue: object, session_id: str, db_factory: Callable[[], PromptDatabase] + self, + queue: object, + session_id: str, + db_factory: Callable[[], PromptDatabase], + *, + high_water: int = -1, ) -> None: - self.writer = HistoryWriter(db_factory) + self.writer = HistoryWriter( + db_factory, on_invalidation=self._invalidate_capture + ) + self.capture = ContextCapture( + id(queue), self.writer.status, self.writer.context_budget + ) + self._capabilities = {} + self._accepted = OrderedDict() + self._high_water = high_water self._session_id = session_id self._lock = threading.Lock() self._active: dict[int, str] = {} self._disposed = False + self._capture_generation = None self._unsubscribers: list[Callable[[], None]] = [] try: self._unsubscribers.append( @@ -41,46 +58,99 @@ class HistoryRegistration: self._unsubscribers.append( subscribe_task_done(queue, "execution_history", self._task_done) ) + try: + self._unsubscribers.append( + subscribe_executor_messages( + f"execution_history_{id(self)}", self._executor_message + ) + ) + except (ImportError, AttributeError, TypeError, ValueError): + # Optional cache evidence must not prevent lifecycle tracking. + pass self.writer.start() atexit.register(self.dispose) except Exception: self.dispose() raise + def _executor_message(self, executor: object, args: tuple, kwargs: dict) -> None: + binding = active_capture() + if binding is None or binding[0] is not self.capture: + return + event = args[0] if args else kwargs.get("event") + data = args[1] if len(args) > 1 else kwargs.get("data") + observe_cached_context(executor, event, data) + + def _invalidate_capture(self) -> None: + self.capture.revoke_all() + 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()): + if status.capture_generation != self._capture_generation or any( + value != epoch for value in self._active.values() + ): self._active.clear() + self._capabilities.clear() + self.capture.revoke_all() + self._capture_generation = status.capture_generation return epoch def _dequeue(self, result: object) -> None: try: with self._lock: + epoch = self._ready_epoch() + command = ( + make_start( + self._session_id, result, epoch, time.time_ns() // 1_000_000 + ) + if epoch + else None + ) + accepted = ( + self._accepted.get(command.queue_item_id) + if command is not None + else None + ) + if command is not None and accepted == ( + command.job_id, + command.storage_epoch, + ): + return + unbind_current() if self._disposed: return - epoch = self._ready_epoch() - if epoch is None: + if command is None or command.queue_item_id in self._active: 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 + fresh = command.queue_item_id > self._high_water + self._high_water = max(self._high_water, command.queue_item_id) if self.writer.submit(command): - self._active[item_id] = epoch + self._accepted[command.queue_item_id] = ( + command.job_id, + command.storage_epoch, + ) + if len(self._accepted) > 128: + self._accepted.popitem(last=False) + self._active[command.queue_item_id] = epoch + if not fresh: + return + queue_item = result[0] + graph = queue_item[2] if len(queue_item) > 2 else None + outputs = queue_item[4] if len(queue_item) > 4 else None + cap = self.capture.begin( + command, self.writer.status().capture_generation, graph, outputs + ) + self._capabilities[command.queue_item_id] = cap except Exception: + unbind_current() self.writer.record_gap() def _task_done(self, args: tuple, kwargs: dict) -> None: + sealed = None try: with self._lock: if self._disposed: @@ -90,6 +160,11 @@ class HistoryRegistration: self.writer.record_gap() return item_id = args[0] if args else kwargs.get("item_id") + if type(item_id) is int: + cap = self._capabilities.pop(item_id, None) + self._active.pop(item_id, None) + if cap is not None: + sealed = self.capture.seal(cap) status = args[2] if len(args) > 2 else kwargs.get("status") command = make_finish( self._session_id, @@ -101,10 +176,16 @@ class HistoryRegistration: if command is None: self.writer.record_gap() return - self._active.pop(command.queue_item_id, None) - self.writer.submit(command) + if sealed is not None: + self.writer.submit_finalization(command, sealed) + sealed = None + else: + self.writer.submit(command) except Exception: self.writer.record_gap() + finally: + if sealed is not None and sealed.lease is not None: + sealed.lease.release() def dispose(self, timeout: float = 2.0) -> bool: """Stop capture, release subscriptions and drain the owned writer boundedly.""" @@ -112,6 +193,8 @@ class HistoryRegistration: if not self._disposed: self._disposed = True self._active.clear() + self._capabilities.clear() + self.capture.revoke_all() for unsubscribe in self._unsubscribers: unsubscribe() self._unsubscribers.clear() @@ -137,7 +220,12 @@ def register_history_hook( previous = state["registration"] if previous is not None and not previous._disposed: return previous - registration = HistoryRegistration(queue, state["session_id"], db_factory) + registration = HistoryRegistration( + queue, + state["session_id"], + db_factory, + high_water=previous._high_water if previous is not None else -1, + ) state["registration"] = registration return registration except Exception: diff --git a/utils/execution_history_writer.py b/utils/execution_history_writer.py index 3cba705..b8583a2 100644 --- a/utils/execution_history_writer.py +++ b/utils/execution_history_writer.py @@ -1,11 +1,13 @@ """A bounded, single-worker transport for best-effort execution history.""" +from contextlib import contextmanager from dataclasses import dataclass +from functools import partial import math import queue import threading import time -from typing import Callable +from typing import Callable, Iterator try: from ..database import PromptDatabase @@ -20,7 +22,14 @@ except ImportError: subscribe_history_restore, ) -from .execution_history_events import RunCommand, RunStart, encode_command +from .execution_context_budget import ( + SHARED_CONTEXT_BUDGET, + ContextBudget, + PayloadLease, + SealedContext, +) +from .execution_contexts import ContextBundle, bundle_payload_size +from .execution_history_events import RunCommand, RunFinish, RunStart, encode_command from .logging_config import get_logger _logger = get_logger("prompt_manager.execution_history_writer") @@ -36,6 +45,13 @@ class HistoryWriterStatus: pending: int recording_gaps: int alive: bool + capture_generation: int + + +@dataclass(frozen=True) +class _Finalization: + command: RunFinish + context: SealedContext class HistoryWriter: @@ -46,12 +62,21 @@ class HistoryWriter: db_factory: Callable[[], PromptDatabase], *, clock: Callable[[], float] = time.monotonic, + context_budget: ContextBudget | None = None, + on_invalidation: Callable[[], None] | None = None, ) -> None: + self.context_budget = ( + SHARED_CONTEXT_BUDGET if context_budget is None else context_budget + ) + self._on_invalidation = on_invalidation self._db_factory = db_factory self._clock = clock + self._context_commit_gate = threading.Lock() self._lock = threading.Lock() self._wake = threading.Event() - self._queue: queue.Queue[RunCommand] = queue.Queue(maxsize=_MAX_PENDING) + self._queue: queue.Queue[RunCommand | _Finalization] = queue.Queue( + maxsize=_MAX_PENDING + ) self._thread: threading.Thread | None = None self._epoch: str | None = None self._accepting = False @@ -94,6 +119,69 @@ class HistoryWriter: self._wake.set() return accepted + @staticmethod + def _release_context(context: SealedContext) -> None: + if isinstance(context, SealedContext) and type(context.lease) is PayloadLease: + context.lease.release() + + def _consume_context_locked(self, context: SealedContext) -> SealedContext: + try: + if type(context) is not SealedContext: + raise ValueError("Invalid sealed context") + context.validate_header() + if context.bundle is None: + if context.lease is not None: + raise ValueError("Unexpected payload lease") + return context + size = bundle_payload_size(context.bundle) + if ( + context.capture_generation != self._generation + or context.capture_generation is None + or type(context.lease) is not PayloadLease + or context.lease._pool is not self.context_budget + or context.lease.byte_count < size + or context.state != context.bundle.state + or context.reason != context.bundle.reason + ): + raise ValueError("Invalid context ownership or generation") + owner = context.lease.transfer() + return SealedContext( + context.bundle, + owner, + context.state, + context.reason, + context.capture_generation, + ) + except ValueError: + self._release_context(context) + return SealedContext(None, None, "unavailable", "invalid_evidence") + + def submit_finalization(self, command: RunFinish, context: SealedContext) -> bool: + try: + if type(command) is not RunFinish: + raise ValueError("Finalization requires finish") + encode_command(command) + except ValueError: + self._release_context(context) + self.record_gap() + return False + with self._lock: + if ( + not self._accepting + or command.storage_epoch != self._epoch + or self._pending >= _MAX_PENDING + ): + self._release_context(context) + self._record_gap_locked() + accepted = False + else: + owned = self._consume_context_locked(context) + self._queue.put_nowait(_Finalization(command, owned)) + self._pending += 1 + accepted = True + self._wake.set() + return accepted + def _record_gap_locked(self) -> None: self._gaps = min(_MAX_GAPS, self._gaps + 1) @@ -103,13 +191,36 @@ class HistoryWriter: self._record_gap_locked() self._wake.set() + @contextmanager + def _context_generation_guard( + self, generation: int, epoch: str, capture_generation: int | None + ) -> Iterator[bool]: + with self._context_commit_gate: + with self._lock: + current = ( + generation == self._generation + and epoch == self._epoch + and (capture_generation is None or capture_generation == generation) + ) + yield current + + def _notify_invalidation(self) -> None: + if self._on_invalidation is not None: + try: + self._on_invalidation() + except Exception: + # Observer failure must not block revocation of stale captures. + pass + 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() + with self._context_commit_gate: + with self._lock: + self._accepting = False + self._epoch = None + self._generation += 1 + self._next_probe = self._clock() + _PROBE_INTERVAL + self._discard_locked() + self._notify_invalidation() self._wake.set() def status(self) -> HistoryWriterStatus: @@ -120,6 +231,7 @@ class HistoryWriter: self._pending, self._gaps, self._thread is not None and self._thread.is_alive(), + self._generation, ) def close(self, timeout: float = 2.0) -> bool: @@ -132,17 +244,31 @@ class HistoryWriter: self._close_deadline = deadline self._accepting = False thread = self._thread + self._notify_invalidation() self._wake.set() if thread is not None: - thread.join(max(0.0, deadline - time.monotonic())) + while thread.is_alive(): + remaining = deadline - time.monotonic() + if remaining <= 0: + break + thread.join(remaining) + if thread.is_alive(): + remaining = deadline - time.monotonic() + if remaining > 0: + time.sleep(min(remaining, 0.001)) + with self._lock: + if self._closing and time.monotonic() >= self._close_deadline: + self._discard_locked() return thread is None or not thread.is_alive() def _discard_locked(self) -> None: while True: try: - self._queue.get_nowait() + item = self._queue.get_nowait() except queue.Empty: return + if isinstance(item, _Finalization): + self._release_context(item.context) self._pending -= 1 self._record_gap_locked() @@ -200,6 +326,10 @@ class HistoryWriter: self._record_gap_locked() if command is not None: failed = False + context = None + if isinstance(command, _Finalization): + context = command.context + command = command.command try: # Restore may have reopened this thread's managed connection. db.model.get_connection().execute("PRAGMA busy_timeout=250") @@ -207,12 +337,49 @@ class HistoryWriter: db.write_execution_start(command) else: db.write_execution_finish(command) + if context is not None: + with self._lock: + current = ( + generation == self._generation + and command.storage_epoch == self._epoch + and ( + context.capture_generation is None + or context.capture_generation + == self._generation + ) + ) + if current: + try: + db.write_execution_contexts_for_finish( + command, + ( + context.bundle + if context.bundle is not None + else ContextBundle( + (), context.state, context.reason + ) + ), + generation_guard=partial( + self._context_generation_guard, + generation, + command.storage_epoch, + context.capture_generation, + ), + ) + except ValueError: + # Validation/collision cannot revoke the outcome. + with self._lock: + self._record_gap_locked() 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 + finally: + if context is not None: + self._release_context(context) + context = None with self._lock: self._pending -= 1 if failed: @@ -222,6 +389,8 @@ class HistoryWriter: self._generation += 1 self._next_probe = self._clock() + _PROBE_INTERVAL self._discard_locked() + if failed: + self._notify_invalidation() with self._lock: gaps = self._gaps has_work = self._epoch is not None and not self._queue.empty() @@ -239,6 +408,7 @@ class HistoryWriter: self._accepting = False self._epoch = None self._discard_locked() + self._notify_invalidation() if unsubscribe is not None: unsubscribe() if db is not None: