# Refactor Plan 2 — ComfyUI-Distributed

All paths relative to `custom_nodes/ComfyUI-Distributed/`.
Read every file listed under "Read first" before making any edits.
After every phase, run `python -c "import custom_nodes.ComfyUI-Distributed"` (or equivalent import check) to confirm nothing is broken.

---

## Phase 1 — Correctness / Data Integrity (do these first)

### Task 1.1 — Fix lock contention in USDU requeue path

**Problem:** `_check_and_requeue_timed_out_workers` in `utils/usdu_managment.py` calls
`await probe_worker(...)` (a network request) while holding `distributed_tile_jobs_lock`.
Any concurrent heartbeat, submit_tiles, or request_image that also needs the lock is blocked
for the full probe timeout.

**Read first:** `utils/usdu_managment.py` (find `_check_and_requeue_timed_out_workers`),
`utils/network.py` (confirm `probe_worker` signature).

**Changes:**

In `_check_and_requeue_timed_out_workers`:

1. **Snapshot under lock, then release.** Collect the list of workers to probe while holding
   the lock, then exit the lock context before any `await probe_worker(...)` call.

   Before (pseudocode):
   ```python
   async with prompt_server.distributed_tile_jobs_lock:
       job_data = ...
       for worker_id, last_seen in timed_out:
           result = await probe_worker(url, timeout=...)   # WRONG: network call under lock
           if not result:
               requeue(...)
   ```

   After:
   ```python
   # Step 1: snapshot under lock
   async with prompt_server.distributed_tile_jobs_lock:
       job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
       if not isinstance(job_data, BaseJobState):
           return 0
       timed_out_snapshot = [
           (wid, ts) for wid, ts in job_data.worker_status.items()
           if (now - float(ts)) > HEARTBEAT_TIMEOUT
       ]
       assigned_snapshot = dict(job_data.assigned_to_workers)

   # Step 2: probe outside lock
   workers_to_requeue = []
   for worker_id, last_seen in timed_out_snapshot:
       task_ids = assigned_snapshot.get(worker_id, [])
       wrec = _find_worker_record(worker_id)
       if wrec:
           result = await probe_worker(build_worker_url(wrec), timeout=2.0)
           if result is None:
               workers_to_requeue.append((worker_id, task_ids))
       else:
           workers_to_requeue.append((worker_id, task_ids))

   # Step 3: re-acquire lock to apply requeue decisions
   if not workers_to_requeue:
       return 0
   async with prompt_server.distributed_tile_jobs_lock:
       job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
       if not isinstance(job_data, BaseJobState):
           return 0
       requeued = 0
       for worker_id, task_ids in workers_to_requeue:
           for task_id in task_ids:
               await job_data.pending_tasks.put(task_id)
               requeued += 1
           job_data.assigned_to_workers.pop(worker_id, None)
           job_data.worker_status.pop(worker_id, None)
       return requeued
   ```

2. Add a private helper `_find_worker_record(worker_id)` that loads config and returns the
   worker dict (or `None`). This keeps the probe-outside-lock section clean.

**Verify:** The function should now only hold the lock during pure dict reads/writes.
No `await` expressions should appear inside any `async with distributed_tile_jobs_lock:` block
in this function.

---

### Task 1.2 — Fix non-idempotent client retries

**Problem:** `web/apiClient.js` retries every request unconditionally. Mutating endpoints
(`/distributed/queue`, launch, stop, update, delete, startTunnel, stopTunnel) can execute
twice on a transient network error, causing duplicate jobs or double-start/stop.

**Read first:** `web/apiClient.js` (full file).

**Changes:**

1. **Add a `retry` option to the internal `request` function.** Default `true` for safe reads,
   `false` for all mutating calls.

   ```javascript
   const request = async (endpoint, options = {}, { retries = TIMEOUTS.MAX_RETRIES, retry = true } = {}) => {
       const maxAttempts = retry ? retries : 1;
       let lastError;
       let delay = TIMEOUTS.RETRY_DELAY;
       for (let attempt = 0; attempt < maxAttempts; attempt++) {
           try {
               const response = await fetch(...);
               if (!response.ok) { ... }
               return await response.json();
           } catch (error) {
               lastError = error;
               if (attempt < maxAttempts - 1) {
                   await new Promise(resolve => setTimeout(resolve, delay));
                   delay *= 2;
               }
           }
       }
       throw lastError;
   };
   ```

2. **Set `retry: false` for all mutating endpoints:**
   - `launchWorker` — `request('/distributed/launch_worker', { method: 'POST', ... }, { retry: false })`
   - `stopWorker` — same
   - `queueDistributed` — same
   - `updateWorker` — same
   - `deleteWorker` — same
   - `updateSetting` — same
   - `updateMaster` — same
   - `startTunnel` — same
   - `stopTunnel` — same
   - `dispatchToWorker` — same

3. **Add idempotency key support for `queueDistributed`.** The caller already generates a
   `trace_execution_id`; pass it as an `X-Idempotency-Key` header so a future backend can
   deduplicate if desired.

   ```javascript
   queueDistributed: (payload) =>
       request('/distributed/queue', {
           method: 'POST',
           headers: {
               'Content-Type': 'application/json',
               ...(payload.trace_execution_id
                   ? { 'X-Idempotency-Key': payload.trace_execution_id }
                   : {}),
           },
           body: JSON.stringify(payload),
       }, { retry: false }),
   ```

4. **Leave `probeWorker`, `checkStatus`, `checkMultipleStatuses`, `getManagedWorkers`,
   `getWorkerLog`, `getTunnelStatus`, `getNetworkInfo` as `retry: true`** — these are
   read-only and safe to retry.

**Verify:** `queueDistributed`, `launchWorker`, `stopWorker`, `startTunnel`, `stopTunnel`
all pass `{ retry: false }`. No mutating call uses the default retry loop.

---

### Task 1.3 — Make config writes atomic and race-safe

**Problem:** `utils/config.py` `save_config()` writes directly to the config file with
`json.dump(..., f)`. A crash mid-write corrupts the file. Concurrent requests from
`api/config_routes.py` do independent load–modify–save cycles with no locking, creating
a race condition.

**Read first:** `utils/config.py` (full), `api/config_routes.py` (full).

**Changes in `utils/config.py`:**

1. Add a module-level `asyncio.Lock`:
   ```python
   import asyncio
   _config_lock = asyncio.Lock()
   ```

2. Rewrite `save_config` to write atomically using a temp file + `os.replace`:
   ```python
   def save_config(config):
       config_path = _get_config_path()
       tmp_path = config_path + '.tmp'
       try:
           with open(tmp_path, 'w') as f:
               json.dump(config, f, indent=2)
           os.replace(tmp_path, config_path)
       except Exception:
           try:
               os.unlink(tmp_path)
           except OSError:
               pass
           raise
   ```

3. Expose a new async context manager `config_transaction()` for use by routes:
   ```python
   from contextlib import asynccontextmanager

   @asynccontextmanager
   async def config_transaction():
       """Acquire config lock, yield loaded config, save on clean exit."""
       async with _config_lock:
           config = load_config()
           yield config
           save_config(config)
   ```

**Changes in `api/config_routes.py`:**

4. Replace every `load_config()` / `save_config(config)` pair in route handlers with
   `async with config_transaction() as config:`. Each handler becomes:
   ```python
   async with config_transaction() as config:
       # mutate config in place
       # no explicit save needed — context manager handles it
   ```
   Apply to: `update_worker_endpoint`, `delete_worker_endpoint`, `update_setting_endpoint`,
   `update_master_endpoint`, and the bulk `POST /distributed/config` handler.

5. The read-only `GET /distributed/config` route does not need the lock — leave it using
   `load_config()` directly.

**Verify:** No `save_config` call outside of the `config_transaction` context manager in
route handlers. Grep for `save_config(` in `api/config_routes.py` and confirm zero results.

---

## Phase 2 — Structural Splits

### Task 2.1 — Split `utils/usdu_managment.py` god module

**Problem:** `utils/usdu_managment.py` mixes job-state CRUD, timeout/requeue policy,
multipart payload parsing, conditioning utilities, and HTTP route registrations in one ~630-line file.

**Read first:** `utils/usdu_managment.py` (full), `upscale/job_state.py`, `upscale/job_models.py`,
`api/worker_routes.py`, `upscale/worker_comms.py` (check `_send_heartbeat_to_master` import).

**Target structure:**

```
upscale/
  job_store.py        ← job init + ensure_tile_jobs_initialized + MAX_PAYLOAD_SIZE
  job_timeout.py      ← heartbeat check, requeue logic (_check_and_requeue_timed_out_workers)
  payload_parsers.py  ← _parse_tiles_from_form, any other multipart parsers
api/
  usdu_routes.py      ← all HTTP route handlers from usdu_managment.py
utils/
  usdu_managment.py   ← thin shim: re-exports for backward compat + _send_heartbeat_to_master
```

**Step-by-step:**

1. **Create `upscale/job_store.py`.**
   Move from `usdu_managment.py`:
   - `ensure_tile_jobs_initialized()`
   - `_init_job_queue(multi_job_id, mode, batch_size, num_tiles_per_image, batched_static)`
   - `init_dynamic_job(...)`
   - `init_static_job_batched(...)`
   - `_drain_results_queue(multi_job_id)`
   - The `MAX_PAYLOAD_SIZE` constant
   Imports needed: `asyncio`, `server`, `upscale.job_models`, `utils.logging`.

2. **Create `upscale/job_timeout.py`.**
   Move from `usdu_managment.py` (after Task 1.1 changes are applied):
   - `_check_and_requeue_timed_out_workers(multi_job_id, batch_size)` (already refactored in Task 1.1)
   - `_find_worker_record(worker_id)` (new helper from Task 1.1)
   Imports needed: `utils.network`, `utils.config`, `utils.constants`, `upscale.job_models`,
   `upscale.job_store` (for `ensure_tile_jobs_initialized`).

3. **Create `upscale/payload_parsers.py`.**
   Move from `usdu_managment.py`:
   - `_parse_tiles_from_form(request)` and any helper it uses
   Imports needed: `aiohttp`, `PIL`, `utils.image`.

4. **Create `api/usdu_routes.py`.**
   Move from `usdu_managment.py` all `@server.PromptServer.instance.routes.*` decorated
   functions:
   - `heartbeat_endpoint` (`POST /distributed/heartbeat`)
   - `submit_tiles_endpoint` (`POST /distributed/submit_tiles`)
   - `submit_image_endpoint` (`POST /distributed/submit_image`)
   - `request_image_endpoint` (`POST /distributed/request_image`)
   - `job_status_endpoint` (`GET /distributed/job_status`)
   Import from new modules above instead of `usdu_managment`.

5. **Rewrite `utils/usdu_managment.py` as a thin shim.**
   Keep only:
   - `_send_heartbeat_to_master(...)` (used by `upscale/worker_comms.py`)
   - `MAX_PAYLOAD_SIZE` re-export from `upscale.job_store`
   - Re-exports of `_check_and_requeue_timed_out_workers` and `ensure_tile_jobs_initialized`
     for any callers that import from here directly.

   ```python
   # utils/usdu_managment.py — backward-compat shim
   from ..upscale.job_store import (
       MAX_PAYLOAD_SIZE,
       ensure_tile_jobs_initialized,
       _init_job_queue,
       init_dynamic_job,
       init_static_job_batched,
   )
   from ..upscale.job_timeout import _check_and_requeue_timed_out_workers

   async def _send_heartbeat_to_master(...):
       ...  # keep this here — it's worker-side only
   ```

6. **Register `api/usdu_routes.py` in `__init__.py` or `distributed.py`.**
   Find where `usdu_managment` is currently imported to trigger route registration, and
   add an import of `api.usdu_routes` in the same place.

7. **Update all import sites.** Grep for `from ..utils.usdu_managment import` and
   `from .utils.usdu_managment import` across all files. Update each to import from the
   appropriate new module.

**Verify:** `python -m py_compile` on all modified files. Grep confirms no direct
`usdu_managment` imports remain except the shim itself and `worker_comms.py`
(which imports `_send_heartbeat_to_master`).

---

### Task 2.2 — Decompose `CloudflareTunnelManager`

**Problem:** `utils/cloudflare.py` is a 387-line class handling binary management, state
persistence, log streaming, process lifecycle, and status reporting.

**Read first:** `utils/cloudflare.py` (full), `api/tunnel_routes.py`, `distributed.py`
(check where `CloudflareTunnelManager` is imported/instantiated).

**Target structure:**

```
utils/
  cloudflare/
    __init__.py         ← re-exports CloudflareTunnelManager for backward compat
    binary.py           ← binary discovery and download
    state.py            ← config persistence, previous-host restoration
    process_reader.py   ← thread-based output reader, URL capture, log buffer
    tunnel.py           ← CloudflareTunnelManager: start/stop lifecycle, status
```

**Step-by-step:**

1. **Create `utils/cloudflare/` package.**

2. **Create `utils/cloudflare/binary.py`.**
   Move from `cloudflare.py`:
   - `_get_cloudflared_dir()`, `_get_binary_path()` (or equivalent path helpers)
   - `_get_platform_binary_name()`
   - `_download_cloudflared()` / `ensure_binary()` — everything related to finding or
     downloading the `cloudflared` binary.
   Export a single `ensure_binary() -> str` function that returns the path to the binary.

3. **Create `utils/cloudflare/state.py`.**
   Move from `cloudflare.py`:
   - `_save_tunnel_url_to_config(url)` / `_restore_previous_master_host()` /
     `_clear_tunnel_state()` — all code that reads/writes the config file to persist or
     restore tunnel URLs and master host settings.
   This module imports from `utils.config`.

4. **Create `utils/cloudflare/process_reader.py`.**
   Move from `cloudflare.py`:
   - The `_read_output` thread / `_OutputReader` thread class that reads `cloudflared`
     stdout line-by-line, extracts the tunnel URL via regex, appends to `recent_logs`,
     and sets the URL event.
   Export: `ProcessReader` class (or equivalent) with `start(process)` / `stop()` /
   `get_url()` / `get_recent_logs()` methods. Keep the `LOG_BUFFER_SIZE` constant here
   (or move to `utils/constants.py` per Task 5.2).

5. **Create `utils/cloudflare/tunnel.py`.**
   This keeps the `CloudflareTunnelManager` class but now delegates to the modules above:
   ```python
   from .binary import ensure_binary
   from .state import save_tunnel_url, restore_previous_host, clear_tunnel_state
   from .process_reader import ProcessReader

   class CloudflareTunnelManager:
       def __init__(self): ...
       async def start(self): ...    # calls ensure_binary, starts process, starts ProcessReader
       async def stop(self): ...     # stops ProcessReader, stops process, calls clear/restore
       def get_status(self): ...
       def get_recent_logs(self): ...
   ```

6. **Create `utils/cloudflare/__init__.py`:**
   ```python
   from .tunnel import CloudflareTunnelManager
   __all__ = ['CloudflareTunnelManager']
   ```

7. **Delete `utils/cloudflare.py`** (the original monolithic file) after confirming all
   import sites resolve to `utils.cloudflare` (the package), which works identically due
   to `__init__.py`.

**Verify:** `from utils.cloudflare import CloudflareTunnelManager` resolves. All tunnel
routes still function. `TUNNEL_START_TIMEOUT` constant is accessible (either via
`utils.constants` after Task 5.2, or defined in `tunnel.py`).

---

## Phase 3 — Reliability / Leak Fixes

### Task 3.1 — Fix UI modal lifecycle leak

**Problem:** `web/ui.js` adds a `keydown` listener when the log modal opens but never
removes it on close or backdrop-click, leaking one listener per open.

**Read first:** `web/ui.js` (full — focus on log modal creation, open/close methods,
and the `keydown` listener).

**Changes:**

1. **Create `web/ui/logModal.js`** (file already has a `web/ui/` directory).
   Extract the log modal HTML construction, open/close logic, and auto-refresh timer into
   a module with explicit lifecycle:

   ```javascript
   export function createLogModal() {
       let _keydownHandler = null;
       let _refreshTimer = null;
       let _modalEl = null;

       function mount(container, { onClose, fetchLog }) {
           // build modal DOM
           _modalEl = buildModalDOM();
           container.appendChild(_modalEl);

           // store listener reference so we can remove it
           _keydownHandler = (e) => { if (e.key === 'Escape') unmount(); };
           document.addEventListener('keydown', _keydownHandler);

           // backdrop click
           _modalEl.querySelector('.modal-backdrop')
               .addEventListener('click', unmount);

           // auto-refresh if checkbox checked
           _refreshTimer = setInterval(() => { if (isAutoRefresh()) fetchLog(); }, 2000);
       }

       function unmount() {
           if (_keydownHandler) {
               document.removeEventListener('keydown', _keydownHandler);
               _keydownHandler = null;
           }
           if (_refreshTimer) {
               clearInterval(_refreshTimer);
               _refreshTimer = null;
           }
           _modalEl?.remove();
           _modalEl = null;
       }

       return { mount, unmount };
   }
   ```

2. **In `web/ui.js`**, replace the inline modal construction with:
   ```javascript
   import { createLogModal } from './ui/logModal.js';
   // ...
   showLogModal(workerId) {
       const modal = createLogModal();
       modal.mount(document.body, {
           onClose: () => modal.unmount(),
           fetchLog: () => this._fetchAndDisplayLog(workerId),
       });
   }
   ```

3. **Move any inline `style="..."` attributes from the modal DOM** into `distributed.css`
   class rules (e.g., `.log-modal`, `.log-modal__body`, `.log-modal__backdrop`).

**Verify:** Open log modal → close with Escape → open again. Confirm only one `keydown`
listener is registered at a time (check with DevTools → Event Listeners on `document`).

---

## Phase 4 — Deduplication

### Task 4.1 — Extract shared trace logging helpers

**Problem:** Three copies of identical `_trace_prefix` / `_trace_debug` / `_trace_info`
helpers exist in `api/orchestration/dispatch.py:11-20`,
`api/orchestration/media_sync.py:24-33`, `api/queue_orchestration.py:32-41`.

**Read first:** All three files (confirm the copies are byte-for-byte identical).

**Changes:**

1. **Create `utils/trace_logger.py`:**
   ```python
   from .logging import debug_log, log

   def trace_prefix(trace_execution_id: str) -> str:
       return f"[Distributed][exec:{trace_execution_id}]"

   def trace_debug(trace_execution_id: str, message: str) -> None:
       debug_log(f"{trace_prefix(trace_execution_id)} {message}")

   def trace_info(trace_execution_id: str, message: str) -> None:
       log(f"{trace_prefix(trace_execution_id)} {message}")
   ```

2. **In each of the three files**, delete the local definitions and add:
   ```python
   from ..utils.trace_logger import trace_prefix, trace_debug, trace_info
   ```
   Then rename call sites: `_trace_prefix` → `trace_prefix`, etc.

**Verify:** Grep for `def _trace_prefix` across the repo — should return zero results.

---

### Task 4.2 — Consolidate positive int/float parsing

**Problem:** Three near-identical implementations:
- `_parse_positive_int(value, default)` in `api/queue_orchestration.py:93-108`
- `_coerce_positive_int(value, default)` in `api/orchestration/dispatch.py:48-53`
- `validate_positive_int(value, field_name)` in `api/schemas.py:28-36`
Plus `_parse_positive_float` in `queue_orchestration.py`.

**Read first:** `api/schemas.py`, `api/queue_orchestration.py`, `api/orchestration/dispatch.py`.

**Changes:**

1. **Add to `api/schemas.py`** (after existing `validate_positive_int`):
   ```python
   def parse_positive_int(value, default: int) -> int:
       """Parse value as positive int, returning default on failure."""
       try:
           parsed = int(value)
       except (TypeError, ValueError):
           return max(1, int(default))
       return max(1, parsed)

   def parse_positive_float(value, default: float) -> float:
       """Parse value as positive float, returning default on failure."""
       try:
           parsed = float(value)
       except (TypeError, ValueError):
           return max(0.0, float(default))
       return max(0.0, parsed)
   ```

2. **In `api/queue_orchestration.py`:**
   - Delete `_parse_positive_int` and `_parse_positive_float`.
   - Add `from .schemas import parse_positive_int, parse_positive_float`.
   - Replace all call sites.

3. **In `api/orchestration/dispatch.py`:**
   - Delete `_coerce_positive_int`.
   - Add `from ..api.schemas import parse_positive_int`.
   - Replace call sites.

**Verify:** `grep -r "_parse_positive\|_coerce_positive"` returns zero results outside of
any test files.

---

### Task 4.3 — Deduplicate conditioning slicing

**Problem:** `upscale/modes/dynamic.py` contains conditioning slicing logic already
present in `upscale/tile_ops.py`.

**Read first:** `upscale/modes/dynamic.py` (full), `upscale/tile_ops.py` (full).
Identify the exact duplicated slicing pattern in each file (look for code that slices
`positive` / `negative` conditioning tensors by batch index).

**Changes:**

1. **Identify** the shared logic. It will look like:
   ```python
   # Slicing positive/negative conditioning for a specific batch index
   pos_sliced = [[t[:, idx:idx+1, ...] for t in cond_pair] for cond_pair in positive]
   neg_sliced = [[t[:, idx:idx+1, ...] for t in cond_pair] for cond_pair in negative]
   ```

2. **If the helper already exists in `tile_ops.py`**, find it and confirm its signature.
   If not, extract the logic from `dynamic.py` into a new function in `tile_ops.py`:
   ```python
   def slice_conditioning(positive, negative, batch_idx: int):
       """Return conditioning sliced to a single batch index."""
       ...
   ```

3. **In `upscale/modes/dynamic.py`**, delete the local implementation and replace with:
   ```python
   from ..tile_ops import slice_conditioning
   ```
   Update all call sites to use `slice_conditioning(positive, negative, idx)`.

**Verify:** The slicing logic exists in exactly one place. `grep -n "idx:idx+1\|batch_idx:batch_idx"`
in `dynamic.py` returns zero results.

---

### Task 4.4 — Unify worker probe/status logic in JavaScript

**Problem:** `web/workerLifecycle.js` and `web/executionUtils.js` both implement probe +
parse flows. `executionUtils.js` already uses `extension.api.probeWorker()` but
`workerLifecycle.js` still has raw `fetch` calls for `checkMasterStatus` and
`checkWorkerStatus`.

**Read first:** `web/workerLifecycle.js` (full), `web/executionUtils.js` (full),
`web/apiClient.js` (confirm `probeWorker` signature).

**Changes:**

1. **In `web/workerLifecycle.js`**, find `checkMasterStatus` and `checkWorkerStatus`.
   Note: These two functions have legitimate reasons for raw `fetch`:
   - `checkMasterStatus` uses a same-origin call (no CORS) to the local server
   - `checkWorkerStatus` may use `AbortSignal.any()` to combine timeout + cancellation

   **Assessment before changing:** If `probeWorker` in `apiClient.js` supports an
   `AbortSignal` parameter, migrate both to use `extension.api.probeWorker()`.
   If it does not, add an optional `signal` parameter to `probeWorker`:
   ```javascript
   probeWorker: async (workerUrl, timeout, signal = null) => {
       const controller = new AbortController();
       const timeoutId = setTimeout(() => controller.abort(), timeout);
       const effectiveSignal = signal
           ? AbortSignal.any([controller.signal, signal])
           : controller.signal;
       try {
           const response = await fetch(`${workerUrl}/prompt`, { signal: effectiveSignal });
           ...
       } finally {
           clearTimeout(timeoutId);
       }
   }
   ```

2. **Create a shared UI mapper** in `web/workerUtils.js` (already exists):
   ```javascript
   export function applyProbeResultToWorkerDot(workerId, probeResult) {
       const dot = document.getElementById(`status-${workerId}`);
       if (!dot) return;
       dot.classList.remove('worker-status--online', 'worker-status--offline',
                           'worker-status--processing', 'worker-status--unknown',
                           'status-pulsing');
       if (!probeResult || !probeResult.ok) {
           dot.classList.add('worker-status--offline');
           dot.title = 'Offline';
       } else if (probeResult.queueRemaining > 0) {
           dot.classList.add('worker-status--processing');
           dot.title = `Processing (${probeResult.queueRemaining} queued)`;
       } else {
           dot.classList.add('worker-status--online');
           dot.title = 'Online';
       }
   }
   ```

3. **Use `applyProbeResultToWorkerDot`** in both `workerLifecycle.js` and
   `executionUtils.js` (the `performPreflightCheck` inactive-worker update block at
   lines 915-922 of `executionUtils.js`).

**Verify:** `grep -n "style\." web/workerLifecycle.js` returns zero results. DOM status
updates flow exclusively through class manipulation helpers.

---

### Task 4.5 — Extract timeout handler duplication in result_collector

**Problem:** `upscale/result_collector.py` lines 116-161 contain nearly identical
worker-status logging blocks for static (`TileJobState`) vs dynamic (`ImageJobState`) modes
inside the `asyncio.TimeoutError` handler.

**Read first:** `upscale/result_collector.py` (full).

**Changes:**

1. **Extract a helper method** on `ResultCollectorMixin`:
   ```python
   def _log_worker_timeout_status(self, job_data, current_time: float, multi_job_id: str) -> None:
       """Log which workers have timed out and for how long."""
       if not isinstance(job_data, BaseJobState):
           return
       worker_status = dict(job_data.worker_status)
       for worker_id, last_seen in worker_status.items():
           elapsed = max(0.0, current_time - float(last_seen))
           log(
               f"[USDU] Timeout: job={multi_job_id} worker={worker_id[:8]} "
               f"elapsed={elapsed:.1f}s"
           )
   ```

2. **In the `asyncio.TimeoutError` block**, replace both the static-mode and dynamic-mode
   worker-status logging sections with a single call:
   ```python
   except asyncio.TimeoutError:
       current_time = time.time()
       self._log_worker_timeout_status(job_data_snapshot, current_time, multi_job_id)
       if mode == 'dynamic':
           if current_time - last_heartbeat_check >= HEARTBEAT_INTERVAL:
               requeued = await self._check_and_requeue_timed_out_workers(...)
               ...
   ```
   The `job_data_snapshot` should be obtained under the lock before the `wait_for` call
   (snapshot pattern from Task 1.1).

**Verify:** The timeout handler block has one code path for status logging, not two.

---

### Task 4.6 — Eliminate config update endpoint boilerplate

**Problem:** `api/config_routes.py` `update_worker_endpoint`, `update_master_endpoint`,
and `update_setting_endpoint` all repeat the same `if "field" in data: target[field] = ...`
pattern with special handling for `None` (pop vs assign) and `host` normalization.

**Read first:** `api/config_routes.py` (full, especially lines 107-275).

**Changes:**

1. **Define a field rule schema** near the top of `config_routes.py`:
   ```python
   # Each entry: (key, normalizer_or_None, remove_on_none: bool)
   _WORKER_FIELDS = [
       ('enabled',     None,           False),
       ('name',        None,           False),
       ('host',        normalize_host, True),   # pop key when value is None
       ('cuda_device', None,           True),
       ('extra_args',  None,           True),
       ('type',        None,           False),
   ]
   _MASTER_FIELDS = [
       ('name',        None,           False),
       ('host',        normalize_host, True),
   ]
   ```

2. **Add a helper function:**
   ```python
   def _apply_field_patch(target: dict, data: dict, field_rules: list) -> None:
       """Apply a partial update to target dict based on field rules."""
       for key, normalizer, remove_on_none in field_rules:
           if key not in data:
               continue
           value = data[key]
           if value is None and remove_on_none:
               target.pop(key, None)
           else:
               target[key] = normalizer(value) if (normalizer and value is not None) else value
   ```

3. **Replace the explicit if-blocks** in `update_worker_endpoint` with:
   ```python
   _apply_field_patch(worker, data, _WORKER_FIELDS)
   ```
   And in `update_master_endpoint`:
   ```python
   _apply_field_patch(config['master'], data, _MASTER_FIELDS)
   ```

**Verify:** `update_worker_endpoint` body no longer contains any `if "enabled" in data:` style
blocks. Behavior is identical — test with a PATCH that sets `host` to `None` (should remove
the key) and one that sets it to a URL string (should normalize it).

---

## Phase 5 — Medium-Low

### Task 5.1 — Centralize distributed state initialization

**Problem:** The `ensure_distributed_state()` / job queue init pattern is repeated in
`api/queue_orchestration.py`, `distributed.py`, `api/job_routes.py` (prepare_job endpoint),
and `nodes/collector.py`. Each touches `prompt_server.distributed_pending_jobs` and the
associated lock independently.

**Read first:** `api/queue_orchestration.py` (`ensure_distributed_state` function),
`distributed.py` (find where state attributes are set on `prompt_server`),
`api/job_routes.py` (`prepare_job_endpoint`, lines 131-146),
`nodes/collector.py` (lines 198-204).

**Changes:**

1. **Ensure `ensure_distributed_state()` in `api/queue_orchestration.py` is the sole place**
   that creates `prompt_server.distributed_pending_jobs`, `distributed_jobs_lock`, and
   any other top-level distributed state attributes. It should be idempotent:
   ```python
   def ensure_distributed_state(server_instance=None):
       ps = server_instance or prompt_server
       if not hasattr(ps, 'distributed_pending_jobs'):
           ps.distributed_pending_jobs = {}
       if not hasattr(ps, 'distributed_jobs_lock'):
           ps.distributed_jobs_lock = asyncio.Lock()
   ```

2. **In `api/job_routes.py` `prepare_job_endpoint`**, replace the inline init:
   ```python
   # Before:
   async with prompt_server.distributed_jobs_lock:
       if multi_job_id not in prompt_server.distributed_pending_jobs:
           prompt_server.distributed_pending_jobs[multi_job_id] = asyncio.Queue()
   ```
   with a call to `ensure_distributed_state()` at startup (it's idempotent) and keep the
   per-job queue init in place (that's correct per-job logic, not state init).

3. **In `nodes/collector.py`** (lines 198-204), the early queue creation is correct
   behavior (race prevention) — leave it. But ensure `ensure_distributed_state()` has been
   called before the node can ever run (it's called at startup in `queue_orchestration.py`
   module import time).

4. **In `distributed.py`**, remove any direct `prompt_server.distributed_* = ...`
   assignments and replace with a call to `ensure_distributed_state()` at startup.

**Verify:** Grep for `distributed_pending_jobs = {}` and `distributed_jobs_lock = ` — both
should appear only in `ensure_distributed_state`. All other sites just call that function.

---

### Task 5.2 — Move all magic numbers to `utils/constants.py`

**Problem:** Several timeout and size values are hard-coded in files rather than defined in
`utils/constants.py`.

**Read first:** `utils/constants.py` (full), then check each file listed below.

**Changes:**

1. **In `utils/constants.py`**, add:
   ```python
   # Cloudflare tunnel
   TUNNEL_START_TIMEOUT: float = float(os.environ.get("TUNNEL_START_TIMEOUT", "25"))
   CLOUDFLARE_LOG_BUFFER_SIZE: int = 200

   # USDU result collection
   DYNAMIC_MODE_MAX_POLL_TIMEOUT: float = 10.0

   # Static mode job poll loop
   JOB_POLL_INTERVAL: float = 1.0
   JOB_POLL_MAX_ATTEMPTS: int = 20
   ```

2. **In `utils/cloudflare.py`** (or `utils/cloudflare/tunnel.py` after Task 2.2):
   - Replace `TUNNEL_START_TIMEOUT = 25` (line ~30) with import from constants.
   - Replace `> 200` log buffer check (line ~178) with `CLOUDFLARE_LOG_BUFFER_SIZE`.

3. **In `upscale/result_collector.py`** (line ~65):
   - Replace `min(10.0, timeout)` with `min(DYNAMIC_MODE_MAX_POLL_TIMEOUT, timeout)`.
   - Add import: `from ..utils.constants import DYNAMIC_MODE_MAX_POLL_TIMEOUT`.

4. **In `upscale/modes/static.py`**:
   - Replace `time.sleep(1.0)` with `time.sleep(JOB_POLL_INTERVAL)`.
   - Replace `max_attempts=20` with `max_attempts=JOB_POLL_MAX_ATTEMPTS`.
   - Add import: `from ...utils.constants import JOB_POLL_INTERVAL, JOB_POLL_MAX_ATTEMPTS`.

**Verify:** `grep -rn "= 25\b\|> 200\b\|10\.0\b\|max_attempts=20\b" utils/ upscale/modes/`
shows no magic number remnants in these files.

---

## Phase 6 — Low / Cleanup

### Task 6.1 — Remove dead code and import drift

**Problem:** `api/worker_routes.py` imports helpers from `workers/detection.py` that
are not used within that file. `workers/detection.py` `get_comms_channel()` may be dormant.

**Read first:** `api/worker_routes.py` (all imports at top of file), `workers/detection.py`
(full). Grep for every symbol imported in `worker_routes.py` to confirm usage.

**Changes:**

1. **In `api/worker_routes.py`**, identify every imported name from `workers.detection`
   (and any other module). For each, search the file body for usage. Remove any import
   line whose symbols are not referenced.

2. **For `workers/detection.py` `get_comms_channel()`**: grep the full codebase for
   `get_comms_channel`. If found only in `nodes/collector.py` imports but never called
   (check `collector.py` body), remove the import in `collector.py` and the function
   in `detection.py`. If it is called, leave it.

3. **Also in `nodes/collector.py`** (line 20): `from ..workers.detection import is_local_worker, get_comms_channel`.
   Check whether `is_local_worker` and `get_comms_channel` are referenced in the file body.
   Remove unused ones.

**Verify:** `python -m py_compile` on each modified file. No `ImportError` on module load.

---

### Task 6.2 — Standardize error response format across API routes

**Problem:** Some endpoints use `await handle_api_error(request, msg, status)`,
others call `web.json_response({"error": ...}, status=400)` directly, creating
inconsistent client-facing error shapes.

**Read first:** `utils/network.py` (`handle_api_error` implementation),
`api/job_routes.py` (lines 276-288 — list-of-errors response),
`api/config_routes.py` (find all direct `web.json_response` error calls).

**Changes:**

1. **Extend `handle_api_error`** (if needed) to accept a list of error strings:
   ```python
   async def handle_api_error(request, error, status: int = 500):
       if isinstance(error, list):
           return web.json_response({"errors": error}, status=status)
       msg = str(error)
       debug_log(f"API error [{status}]: {msg}")
       return web.json_response({"error": msg}, status=status)
   ```

2. **In `api/job_routes.py`** (lines ~287-288):
   ```python
   # Before:
   return web.json_response({"error": errors}, status=400)
   # After:
   return await handle_api_error(request, errors, 400)
   ```

3. **In `api/config_routes.py`**, replace all `return web.json_response({"error": ...}, status=...)`
   calls with `return await handle_api_error(request, ..., status)`.

4. **In `api/worker_routes.py`**, check for any direct `web.json_response` error patterns
   and standardize them.

**Verify:** `grep -rn 'json_response.*"error"' api/` returns zero results (all error
responses now flow through `handle_api_error`).

---

### Task 6.3 — Fix redundant server import alias

**Problem:** `api/job_routes.py` lines 10-11 have both:
```python
import server
import server as _server
```
Only one is needed.

**Read first:** `api/job_routes.py` (check which alias is used throughout the file).

**Changes:**

1. Find all usages of `server` and `_server` in `api/job_routes.py`.
2. Pick one (the dominant one, or `server` for consistency with other files).
3. Remove the unused import and update any references to the removed alias.

---

### Task 6.4 — Document mixin inheritance chain

**Problem:** `ResultCollectorMixin`, `StaticModeMixin`, `DynamicModeMixin` in
`upscale/result_collector.py` and `upscale/modes/` have no documentation of the required
method interface, making them hard to extend.

**Read first:** `upscale/result_collector.py`, `upscale/modes/static.py`,
`upscale/modes/dynamic.py`, `upscale/worker_comms.py`.

**Changes:**

1. **In `upscale/result_collector.py`**, add a docstring to `ResultCollectorMixin`
   listing abstract methods / expected attributes:
   ```python
   class ResultCollectorMixin:
       """
       Mixin for master-side result collection in USDU distributed jobs.

       Requires the consuming class to also mix in:
       - JobStateMixin (upscale.job_state) for job state access
       - WorkerCommsMixin (upscale.worker_comms) is NOT required here

       Expected to be used via multiple inheritance in DistributedUSDU node class.
       """
   ```

2. **In `upscale/modes/static.py`** and `upscale/modes/dynamic.py`**, add docstrings
   to `StaticModeMixin` and `DynamicModeMixin` listing what they call on `self` and what
   they require from co-mixins.

3. **No behavior changes** — documentation only.

---

## Verification Checklist (run after all phases)

```bash
# Python syntax check all modified files
find custom_nodes/ComfyUI-Distributed -name "*.py" | xargs python -m py_compile

# Confirm no bare usdu_managment god-module imports remain
grep -r "from.*usdu_managment import" custom_nodes/ComfyUI-Distributed \
  --include="*.py" | grep -v "usdu_managment.py"

# Confirm trace helpers are in one place
grep -rn "def _trace_prefix\|def _trace_debug\|def _trace_info" \
  custom_nodes/ComfyUI-Distributed --include="*.py"
# Expected: 0 results (all moved to utils/trace_logger.py as trace_prefix etc.)

# Confirm no lock-held network calls remain
grep -A5 "distributed_tile_jobs_lock" custom_nodes/ComfyUI-Distributed/utils/usdu_managment.py \
  | grep "await probe_worker"
# Expected: 0 results

# Confirm mutating JS calls don't retry
grep -A3 "queueDistributed\|launchWorker\|stopWorker\|startTunnel\|stopTunnel" \
  custom_nodes/ComfyUI-Distributed/web/apiClient.js | grep "retry: false"
# Expected: one hit per mutating call

# Confirm no raw web.json_response error patterns remain
grep -rn 'web\.json_response.*"error"' custom_nodes/ComfyUI-Distributed/api/ --include="*.py"
# Expected: 0 results
```

---

## Notes for the implementing agent

- **Do not change any HTTP endpoint paths or request/response JSON schemas.** All changes
  are internal to Python modules and JS modules — the public API surface is frozen.
- **Do not change `upscale/job_models.py`** — the typed dataclasses are already correct.
- **Phase 1 tasks are prerequisites for Phase 2.** Specifically, Task 1.1 must be done
  before Task 2.1 since the requeued function moves in 2.1.
- **Task 2.1 is the most complex.** Read all import sites before moving any code. Use
  re-exports in the shim `usdu_managment.py` generously to avoid breaking callers.
- **Tasks within the same phase are independent** and can be done in any order or in
  parallel.
- **After each task**, run `python -m py_compile` on all touched files before moving on.