EricRollei-AAA_Metadata_System/documentation/runtime_hook_implementation_plan.md

# Runtime Hook Implementation Plan

## Overview

We will introduce an optional execution-hook capture layer inspired by
[`ComfyUI_SaveImageWithMetaDataUniversal`](https://github.com/xxmjskxx/ComfyUI_SaveImageWithMetaDataUniversal)
to gather runtime-resolved parameters that our static `WorkflowParser` cannot
always see. The hook facility must remain opt-in, conflict-aware, and safe to
use alongside other custom node packs.

## Objectives

- Collect additional runtime data (resolved sampler settings, dynamic inputs,
  execution order, timestamps) without destabilising ComfyUI.
- Merge hook-captured values into our canonical metadata structure while
  preserving schema integrity.
- Ensure all additional data can be embedded in images (PNG/JPEG/WEBP), XMP
  sidecars, and optional JSON sidecars.
- Provide clear toggles, logging, and documentation so users understand the
  trade-offs of enabling hooks.

## Out of Scope (for now)

- Replacing the current parser-based discovery (parser stays as baseline).
- Enabling hooks by default. The feature ships disabled and advertised as
  advanced.
- Building user-editable declarative capture rules; that refactor remains a
  separate initiative.

## Safety Principles

1. **Opt-in only**: require explicit configuration (environment variable, node
   checkbox, or settings JSON) before patching ComfyUI internals.
2. **Minimal work inside hooks**: hooks collect references and lightweight
   values only; heavy processing happens after the run finishes.
3. **Scope isolation**: maintain state keyed by `prompt_id` so queued/batched
   runs cannot bleed data across sessions.
4. **Conflict detection**: if another extension has already wrapped the target
   dispatcher methods, log a warning and skip hook registration.
5. **Fail-safe design**: any exception inside hook code is caught, logged once,
   and auto-disables the hook for the remainder of the session.
6. **Graceful teardown**: provide a function to unhook and restore the original
   dispatcher when ComfyUI shuts down or the feature is disabled.

## Hook Strategy

### Targeted Dispatcher Patches

- Mirror the hook points used upstream: `PromptExecutor.pre_execute` and
  `PromptExecutor.pre_get_input_data` (confirm the exact import paths on the
  current ComfyUI version).
- Create a module `eric_metadata/hooks/runtime_capture.py` that exports
  `enable_hooks()` and `disable_hooks()` helpers.
- Store original functions in module-level variables so they can be restored
  later.

### Session Context Management

- Introduce a `HookSession` dataclass containing:
  - `prompt_id`
  - `prompt_json`
  - `start_ts`
  - `executed_nodes` (ordered list of `(node_id, class_type, ts)`)
  - `resolved_inputs` (map of node_id to resolved parameter values)
- Maintain an in-memory registry `{prompt_id: HookSession}` inside the hook
  module. Remove entries once the save node finalises metadata for that
  prompt.
- When the hook sees `SaveImageWithMetadata` class IDs, mark the active save
  node so the post-save integration knows which session to consume.

### Enabling and Disabling

- Activation flow:
  1. On node initialise (or top-level package `__init__`), read configuration
     (env var `AAA_METADATA_ENABLE_HOOKS`, config file, or ComfyUI settings).
  2. If enabled, call `enable_hooks()`; otherwise keep parser-only mode.
  3. Emit a debug log summarising whether hooks are active and if conflicts were
     detected.
- Deactivation flow: provide `disable_hooks()` that restores the dispatcher and
  clears state.

## Data Capture and Merge

### Data Collected via Hooks

- Full prompt JSON as resolved at execution time (guaranteed to match the run).
- Execution order of nodes leading up to the save node.
- Runtime input values (e.g., sampler steps after conditional logic, actual
  checkpoint file chosen by switch nodes, live LoRA strengths).
- Optional timings (start/finish timestamps per node) for provenance.

### Integration with Canonical Metadata

- Extend our metadata schema with a `runtime` section under `ai_info`, for
  example, `ai_info.runtime.sampler.steps_resolved`,
  `ai_info.runtime.nodes_executed[]`, `ai_info.runtime.inputs.<node_id>`.
- During save:
  1. Run the existing `WorkflowParser` to populate canonical fields.
  2. Query the `HookSession` for the current `prompt_id`.
  3. Merge the hook data, favouring hook values when both sources exist (they
     represent runtime truth). Keep parser output as fallback for missing items.
  4. Record a provenance flag `provenance.capture_mode = "parser+hooks"` (or
     similar) so downstream tooling can detect that runtime data was available.
  5. If runtime payloads exceed our safe inline size budget, emit a compact
     pointer (e.g., `ai_info.runtime.sidecar_ref`) to a JSON sidecar that holds
     the full structure. Include the same pointer inside the PNG/JPEG metadata
     for discoverability.

### Storage and Serialization Updates

- **Canonical JSON (PNG/sidecar)**: add the `ai_info.runtime` block and ensure
  it serialises alongside existing fields. Compress large payloads when
  embedding in PNG iTXt chunks.
- **XMP mapping**: surface key runtime parameters (sampler, steps, cfg, seed,
  checkpoint name/hash) using the existing `aaa:` namespace. Keep heavyweight
  arrays (execution trace, resolved inputs) in the canonical JSON only.
- **JPEG fallback**: when EXIF size limits are hit, include the runtime pointer
  in the minimal metadata payload and drop the heavy structures into the
  sidecar.

## Implementation Phases

### Phase 0 – Discovery and Design (1–2 days)

1. Confirm current ComfyUI dispatcher signatures and the exact hook points used
   by `ComfyUI_SaveImageWithMetaDataUniversal`.
2. Document any differences between their fork and upstream ComfyUI that we
   must account for (parameter names, execution context objects, etc.).
3. Finalise schema extensions (`ai_info.runtime`, provenance flag, sidecar
   pointer field names).

### Phase 1 – Hook Infrastructure (2–3 days)

1. Implement `runtime_capture.py` with enable/disable helpers, conflict
   detection, and per-prompt session registry.
2. Add configuration surface (env var + node UI toggle) and debug logging.
3. Write unit tests covering:
   - Hook registration/unregistration restores originals.
   - Conflict detection triggers when dispatcher already patched.
   - Session lifecycle (create on pre_execute, cleanup on teardown).

### Phase 2 – Data Merge and Serialization (2–4 days)

1. Update the save node pipeline to request hook data, merge into canonical
   metadata, and set provenance flags.
2. Extend serializers (PNG, XMP, JPEG fallback) to embed runtime values and
   pointer references.
3. Update documentation of metadata schema and add examples demonstrating the
   new fields in embedded JSON/XMP.

### Phase 3 – Validation and Hardening (1–3 days)

1. Run manual workflows covering:
   - Standard linear graph (baseline comparison).
   - Conditional sampler selection / switch nodes.
   - Queue of multiple prompts to validate isolation.
2. Measure performance impact with hooks enabled vs disabled.
3. Add integration tests ensuring parser-only mode still works when hooks are
   disabled.

## Validation and Testing Strategy

- **Unit tests**: hook enable/disable, conflict detection, session registry,
  merge logic (`runtime` block assembled correctly when both parser and hook
  provide data).
- **Integration tests**: simulated ComfyUI execution harness that invokes the
  hook functions around a stubbed save node to verify captured values end up in
  the saved metadata structure.
- **Regression suite**: re-run existing metadata writer tests to make sure
  nothing regresses when hooks are off.
- **Manual QA**: inspect saved PNG/JPEG/XMP files with external tools to
  confirm new fields are present and correctly mapped.

## Risks and Mitigations

- **Hook conflicts with other packs**: detect patched dispatcher signatures,
  warn user, auto-disable our hooks.
- **Runtime exceptions inside hooks**: wrap hook bodies in try/except, log a
  concise error (once per session), disable hooks for safety.
- **Performance overhead**: keep hook-side work minimal and defer heavy tasks
  to post-save. Provide a debug metric in logs for hook execution time.
- **Schema bloat**: document new fields, cap inline payload size, and favour
  sidecars for verbose data.
- **User surprise**: update README/CHANGELOG, highlight that hooks are optional
  and how to enable them.

## Open Questions

1. Should we expose a per-run toggle in the save node UI (e.g., "Prefer runtime
   capture if available") in addition to global config?
2. What is the maximum acceptable size for the embedded `ai_info.runtime`
   section before we force a sidecar (e.g., 32 KB)?
3. Do we want to log sampler/controlnet resolution differences when parser and
   hook disagree, to aid debugging?
4. How should we present hook status to users in the UI (icon, text label,
   metadata note)?

---

**Next step suggestion**: review these open questions, decide on schema field
names (`capture_mode`, `runtime.sidecar_ref`, etc.), and then green-light Phase
1 implementation.