mirror of
https://github.com/Comfy-Org/ComfyUI.git
synced 2026-08-25 18:32:35 +08:00
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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user