fix(server): bound metadata envelope and clean up on cancel paths

Addresses review feedback on the per-prompt metadata envelope:

- Sanitize at the boundary: reject envelopes larger than 16 keys, keys
  over 64 chars, values over 256 chars, or anything that isn't a flat
  ``dict[str, str]``. Logs a warning so abuse is observable. Stops a
  malicious client from inflating broadcast volume by stamping a 10 MB
  metadata blob onto every WS event.
- Cap the in-memory store at 4096 concurrent envelopes with FIFO
  eviction. Acts as a backstop if any cleanup hook is skipped.
- Drop envelopes when prompts are cancelled before reaching the worker:
  ``PromptQueue.wipe_queue`` and ``delete_queue_item`` now call
  ``server.unregister_prompt_metadata`` for every removed item.
- Drop envelopes on hard execution failures: the worker now wraps
  ``e.execute()`` in ``try/finally``, so an uncaught exception in
  execution no longer leaks the envelope.
- Guard the WS reconnect handler: only include ``prompt_id`` in the
  ``executing`` payload when ``last_prompt_id`` is set, so clients
  with strict schemas (zod ``prompt_id: zJobId``) don't reject the
  message with a null id.
- Extract a ``PromptMetadataStore`` class that owns the dict and the
  bounds, so ``PromptServer`` becomes a thin delegating layer and the
  full register/inject/unregister cycle (plus FIFO eviction and
  sanitization) is unit-tested without torch.

44 tests passing; ruff clean on all touched files.
This commit is contained in:
Deep Mehta
2026-05-14 21:03:38 -07:00
parent 74cfcaa318
commit fd89498eac
5 changed files with 362 additions and 83 deletions

View File

@@ -1296,7 +1296,10 @@ class PromptQueue:
def wipe_queue(self):
with self.mutex:
dropped_prompt_ids = [item[1] for item in self.queue]
self.queue = []
for prompt_id in dropped_prompt_ids:
self.server.unregister_prompt_metadata(prompt_id)
self.server.queue_updated()
def delete_queue_item(self, function):
@@ -1306,8 +1309,9 @@ class PromptQueue:
if len(self.queue) == 1:
self.wipe_queue()
else:
self.queue.pop(x)
deleted = self.queue.pop(x)
heapq.heapify(self.queue)
self.server.unregister_prompt_metadata(deleted[1])
self.server.queue_updated()
return True
return False