diff --git a/README.md b/README.md index 410d6ae..d3f8b87 100644 --- a/README.md +++ b/README.md @@ -43,17 +43,19 @@ bridge failure never blocks ComfyUI from loading the nodes. **Agent Emit (-> agent)** — send from the graph to the agent. - Inputs: `channel` (STRING, default `main`); optional `text` (multiline STRING), - `image` (IMAGE). -- Writes the latest text/image to the channel inbox (image is saved to a PNG in - the temp dir and the path is stored). Passes `(text, image)` straight through - as outputs, so you can wire it inline. + `image` (IMAGE), `seed` (INT, socket input — e.g. wire a KSampler seed). +- Writes the latest text/image/seed to the channel inbox (image is saved to a PNG + in the temp dir and the path is stored). Passes `(text, image, seed)` straight + through as outputs, so you can wire it inline. **Agent Receive (<- agent)** — receive from the agent into the graph. - Inputs: `channel` (STRING, default `main`); `wait_seconds` (FLOAT, default `30.0`, max `86400`); `keep_last` (BOOLEAN, default `true`); `stop_on_timeout` (BOOLEAN, default `true`). -- Blocks up to `wait_seconds` for a new `comfy_push` on the channel, then - outputs `(text, image)`. Each pushed value is **consumed once** (turn-based), +- Outputs `(text, image, seed)` — `seed` (INT) is whatever the agent passed to + `comfy_push(seed=...)`, or `0` if none. Wire it into a KSampler, etc. +- Blocks up to `wait_seconds` for a new `comfy_push` on the channel. Each pushed + value is **consumed once** (turn-based), so a second receive with no new push won't replay it. - **`keep_last`** — on timeout (no new message), output the *last* message again instead of blanking. Great for Auto Queue so the display holds steady between @@ -180,8 +182,8 @@ this — it's inline. | Tool | Direction | Purpose | | --- | --- | --- | -| `comfy_pull(channel="main")` | graph -> agent | Read the latest text/image emitted on a channel (returns `{turn, text, image_path}`). | -| `comfy_push(channel="main", text=None, image_path=None)` | agent -> graph | Send text/image to the channel's `Agent Receive` node. | +| `comfy_pull(channel="main")` | graph -> agent | Read the latest text/image/seed emitted on a channel (returns `{turn, text, image_path, seed}`). | +| `comfy_push(channel="main", text=None, image_path=None, seed=None)` | agent -> graph | Send text/image/seed to the channel's `Agent Receive` node. | | `comfy_list_channels()` | — | List active channels and their turn counters. | | `comfy_run_workflow(name, inputs=None, wait=True)` | — | Load a saved workflow from `workflows/`, optionally inject inputs, queue it via the ComfyUI API, optionally poll for results. | | `comfy_get_result(prompt_id)` | — | Fetch produced images/outputs for a queued prompt id. | diff --git a/src/bridge/mcp_tools.py b/src/bridge/mcp_tools.py index 0f1662d..02537ca 100644 --- a/src/bridge/mcp_tools.py +++ b/src/bridge/mcp_tools.py @@ -15,11 +15,12 @@ def comfy_pull(channel: str = "main") -> dict: def comfy_push(channel: str = "main", text: str | None = None, - image_path: str | None = None) -> dict: - """Send text/image to the graph's Agent Receive node on a channel.""" + image_path: str | None = None, seed: int | None = None) -> dict: + """Send text/image/seed to the graph's Agent Receive node on a channel.""" # translate the agent-visible path back to the bridge's container-side path image_path = paths.to_local(image_path) - turn = ChannelStore.instance().push(channel, text=text, image_path=image_path) + turn = ChannelStore.instance().push(channel, text=text, + image_path=image_path, seed=seed) return {"channel": channel, "out_turn": turn} diff --git a/src/bridge/store.py b/src/bridge/store.py index 71feffd..e7ad85f 100644 --- a/src/bridge/store.py +++ b/src/bridge/store.py @@ -16,6 +16,7 @@ _GLOBAL_ATTR = "_comfyui_agent_bridge_channel_store" class _Slot: text: str | None = None image_path: str | None = None + seed: int | None = None turn: int = 0 @@ -52,14 +53,20 @@ class ChannelStore: self._channels[name] = ch return ch + @staticmethod + def _payload(s: "_Slot") -> dict: + return {"turn": s.turn, "text": s.text, "image_path": s.image_path, + "seed": s.seed} + # --- inbox: graph -> agent --- def emit(self, channel: str, text: str | None = None, - image_path: str | None = None) -> int: + image_path: str | None = None, seed: int | None = None) -> int: with self._cond: ch = self._chan(channel) ch.inbox.turn += 1 ch.inbox.text = text ch.inbox.image_path = image_path + ch.inbox.seed = seed self._cond.notify_all() return ch.inbox.turn @@ -67,33 +74,31 @@ class ChannelStore: with self._lock: ch = self._channels.get(channel) if ch is None: - return {"turn": 0, "text": None, "image_path": None} - s = ch.inbox - return {"turn": s.turn, "text": s.text, "image_path": s.image_path} + return {"turn": 0, "text": None, "image_path": None, "seed": None} + return self._payload(ch.inbox) # --- outbox: agent -> graph --- def push(self, channel: str, text: str | None = None, - image_path: str | None = None) -> int: + image_path: str | None = None, seed: int | None = None) -> int: with self._cond: ch = self._chan(channel) ch.outbox.turn += 1 ch.outbox.text = text ch.outbox.image_path = image_path + ch.outbox.seed = seed self._cond.notify_all() return ch.outbox.turn def receive(self, channel: str, wait_seconds: float = 0.0, should_abort=None, poll_interval: float = 0.1) -> dict: - empty = {"turn": 0, "text": None, "image_path": None} + empty = {"turn": 0, "text": None, "image_path": None, "seed": None} deadline = time.monotonic() + max(0.0, wait_seconds) with self._cond: while True: ch = self._channels.get(channel) if ch is not None and ch.outbox.turn > ch.last_consumed_out_turn: ch.last_consumed_out_turn = ch.outbox.turn - s = ch.outbox - return {"turn": s.turn, "text": s.text, - "image_path": s.image_path} + return self._payload(ch.outbox) remaining = deadline - time.monotonic() if remaining <= 0: return empty @@ -111,9 +116,8 @@ class ChannelStore: with self._lock: ch = self._channels.get(channel) if ch is None: - return {"turn": 0, "text": None, "image_path": None} - s = ch.outbox - return {"turn": s.turn, "text": s.text, "image_path": s.image_path} + return {"turn": 0, "text": None, "image_path": None, "seed": None} + return self._payload(ch.outbox) def list_channels(self) -> list[dict]: with self._lock: diff --git a/src/nodes/emit.py b/src/nodes/emit.py index 88ddd53..2482052 100644 --- a/src/nodes/emit.py +++ b/src/nodes/emit.py @@ -3,12 +3,12 @@ from ..bridge import images, paths class AgentEmit: - """Send text/image from the graph to an external agent on a named channel.""" + """Send text/image/seed from the graph to an external agent on a channel.""" CATEGORY = "agents/bridge" FUNCTION = "run" - RETURN_TYPES = ("STRING", "IMAGE") - RETURN_NAMES = ("text", "image") + RETURN_TYPES = ("STRING", "IMAGE", "INT") + RETURN_NAMES = ("text", "image", "seed") @classmethod def INPUT_TYPES(cls): @@ -19,13 +19,16 @@ class AgentEmit: "optional": { "text": ("STRING", {"default": "", "multiline": True}), "image": ("IMAGE",), + "seed": ("INT", {"default": 0, "min": 0, + "max": 0xFFFFFFFFFFFFFFFF, "forceInput": True}), }, } - def run(self, channel="main", text=None, image=None): + def run(self, channel="main", text=None, image=None, seed=None): image_path = None if image is not None: image_path = images.save_tensor_png(image, paths.tmp_dir()) text_val = text if text else None - ChannelStore.instance().emit(channel, text=text_val, image_path=image_path) - return (text_val, image) + ChannelStore.instance().emit(channel, text=text_val, + image_path=image_path, seed=seed) + return (text_val, image, seed if seed is not None else 0) diff --git a/src/nodes/receive.py b/src/nodes/receive.py index ce7387a..0378ddb 100644 --- a/src/nodes/receive.py +++ b/src/nodes/receive.py @@ -7,8 +7,8 @@ class AgentReceive: CATEGORY = "agents/bridge" FUNCTION = "run" - RETURN_TYPES = ("STRING", "IMAGE") - RETURN_NAMES = ("text", "image") + RETURN_TYPES = ("STRING", "IMAGE", "INT") + RETURN_NAMES = ("text", "image", "seed") @classmethod def INPUT_TYPES(cls): @@ -36,7 +36,8 @@ class AgentReceive: image = images.load_png_tensor(payload["image_path"]) else: image = images.empty_image() - return (text, image) + seed = payload.get("seed") + return (text, image, seed if seed is not None else 0) @staticmethod def _interrupt_check(): @@ -75,4 +76,4 @@ class AgentReceive: last = store.peek(channel) if last["turn"] > 0: return self._format(last) - return ("", images.empty_image()) + return ("", images.empty_image(), 0) diff --git a/tests/bridge/test_store_inbox.py b/tests/bridge/test_store_inbox.py index eb1ca51..01a3f51 100644 --- a/tests/bridge/test_store_inbox.py +++ b/tests/bridge/test_store_inbox.py @@ -2,16 +2,17 @@ from src.bridge.store import ChannelStore def test_emit_then_pull_returns_payload(): s = ChannelStore.instance() - s.emit("main", text="hello", image_path="/tmp/a.png") + s.emit("main", text="hello", image_path="/tmp/a.png", seed=7) got = s.pull("main") assert got["text"] == "hello" assert got["image_path"] == "/tmp/a.png" + assert got["seed"] == 7 assert got["turn"] == 1 def test_pull_unknown_channel_is_empty_not_error(): s = ChannelStore.instance() got = s.pull("never") - assert got == {"turn": 0, "text": None, "image_path": None} + assert got == {"turn": 0, "text": None, "image_path": None, "seed": None} def test_emit_bumps_turn_each_time(): s = ChannelStore.instance() diff --git a/tests/bridge/test_store_outbox.py b/tests/bridge/test_store_outbox.py index 2048267..9d9b503 100644 --- a/tests/bridge/test_store_outbox.py +++ b/tests/bridge/test_store_outbox.py @@ -48,7 +48,7 @@ def test_peek_does_not_consume(): def test_receive_nonblocking_empty_channel(): s = ChannelStore.instance() got = s.receive("nope", wait_seconds=0) - assert got == {"turn": 0, "text": None, "image_path": None} + assert got == {"turn": 0, "text": None, "image_path": None, "seed": None} def test_receive_blocks_until_push(): s = ChannelStore.instance() diff --git a/tests/nodes/test_emit.py b/tests/nodes/test_emit.py index 9e142c6..8735394 100644 --- a/tests/nodes/test_emit.py +++ b/tests/nodes/test_emit.py @@ -21,4 +21,10 @@ def test_emit_image_materializes_path(tmp_path, monkeypatch): def test_emit_returns_passthrough(): node = AgentEmit() out = node.run(channel="main", text="hi", image=None) - assert out == ("hi", None) # (text, image) passthrough + assert out == ("hi", None, 0) # (text, image, seed) passthrough + +def test_emit_sends_seed(): + node = AgentEmit() + node.run(channel="main", text="hi", image=None, seed=999) + got = ChannelStore.instance().pull("main") + assert got["seed"] == 999 diff --git a/tests/nodes/test_receive.py b/tests/nodes/test_receive.py index 39d6e83..d9dfdac 100644 --- a/tests/nodes/test_receive.py +++ b/tests/nodes/test_receive.py @@ -17,49 +17,57 @@ def test_is_changed_is_always_dirty(): def test_receive_text(monkeypatch): ChannelStore.instance().push("main", text="done") node = AgentReceive() - text, image = node.run(channel="main", wait_seconds=0.0) + text, image, seed = node.run(channel="main", wait_seconds=0.0) assert text == "done" # no image pushed -> placeholder empty image assert image.shape[0] == 1 and image.shape[3] == 3 + assert seed == 0 # no seed pushed -> 0 def test_receive_image(tmp_path): p = images.save_tensor_png(torch.zeros(1, 4, 4, 3), str(tmp_path)) ChannelStore.instance().push("main", image_path=p) node = AgentReceive() - text, image = node.run(channel="main", wait_seconds=0.0) + text, image, seed = node.run(channel="main", wait_seconds=0.0) assert text == "" # no text -> empty string for STRING output assert image.shape == (1, 4, 4, 3) +def test_receive_seed(): + ChannelStore.instance().push("main", seed=123456789) + node = AgentReceive() + text, image, seed = node.run(channel="main", wait_seconds=0.0) + assert seed == 123456789 + def test_receive_timeout_returns_empty(): node = AgentReceive() - text, image = node.run(channel="main", wait_seconds=0.2, keep_last=False) + text, image, seed = node.run(channel="main", wait_seconds=0.2, keep_last=False) assert text == "" assert image.shape[0] == 1 + assert seed == 0 def test_keep_last_holds_value_after_consume(): - ChannelStore.instance().push("main", text="held-msg") + ChannelStore.instance().push("main", text="held-msg", seed=42) node = AgentReceive() - t1, _ = node.run(channel="main", wait_seconds=0.0) - assert t1 == "held-msg" # consumed the fresh message - # no new push -> keep_last holds the last value instead of blanking - t2, _ = node.run(channel="main", wait_seconds=0.0, keep_last=True, - stop_on_timeout=False) - assert t2 == "held-msg" + t1, _, s1 = node.run(channel="main", wait_seconds=0.0) + assert t1 == "held-msg" and s1 == 42 # consumed the fresh message + # no new push -> keep_last holds the last value (text AND seed) + t2, _, s2 = node.run(channel="main", wait_seconds=0.0, keep_last=True, + stop_on_timeout=False) + assert t2 == "held-msg" and s2 == 42 def test_keep_last_false_blanks_after_consume(): ChannelStore.instance().push("main", text="one") node = AgentReceive() node.run(channel="main", wait_seconds=0.0) # consume - t2, _ = node.run(channel="main", wait_seconds=0.0, keep_last=False, - stop_on_timeout=False) + t2, _, _ = node.run(channel="main", wait_seconds=0.0, keep_last=False, + stop_on_timeout=False) assert t2 == "" def test_stop_on_timeout_without_server_does_not_crash(): node = AgentReceive() - text, image = node.run(channel="main", wait_seconds=0.0, - stop_on_timeout=True) + text, image, seed = node.run(channel="main", wait_seconds=0.0, + stop_on_timeout=True) assert text == "" assert image.shape[0] == 1