feat(history): bind exact runtime context to execution attempts

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(

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(

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(

tests/test_execution_context_capture.py

@@ -0,0 +1,475 @@
+"""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 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_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_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)

tests/test_execution_context_graph.py

@@ -230,3 +230,17 @@ class ProjectionTests(unittest.TestCase):
             {"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, ())

tests/test_execution_history_hooks.py

@@ -356,7 +356,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 +448,180 @@ class HistoryHooksTests(unittest.TestCase):
         self.drain(registration)
         self.assertEqual(seen, [JOB])
         self.assertEqual(len(self.rows()), 1)
+
+    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]))
+            release.set()
+
+    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)

tests/test_execution_history_writer.py

@@ -905,3 +905,26 @@ class HistoryWriterTests(unittest.TestCase):
 
 if __name__ == "__main__":
     unittest.main()
+
+    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)])

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)

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)

utils/execution_context_capture.py

@@ -0,0 +1,401 @@
+"""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 = ""
+
+
+_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(
+                        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 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]
+                    if role.inferred is not None:
+                        s = role.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

utils/execution_context_graph.py

@@ -48,6 +48,10 @@ class _Unknown(ValueError):
     pass
 
 
+class _Limit(_Unknown):
+    pass
+
+
 def _identifier(value: object) -> str:
     if type(value) is int:
         if abs(value) > 10**127:
@@ -79,6 +83,7 @@ class _Projector:
     def __init__(self, graph: dict) -> None:
         self.graph = graph
         self.visits = 0
+        self.retained = 0
         self.transformed = False
         self.selected_seen = set()
 
@@ -179,24 +184,29 @@ class _Projector:
             )
             raw = child.effective.value
             prepend, append = before.effective.value, after.effective.value
-            effective = " ".join(
-                part for part in (prepend.strip(), raw, append.strip()) if part
-            )
+            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 = PromptSnapshot(
-                child.effective,
-                TextEvidence(effective, "inferred", "graph_semantics", None),
-                before.effective,
-                after.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 = PromptSnapshot(
-                child.effective,
-                child.effective,
-                TextEvidence("", "inferred", "graph_literal", None),
-                TextEvidence("", "inferred", "graph_literal", None),
+            snapshot = self.retain(
+                PromptSnapshot(
+                    child.effective,
+                    child.effective,
+                    TextEvidence("", "inferred", "graph_literal", None),
+                    TextEvidence("", "inferred", "graph_literal", None),
+                )
             )
         else:
             snapshot = child
@@ -254,8 +264,23 @@ class _Projector:
             return tuple(sources)
         return ()
 
-    @staticmethod
+    def retain(self, snapshot: PromptSnapshot) -> PromptSnapshot:
+        size = sum(
+            text_size(field.value) + 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,
@@ -263,19 +288,21 @@ class _Projector:
         composed: bool = False,
         raw_composed: bool = False,
     ) -> PromptSnapshot:
-        return PromptSnapshot(
-            *(
-                TextEvidence(
-                    value,
-                    "inferred",
-                    (
-                        "graph_semantics"
-                        if (composed and index == 1 or raw_composed and index == 0)
-                        else "graph_literal"
-                    ),
-                    None,
+        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))
                 )
-                for index, value in enumerate((raw, effective, prepend, append))
             )
         )
 
@@ -288,6 +315,8 @@ class _Projector:
             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(

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,7 @@ except ImportError:
     from database import PromptDatabase
 
 from .execution_history_events import make_finish, make_start
+from .execution_context_capture import ContextCapture, unbind_current
 from .execution_history_writer import HistoryWriter
 from .logging_config import get_logger
 from .queue_observers import subscribe_dequeue, subscribe_task_done
@@ -26,13 +28,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(
@@ -47,40 +63,76 @@ class HistoryRegistration:
             self.dispose()
             raise
 
+    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 +142,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 +158,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 +175,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 +202,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:

utils/execution_history_writer.py

@@ -63,10 +63,12 @@ class HistoryWriter:
         *,
         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()
@@ -202,6 +204,13 @@ class HistoryWriter:
                 )
             yield current
 
+    def _notify_invalidation(self) -> None:
+        if self._on_invalidation is not None:
+            try:
+                self._on_invalidation()
+            except Exception:
+                pass
+
     def invalidate(self) -> None:
         with self._context_commit_gate:
             with self._lock:
@@ -210,6 +219,7 @@ class HistoryWriter:
                 self._generation += 1
                 self._next_probe = self._clock() + _PROBE_INTERVAL
                 self._discard_locked()
+        self._notify_invalidation()
         self._wake.set()
 
     def status(self) -> HistoryWriterStatus:
@@ -233,6 +243,7 @@ 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()))
@@ -369,6 +380,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()
@@ -386,6 +399,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: