diff --git a/services/queue_submit.py b/services/queue_submit.py index 28481d4..0c2c5c7 100644 --- a/services/queue_submit.py +++ b/services/queue_submit.py @@ -57,6 +57,24 @@ COMFYUI_URL = ( or os.environ.get("MOLTBOT_COMFYUI_URL") or "http://127.0.0.1:8188" ) +# IMPORTANT: keep this fixed and prompt-free; ComfyUI forwards it into API-node hidden inputs. +COMFY_USAGE_SOURCE = "comfyui-openclaw" + + +def _build_queue_extra_data( + extra_data: Optional[Dict[str, Any]] = None, + *, + tenant_id: Optional[str] = None, +) -> Dict[str, Any]: + payload_extra = dict(extra_data or {}) + payload_extra.setdefault("comfy_usage_source", COMFY_USAGE_SOURCE) + + openclaw_value = payload_extra.get("openclaw") + openclaw_extra = dict(openclaw_value) if isinstance(openclaw_value, dict) else {} + openclaw_extra.setdefault("tenant_id", tenant_id or get_current_tenant_id()) + payload_extra["openclaw"] = openclaw_extra + + return payload_extra async def submit_prompt( @@ -111,13 +129,7 @@ async def submit_prompt( client_id = str(uuid.uuid4()) payload = {"prompt": prompt_workflow, "client_id": client_id} - - if extra_data: - payload["extra_data"] = extra_data - - # S49: keep tenant context in queue metadata for cross-service traceability. - openclaw_extra = payload.setdefault("extra_data", {}).setdefault("openclaw", {}) - openclaw_extra.setdefault("tenant_id", tenant_id or get_current_tenant_id()) + payload["extra_data"] = _build_queue_extra_data(extra_data, tenant_id=tenant_id) # NOTE: Debug-only full payload logging for troubleshooting mismatched outputs. # Enable with OPENCLAW_DEBUG_PROMPT_PAYLOAD=1. This may include sensitive prompt content. diff --git a/tests/test_r122_real_backend_lane.py b/tests/test_r122_real_backend_lane.py index acf5e4c..c67e080 100644 --- a/tests/test_r122_real_backend_lane.py +++ b/tests/test_r122_real_backend_lane.py @@ -159,6 +159,14 @@ class TestR122RealBackendLane(AioHTTPTestCase): self.assertEqual( self._prompt_payload["prompt"]["3"]["inputs"]["text"], "real backend lane" ) + self.assertEqual( + self._prompt_payload["extra_data"]["comfy_usage_source"], + "comfyui-openclaw", + ) + self.assertNotIn( + "real backend lane", + self._prompt_payload["extra_data"]["comfy_usage_source"], + ) if __name__ == "__main__": diff --git a/tests/test_r62_queue_degrade.py b/tests/test_r62_queue_degrade.py index f3d0627..eba8ed1 100644 --- a/tests/test_r62_queue_degrade.py +++ b/tests/test_r62_queue_degrade.py @@ -14,7 +14,11 @@ from api.errors import APIError, ErrorCode # We import submit_prompt inside tests to ensure mocks applied before import if needed, # although the lazy import inside the function makes it easier. -from services.queue_submit import submit_prompt +from services.queue_submit import ( + COMFY_USAGE_SOURCE, + _build_queue_extra_data, + submit_prompt, +) class TestR62QueueDegrade(unittest.IsolatedAsyncioTestCase): @@ -53,6 +57,72 @@ class TestR62QueueDegrade(unittest.IsolatedAsyncioTestCase): self.assertEqual(result["prompt_id"], "123") mock_session_inst.post.assert_called_once() + sent_payload = mock_session_inst.post.call_args.kwargs["json"] + self.assertEqual( + sent_payload["extra_data"]["comfy_usage_source"], COMFY_USAGE_SOURCE + ) + self.assertIn("tenant_id", sent_payload["extra_data"]["openclaw"]) + + def test_queue_extra_data_sets_stable_usage_source_without_prompt_leak(self): + prompt_text = "private prompt body should not become attribution" + + extra = _build_queue_extra_data( + {"openclaw": {"trace_id": "trace-1"}, "moltbot": {"trace_id": "trace-1"}}, + tenant_id="tenant-a", + ) + + self.assertEqual(extra["comfy_usage_source"], COMFY_USAGE_SOURCE) + self.assertEqual(extra["openclaw"]["trace_id"], "trace-1") + self.assertEqual(extra["openclaw"]["tenant_id"], "tenant-a") + self.assertEqual(extra["moltbot"]["trace_id"], "trace-1") + self.assertNotIn(prompt_text, extra["comfy_usage_source"]) + self.assertNotIn("tenant-a", extra["comfy_usage_source"]) + self.assertNotIn("trace-1", extra["comfy_usage_source"]) + + async def test_submit_preserves_caller_usage_source_and_openclaw_metadata(self): + """Caller-provided ComfyUI attribution takes precedence.""" + + mock_response = MagicMock() + mock_response.status = 200 + mock_response.json = AsyncMock(return_value={"prompt_id": "123", "number": 1}) + + mock_session_inst = MagicMock() + mock_session_inst.post.return_value.__aenter__.return_value = mock_response + + mock_session_cls = MagicMock() + mock_session_cls.return_value.__aenter__.return_value = mock_session_inst + + mock_aiohttp = MagicMock() + mock_aiohttp.ClientSession = mock_session_cls + + caller_extra = { + "comfy_usage_source": "caller-owned-source", + "openclaw": {"trace_id": "trace-caller"}, + "moltbot": {"trace_id": "trace-caller"}, + } + workflow = { + "1": { + "class_type": "CLIPTextEncode", + "inputs": {"text": "private prompt text"}, + } + } + + with patch.dict(sys.modules, {"aiohttp": mock_aiohttp}): + result = await submit_prompt( + workflow, + extra_data=caller_extra, + tenant_id="tenant-caller", + ) + + self.assertEqual(result["prompt_id"], "123") + sent_payload = mock_session_inst.post.call_args.kwargs["json"] + sent_extra = sent_payload["extra_data"] + self.assertEqual(sent_extra["comfy_usage_source"], "caller-owned-source") + self.assertEqual(sent_extra["openclaw"]["trace_id"], "trace-caller") + self.assertEqual(sent_extra["openclaw"]["tenant_id"], "tenant-caller") + self.assertEqual(sent_extra["moltbot"]["trace_id"], "trace-caller") + self.assertNotIn("private prompt text", sent_extra["comfy_usage_source"]) + self.assertEqual(caller_extra["openclaw"], {"trace_id": "trace-caller"}) async def test_upstream_failure(self): """Test standard APIError when upstream returns non-200.""" diff --git a/tests/test_r68_integration_flow.py b/tests/test_r68_integration_flow.py index 487430d..79153fa 100644 --- a/tests/test_r68_integration_flow.py +++ b/tests/test_r68_integration_flow.py @@ -201,6 +201,18 @@ class TestR68IntegrationFlow(AioHTTPTestCase): # Correct inputs injected? workflow = sent_payload["prompt"] self.assertEqual(workflow["3"]["inputs"]["text"], "integration flow") + self.assertEqual( + sent_payload["extra_data"]["comfy_usage_source"], + "comfyui-openclaw", + ) + self.assertEqual( + sent_payload["extra_data"]["openclaw"]["tenant_id"], "default" + ) + self.assertIn("trace_id", sent_payload["extra_data"]["moltbot"]) + self.assertNotIn( + "integration flow", + sent_payload["extra_data"]["comfy_usage_source"], + ) # Verify Trace ID propagation self.assertIn("trace_id", data)