diff --git a/README.md b/README.md index 78b0cf5..5620ffb 100644 --- a/README.md +++ b/README.md @@ -189,6 +189,31 @@ variable. Extension scripts and helper scripts revalidate their browser cache on reload. Only the refresh extension has a versioned URL, preserving other plugins' module identities and relative imports. A reverse proxy must honor these cache headers. +When the connection drops, the page displays “服务器连接已断开，正在等待恢复” +once and waits without refreshing. Reconnecting clears the notice and checks the +frontend revision again. Load the updated script while the server is available +before testing a shutdown; an already open page keeps its previously loaded code. + +## Interrupted API task recovery + +With the distributed queue enabled, successful asynchronous submissions made through +the audited `requests`/`httpx` clients create encrypted recovery checkpoints. After a +job becomes `worker_lost`, a single recovery worker queries the original task ID, +saves returned videos under the owner's output directory, and adds them to that +owner's history. It never resubmits generation or runs downstream workflow nodes; +the history remains failed with an explicit recovery message. + +Supported protocols are Seedance Ark, Wan OPC/Vapeur, and Kling official/Vapeur video +endpoints using bearer authentication. Other protocols, synchronous calls, and old +submissions without checkpoints are not automatically recovered. Short-lived bearer +tokens must still be valid; expired credentials or missing tasks require manual lookup. + +Network/download failures retry with backoff, stopping after 20 consecutive failures +or seven days since submission. Keep the scheduler database, history database, +internal signing secret, and shared output directory across restarts. The secret is +needed to decrypt pending query credentials; terminal checkpoints discard credentials. +Multiple workers must share these paths and support the recovery file lock. + ## API Access Authenticated API calls can include either: diff --git a/account_manager.py b/account_manager.py index 292cad8..1b0d93d 100644 --- a/account_manager.py +++ b/account_manager.py @@ -6,11 +6,14 @@ import folder_paths from aiohttp import web from server import PromptServer, is_loopback +from comfy_execution.utils import get_executing_context from .utils import * from .utils.api_audit import install_api_audit from .utils.frontend import create_jobs_middleware, install_origin_middleware from .utils.frontend_reload import install_frontend_reload +from .utils.api_recovery import ApiRecoveryManager +from .utils.history_store import HistoryStore instance = PromptServer.instance app = instance.app @@ -85,6 +88,32 @@ jwt_auth = JWTAuth( ) access_control.sync_comfy_users() +api_recovery = None +if scheduler: + api_recovery = ApiRecoveryManager( + scheduler, HistoryStore(HISTORY_FILE), internal_signer.secret, + folder_paths.get_output_directory(), + register_assets=access_control.register_recovered_assets, + notify=instance.queue_updated, + ) + + def observe_api_submission(prompt_id, request, response): + context = get_executing_context() + node_id = context.node_id if context and context.prompt_id == prompt_id else "" + api_recovery.observe(prompt_id, node_id, request, response) + + install_api_audit(scheduler.record_api_call, observe_api_submission) + + async def start_api_recovery(app): + api_recovery.start() + + app.on_startup.append(start_api_recovery) + + async def stop_api_recovery(app): + api_recovery.stop() + + app.on_cleanup.append(stop_api_recovery) + if scheduler: event_relay = EventRelay(instance, scheduler, internal_signer, INSTANCE_PORT) event_relay.install() diff --git a/docs/superpowers/plans/2026-10-28-api-recovery.md b/docs/superpowers/plans/2026-10-28-api-recovery.md new file mode 100644 index 0000000..05da61a --- /dev/null +++ b/docs/superpowers/plans/2026-10-28-api-recovery.md @@ -0,0 +1,29 @@ +# API result recovery implementation plan + +**Goal:** Recover already submitted asynchronous media results after a worker is lost, without resubmitting generation or resuming downstream nodes. + +**Design:** Capture successful, recognized submission responses at the existing HTTP audit boundary. Store task identity and encrypted query credentials in a private scheduler table. A cross-process file lock serializes recovery. Only worker_lost jobs are eligible. Download to deterministic owner-scoped paths using temporary files and atomic rename; merge results idempotently into persistent history. Keep the interrupted workflow failed and label recovered outputs explicitly. + +**Approaches:** Re-running a workflow risks duplicate charges. Replaying audit logs cannot authenticate because secrets are redacted. Selected approach is submission checkpoints plus explicit read-only provider query adapters. The user approved this recovery scope in the preceding conversation and requested implementation. + +**Tech:** Python, SQLite, existing filelock and safe media downloader; cryptography Fernet for credentials, derived from the existing internal secret with domain separation. + +**Scope:** Ark contents/generations/tasks, Wan OPC and Vapeur, Kling official video endpoints and Vapeur submit/describe. Signed Tencent/LAS, synchronous responses, and old submissions without checkpoints are outside automatic recovery. No changes outside Account Manager; no remote restart/deployment. + +## Tasks +- [x] Add explicit provider recognition/query/result parsers, with fixture tests. +- [x] Add durable recovery store/worker and audit observer with encrypted credentials, tests for restart, duplicate processing, retry, isolation, failed provider, and local completion. +- [x] Wire startup and persistent history/assets, preserving failed workflow status and existing outputs. +- [x] Run focused and complete suites, review the patch, document supported protocols and operational limits. + +Validation: 129 Python tests, 6 frontend unit tests, and the browser refresh suite +passed against the installed ComfyUI frontend, including WebSocket disconnect/reconnect. +Provider responses/downloads were simulated; no live paid generation, deployment, +or production restart was performed. + +## Review focus +- Never submit a second generation request; only known GET paths or Vapeur describe POST. +- Credentials never enter ordinary audit logs, history, manifests, or error messages. +- Multiple workers and crashes must not duplicate history entries or expose another user's outputs. +- Provider timeout/authentication failure must not become a completed workflow. +- Live jobs and cancelled jobs must not be recovered; unknown protocols must not be guessed. diff --git a/requirements.txt b/requirements.txt index 71c658c..1260591 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,3 +4,4 @@ PyJWT bcrypt bleach filelock +cryptography>=41 diff --git a/tests/test_api_audit.py b/tests/test_api_audit.py index f3d05fd..fc15388 100644 --- a/tests/test_api_audit.py +++ b/tests/test_api_audit.py @@ -20,6 +20,24 @@ class ApiAuditTests(unittest.TestCase): def tearDown(self): api_audit.clear_current_job() api_audit._recorder = None + api_audit._observer = None + + def test_checkpoint_observer_keeps_audit_redacted_and_recovery_is_not_misattributed(self): + records, observed = [], [] + api_audit._recorder = lambda **item: records.append(item) + api_audit._observer = lambda *args: observed.append(args) + api_audit.set_current_job("job-a") + request = type("Request", (), {"method": "POST", "url": "https://api.example/tasks", + "headers": {"Authorization": "Bearer secret"}, "body": b'{}'})() + response = type("Response", (), {"status_code": 200, "headers": {}, "content": b'{"id":"task-a"}'})() + api_audit._record("requests", request, response) + self.assertEqual("job-a", observed[0][0]) + self.assertEqual("Bearer secret", observed[0][1].headers["Authorization"]) + self.assertNotIn("Bearer secret", str(records)) + with api_audit.suspend_api_audit(): + api_audit._record("requests", request, response) + self.assertEqual(1, len(records)) + self.assertEqual(1, len(observed)) def test_records_original_request_and_response_bodies(self): records = [] diff --git a/tests/test_api_recovery.py b/tests/test_api_recovery.py new file mode 100644 index 0000000..d58f203 --- /dev/null +++ b/tests/test_api_recovery.py @@ -0,0 +1,245 @@ +import importlib +import json +import sys +import tempfile +import types +import unittest +from contextlib import closing +from pathlib import Path +from unittest.mock import patch + +import httpx + +package = types.ModuleType("recovery_test_utils") +package.__path__ = [str(Path(__file__).parents[1] / "utils")] +sys.modules[package.__name__] = package +recovery = importlib.import_module("recovery_test_utils.api_recovery") +SchedulerStore = importlib.import_module("recovery_test_utils.scheduler_store").SchedulerStore +HistoryStore = importlib.import_module("recovery_test_utils.history_store").HistoryStore + + +class ApiRecoveryTests(unittest.TestCase): + def setUp(self): + self.temp = tempfile.TemporaryDirectory() + self.root = Path(self.temp.name) + self.scheduler = SchedulerStore(str(self.root / "scheduler.db")) + self.history = HistoryStore(str(self.root / "history.db")) + self.scheduler.enqueue((0, "job-1", {}, {}, [], {}), "user-a", 8180) + self.set_job_status("running") + self.manager = self.make_manager() + + def tearDown(self): + self.temp.cleanup() + + def make_manager(self): + return recovery.ApiRecoveryManager(self.scheduler, self.history, b"test-secret", str(self.root / "output")) + + def set_job_status(self, status): + with closing(self.scheduler._connect()) as db: + db.execute("UPDATE scheduler_jobs SET status=? WHERE prompt_id='job-1'", (status,)) + + def capture(self): + request = types.SimpleNamespace( + method="POST", url="https://api.example/api/v3/contents/generations/tasks", + headers={"Authorization": "Bearer private-key"}, + ) + response = types.SimpleNamespace(status_code=200, content=b'{"id":"task-1"}') + self.manager.observe("job-1", "7", request, response) + + def rows(self): + with closing(self.scheduler._connect()) as db: + return [dict(row) for row in db.execute("SELECT * FROM scheduler_api_recovery")] + + def download(self, url, path, max_bytes): + Path(path).write_bytes(b"video-data") + return "video/mp4", 10 + + def test_restart_recovers_only_original_task_and_is_idempotent(self): + self.capture() + self.assertNotIn("private-key", str(self.rows())) + self.set_job_status("worker_lost") + restarted = self.make_manager() + with patch.object(restarted, "_query", return_value={"status": "succeeded", "content": {"video_url": "https://cdn.example/a.mp4"}}) as query, patch.object(recovery, "download_remote", side_effect=self.download) as download: + self.assertTrue(restarted.run_once()) + self.assertFalse(restarted.run_once()) + self.assertEqual(1, query.call_count) + self.assertEqual(1, download.call_count) + self.assertEqual("recovered", self.rows()[0]["state"]) + self.assertIsNone(self.rows()[0]["credentials"]) + visible = self.history.query(owner_id="user-a")["job-1"] + self.assertEqual("error", visible["status"]["status_str"]) + self.assertFalse(visible["status"]["completed"]) + self.assertEqual({}, self.history.query(owner_id="user-b")) + videos = [ref for output in visible["outputs"].values() for ref in output.get("video", [])] + self.assertEqual(1, len(videos)) + self.assertTrue((self.root / "output" / videos[0]["subfolder"] / videos[0]["filename"]).is_file()) + + def test_running_completed_cancelled_tasks_never_requery(self): + for status in ("running", "completed", "cancelled", "failed"): + with self.subTest(status=status): + self.set_job_status(status) + self.capture() + with patch.object(self.manager, "_query") as query: + self.manager.run_once() + query.assert_not_called() + + def test_pending_and_network_failure_retry_without_duplicate_submit(self): + self.capture() + self.set_job_status("worker_lost") + with patch.object(self.manager, "_query", return_value={"status": "running"}): + self.manager.run_once() + self.assertEqual("pending", self.rows()[0]["state"]) + with closing(self.scheduler._connect()) as db: + db.execute("UPDATE scheduler_api_recovery SET next_check=0") + with patch.object(self.manager, "_query", side_effect=TimeoutError("private-key must not be logged")): + self.manager.run_once() + self.assertEqual("pending", self.rows()[0]["state"]) + self.assertNotIn("private-key", str(self.rows())) + + def test_provider_failure_is_visible_and_stops_polling(self): + self.capture() + self.set_job_status("worker_lost") + with patch.object(self.manager, "_query", return_value={"status": "failed"}): + self.manager.run_once() + self.assertEqual("failed", self.rows()[0]["state"]) + item = self.history.query(owner_id="user-a")["job-1"] + self.assertIn("第三方任务失败", str(item)) + + def test_cross_worker_lock_prevents_duplicate_queries(self): + self.capture() + self.set_job_status("worker_lost") + sibling = self.make_manager() + with self.manager.lock: + with patch.object(sibling, "_query") as query: + self.assertFalse(sibling.run_once()) + query.assert_not_called() + + def test_duplicate_capture_does_not_reset_recovery(self): + self.capture() + self.capture() + self.assertEqual(1, len(self.rows())) + + def test_audit_boundary_checkpoints_submission_without_logging_credentials(self): + audit = importlib.import_module("recovery_test_utils.api_audit") + records = [] + with patch.object(audit, "_recorder", lambda **record: records.append(record)), \ + patch.object(audit, "_observer", lambda prompt, request, response: self.manager.observe(prompt, "7", request, response)), \ + patch.object(audit, "current_job", return_value="job-1"): + # Use the same HTTP response type consumed by the real audit boundary. + with httpx.Client(transport=httpx.MockTransport( + lambda request: httpx.Response(200, json={"id": "task-1"}) + )) as client: + response = client.post("https://api.example/api/v3/contents/generations/tasks", + headers={"Authorization": "Bearer private-key"}) + audit._record("httpx", response.request, response) + self.assertEqual("task-1", self.rows()[0]["task_id"]) + self.assertNotIn("private-key", str(records)) + + def test_recovered_video_is_previewable_in_installed_jobs_api(self): + sys.path.insert(0, str(Path(__file__).parents[3])) + from comfy_execution.jobs import normalize_history_item + + self.capture() + self.set_job_status("worker_lost") + with patch.object(self.manager, "_query", return_value={"status": "succeeded", "content": {"video_url": "https://cdn.example/a.mp4"}}), patch.object(recovery, "download_remote", side_effect=self.download): + self.manager.run_once() + history = self.history.query(owner_id="user-a", for_jobs=True)["job-1"] + job = normalize_history_item("job-1", history, include_outputs=True) + self.assertEqual("failed", job["status"]) + self.assertEqual(1, job["previewable_outputs_count"]) + self.assertEqual("000.mp4", job["preview_output"]["filename"]) + + def test_crash_after_history_commit_reuses_file_and_history(self): + self.capture() + self.set_job_status("worker_lost") + body = {"status": "succeeded", "content": {"video_url": "https://cdn.example/a.mp4"}} + self.manager.notify = lambda: (_ for _ in ()).throw(RuntimeError("Simulated exit")) + with patch.object(self.manager, "_query", return_value=body), patch.object(recovery, "download_remote", side_effect=self.download) as download: + with self.assertRaises(RuntimeError): + self.manager.run_once() + self.manager.notify = None + self.manager.run_once() + self.assertEqual(1, download.call_count) + item = self.history.query(owner_id="user-a")["job-1"] + self.assertEqual(1, len(item["api_recovery"])) + self.assertEqual(1, len(item["outputs"])) + + def test_history_preserves_previous_outputs(self): + self.history.save("job-1", {"user_id": "user-a", "outputs": {"3": {"text": ["original"]}}, + "status": {"status_str": "error", "completed": False}}, 100) + self.capture() + self.set_job_status("worker_lost") + with patch.object(self.manager, "_query", return_value={"status": "running"}): + self.manager.run_once() + item = self.history.query(owner_id="user-a")["job-1"] + self.assertEqual(["original"], item["outputs"]["3"]["text"]) + + def test_query_uses_original_id_and_never_follows_auth_redirect(self): + self.capture() + row = self.rows()[0] + response = unittest.mock.MagicMock() + response.__enter__.return_value = response + response.status_code = 200 + response.json.return_value = {"status": "running"} + with patch.object(recovery, "_validate_public_http_url"), patch.object(recovery.requests, "request", return_value=response) as request: + self.manager._query(row) + args, kwargs = request.call_args + self.assertEqual(("GET", "https://api.example/api/v3/contents/generations/tasks/task-1"), args) + self.assertFalse(kwargs["allow_redirects"]) + self.assertEqual("Bearer private-key", kwargs["headers"]["Authorization"]) + + def test_auth_failure_stops_and_retains_task_identity_without_credentials(self): + self.capture() + self.set_job_status("worker_lost") + with patch.object(self.manager, "_query", side_effect=recovery.RecoveryQueryError("恢复查询凭证已失效")): + self.manager.run_once() + row = self.rows()[0] + self.assertEqual("failed", row["state"]) + self.assertEqual("task-1", row["task_id"]) + self.assertIsNone(row["credentials"]) + + def test_partial_download_is_removed_and_retried(self): + self.capture() + self.set_job_status("worker_lost") + body = {"status": "succeeded", "content": {"video_url": "https://cdn.example/a.mp4"}} + def interrupted(url, destination, maximum): + Path(destination).write_bytes(b"incomplete") + raise TimeoutError() + with patch.object(self.manager, "_query", return_value=body), patch.object(recovery, "download_remote", side_effect=interrupted): + self.manager.run_once() + self.assertEqual([], list((self.root / "output").rglob("*.mp4"))) + self.assertEqual([], list((self.root / "output").rglob("*.part"))) + self.assertEqual("pending", self.rows()[0]["state"]) + + def test_captured_body_without_auth_or_with_invalid_json_is_ignored(self): + request = types.SimpleNamespace(method="POST", url="https://api.example/api/v3/contents/generations/tasks", headers={}) + for content in (b'{"id":"x"}', b'not json'): + self.manager.observe("job-1", "7", request, types.SimpleNamespace(status_code=200, content=content)) + self.assertEqual([], self.rows()) + + def test_original_worker_finishing_during_query_cancels_recovery(self): + self.capture() + self.set_job_status("worker_lost") + def complete(row): + self.set_job_status("completed") + return {"status": "succeeded", "content": {"video_url": "https://cdn.example/a.mp4"}} + with patch.object(self.manager, "_query", side_effect=complete), patch.object(recovery, "download_remote") as download: + self.manager.run_once() + download.assert_not_called() + self.assertEqual("ignored", self.rows()[0]["state"]) + self.assertEqual({}, self.history.query()) + + def test_successful_history_is_not_overwritten_before_scheduler_completion(self): + self.capture() + self.set_job_status("worker_lost") + self.history.save("job-1", {"user_id": "user-a", "outputs": {"3": {"text": ["success"]}}, + "status": {"status_str": "success", "completed": True}}, 100) + with patch.object(self.manager, "_query", return_value={"status": "running"}): + self.manager.run_once() + item = self.history.query()["job-1"] + self.assertNotIn("api_recovery", item) + self.assertEqual("ignored", self.rows()[0]["state"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_recovery_providers.py b/tests/test_recovery_providers.py new file mode 100644 index 0000000..d1bece7 --- /dev/null +++ b/tests/test_recovery_providers.py @@ -0,0 +1,55 @@ +import importlib.util +import unittest +from pathlib import Path + +spec = importlib.util.spec_from_file_location( + "recovery_providers", Path(__file__).parents[1] / "utils/recovery_providers.py" +) +providers = importlib.util.module_from_spec(spec) +spec.loader.exec_module(providers) + + +class RecoveryProviderTests(unittest.TestCase): + def test_submissions_map_to_query_only_endpoints(self): + cases = [ + ("/api/v3/contents/generations/tasks", {"id": "task-1"}, "ark", "GET", "/api/v3/contents/generations/tasks/task-1", None), + ("/v1/videos/generations", {"output": {"task_id": "task-1"}}, "wan", "GET", "/v1/tasks/task-1", None), + ("/qwen/api/v1/services/aigc/video-generation/generation", {"output": {"task_id": "task-1"}}, "wan", "GET", "/qwen/api/v1/tasks/video/task-1", None), + ("/v1/videos/image2video", {"code": 0, "data": {"task_id": "task-1"}}, "kling", "GET", "/v1/videos/image2video/task-1", None), + ("/v1/kling/text-to-video/submit", {"Response": {"JobId": "task-1"}}, "vapeur_kling", "POST", "/v1/kling/text-to-video/describe", {"JobId": "task-1"}), + ("/v1/videos/generations", {"data": {"task_id": "task-1"}}, "wan", "GET", "/v1/tasks/task-1", None), + ("/v1/videos/generations", {"data": {"id": "task-1"}}, "wan", "GET", "/v1/tasks/task-1", None), + ] + for path, response, provider, method, query, body in cases: + with self.subTest(path=path): + result = providers.submission("https://api.example" + path, response) + self.assertEqual((provider, "task-1", method, "https://api.example" + query, body), result) + + def test_unknown_invalid_and_failed_submissions_are_not_replayed(self): + for path, response in [ + ("/unknown/tasks", {"id": "task-1"}), + ("/api/v3/contents/generations/tasks", {"id": "../delete"}), + ("/v1/videos/image2video", {"code": 12, "data": {"task_id": "task-1"}}), + ]: + self.assertIsNone(providers.submission("https://api.example" + path, response)) + + def test_results_extract_only_generated_media(self): + cases = [ + ("ark", {"status": "succeeded", "content": {"video_url": "https://cdn.example/a.mp4"}}), + ("wan", {"output": {"task_status": "SUCCEEDED", "video_url": "https://cdn.example/a.mp4"}}), + ("kling", {"code": 0, "data": {"task_status": "succeed", "task_result": {"videos": [{"url": "https://cdn.example/a.mp4"}]}}}), + ("vapeur_kling", {"Response": {"Status": "DONE", "ResultVideoUrl": "https://cdn.example/a.mp4"}}), + ] + for provider, body in cases: + body["input"] = {"video_url": "https://cdn.example/input.mp4"} + self.assertEqual(("succeeded", ["https://cdn.example/a.mp4"]), providers.result(provider, body)) + + def test_pending_failed_and_missing_output_are_distinct(self): + self.assertEqual(("pending", []), providers.result("ark", {"status": "running"})) + self.assertEqual(("failed", []), providers.result("ark", {"status": "expired"})) + with self.assertRaises(ValueError): + providers.result("ark", {"status": "succeeded"}) + + +if __name__ == "__main__": + unittest.main() diff --git a/utils/access_control.py b/utils/access_control.py index c14443a..44b90f5 100644 --- a/utils/access_control.py +++ b/utils/access_control.py @@ -44,6 +44,7 @@ class AccessControl: self.instance_port = int(instance_port or 0) self.worker_resource_class = str(worker_resource_class or "default") self.cloud_archive = cloud_archive + self.asset_access = None self._worker_heartbeat = None self._current_user = contextvars.ContextVar("user_id", default=None) @@ -620,6 +621,34 @@ class AccessControl: except Exception: logger.exception("Failed to register persistent history asset %s", path) + def register_recovered_assets(self, assets, prompt_id: str, user_id: str) -> None: + if self.asset_access: + from app.assets.database.queries.records import get_record_by_path_or_none + from app.assets.services.ingest import register_executed_output + from app.database.db import create_session + + with self.asset_access.scope(user_id, write=True): + for reference, path in assets: + with create_session() as session: + existing = get_record_by_path_or_none(session, path) + asset_id = existing.id if existing else None + if not asset_id: + registered = register_executed_output(path, job_id=prompt_id) + if not registered: + raise RuntimeError("Recovered asset registration failed") + asset_id = registered.id + reference["id"] = asset_id + else: + self._register_history_assets(assets, prompt_id, user_id) + if any(not reference.get("id") for reference, _ in assets): + raise RuntimeError("Recovered asset registration failed") + for reference, path in assets: + if reference.get("id"): + self.scheduler.upsert_asset( + reference["id"], user_id, prompt_id, self.instance_port, + dict(reference, file_path=path), + ) + def _save_history_item(self, prompt_id: str) -> None: if not self.__history_store: return diff --git a/utils/api_audit.py b/utils/api_audit.py index 47645a6..19b3b12 100644 --- a/utils/api_audit.py +++ b/utils/api_audit.py @@ -1,8 +1,10 @@ import contextvars import hashlib import json +import logging import threading import time +from contextlib import contextmanager from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit @@ -31,6 +33,17 @@ _fallback_prompt_id = None _fallback_lock = threading.Lock() _recorder = None _installed = False +_observer = None +_suspended = contextvars.ContextVar("account_manager_audit_suspended", default=False) + + +@contextmanager +def suspend_api_audit(): + token = _suspended.set(True) + try: + yield + finally: + _suspended.reset(token) MEDIA_CONTENT_TYPES = ( "image/", @@ -211,9 +224,17 @@ def _audit_body(value, headers=None) -> bytes | None: def _record(client, request, response=None, error: BaseException = None, *, stream=False) -> None: + if _suspended.get(): + return prompt_id = current_job() if not prompt_id or _recorder is None: return + if _observer and response is not None and not stream: + try: + _observer(prompt_id, request, response) + except Exception: + # Audit and execution must continue even if a checkpoint cannot be saved. + logging.getLogger("ComfyUI-Account-Manager").error("API recovery checkpoint could not be saved") try: response_body = None # Reading .content consumes the raw stream used by COS and other downloaders. @@ -257,9 +278,10 @@ def _fallback_request(method, url, kwargs): return request -def install_api_audit(recorder) -> None: - global _installed, _recorder +def install_api_audit(recorder, observer=None) -> None: + global _installed, _recorder, _observer _recorder = recorder + _observer = observer if _installed: return diff --git a/utils/api_recovery.py b/utils/api_recovery.py new file mode 100644 index 0000000..2051393 --- /dev/null +++ b/utils/api_recovery.py @@ -0,0 +1,234 @@ +"""Recover submitted API results, never re-execute a workflow or generation POST.""" + +import base64 +import hashlib +import json +import logging +import os +import re +import threading +import time +from contextlib import closing +from pathlib import Path +from urllib.parse import urlsplit + +import requests +from cryptography.fernet import Fernet +from filelock import FileLock, Timeout + +from .api_audit import suspend_api_audit +from .cloud_archive import download_remote, _validate_public_http_url +from .recovery_providers import submission, result + + +logger = logging.getLogger("ComfyUI-Account-Manager") + + +class RecoveryQueryError(Exception): + pass + + +class RecoverySuperseded(Exception): + pass + + +class ApiRecoveryManager: + def __init__(self, scheduler, history, secret: bytes, output_dir: str, + register_assets=None, notify=None): + self.scheduler = scheduler + self.history = history + self.output_dir = Path(output_dir).resolve() + self.register_assets = register_assets + self.notify = notify + key = hashlib.sha256(b"AccountManager.APIRecovery.v1\0" + secret).digest() + self.cipher = Fernet(base64.urlsafe_b64encode(key)) + self.lock = FileLock(scheduler.database + ".recovery.lock", timeout=0) + self._stop = threading.Event() + self._thread = None + with closing(scheduler._connect()) as db: + db.executescript(""" + CREATE TABLE IF NOT EXISTS scheduler_api_recovery ( + id TEXT PRIMARY KEY, prompt_id TEXT NOT NULL, node_id TEXT NOT NULL, + provider TEXT NOT NULL, task_id TEXT NOT NULL, + query_method TEXT NOT NULL, query_url TEXT NOT NULL, query_body TEXT, + credentials BLOB, state TEXT NOT NULL DEFAULT 'pending', + created_at REAL NOT NULL, next_check REAL NOT NULL DEFAULT 0, + failures INTEGER NOT NULL DEFAULT 0, message TEXT + ); + CREATE INDEX IF NOT EXISTS scheduler_api_recovery_due + ON scheduler_api_recovery(state, next_check); + """) + + def observe(self, prompt_id, node_id, request, response): + if str(request.method).upper() != "POST" or not 200 <= response.status_code < 300: + return + try: + body = json.loads(response.content) + except (ValueError, UnicodeDecodeError): + return + parsed = submission(str(request.url), body) + if parsed is None: + return + provider, task_id, method, url, query_body = parsed + headers = {str(k).lower(): str(v) for k, v in request.headers.items()} + authorization = headers.get("authorization", "") + if not authorization.lower().startswith("bearer "): + return # Signed requests require a provider-specific signer, not replay. + key = hashlib.sha256(f"{prompt_id}\0{url}\0{task_id}".encode()).hexdigest() + credentials = self.cipher.encrypt(authorization.encode()) + with closing(self.scheduler._connect()) as db: + db.execute("PRAGMA synchronous=FULL") + # Owner identity comes exclusively from the queue, never the response. + db.execute(""" + INSERT OR IGNORE INTO scheduler_api_recovery + (id,prompt_id,node_id,provider,task_id,query_method,query_url,query_body,credentials,created_at) + SELECT ?,prompt_id,?,?,?,?,?,?,?,? FROM scheduler_jobs + WHERE prompt_id=? AND status='running' + """, (key, str(node_id or ""), provider, task_id, method, url, + json.dumps(query_body), credentials, time.time(), prompt_id)) + + def start(self): + self._thread = threading.Thread(target=self._run, name="api-result-recovery", daemon=True) + self._thread.start() + + def stop(self): + self._stop.set() + if self._thread: + self._thread.join(timeout=5) + + def _run(self): + while not self._stop.wait(5): + try: + self.run_once() + except Exception: + # Do not log exception text: HTTP errors can contain signed URLs. + logger.error("API result recovery iteration failed; retrying") + + def _query(self, row): + _validate_public_http_url(row["query_url"]) + authorization = self.cipher.decrypt(row["credentials"]).decode() + with requests.request( + row["query_method"], row["query_url"], + headers={"Authorization": authorization}, + json=json.loads(row["query_body"]), timeout=(10, 30), allow_redirects=False, + ) as response: + if response.status_code in (401, 403): + raise RecoveryQueryError("恢复查询凭证已失效或无权限，请人工查询原任务 ID。") + if response.status_code in (404, 410): + raise RecoveryQueryError("第三方任务不存在或已过期，请检查原任务 ID。") + response.raise_for_status() + if not 200 <= response.status_code < 300: + raise ValueError("Unexpected query redirect") + return response.json() + + def _download(self, row, urls): + owner = row["owner_id"] + if not re.fullmatch(r"[A-Za-z0-9_-]+", owner): + raise ValueError("Invalid owner path") + relative = Path(owner) / "recovered" / row["id"] + directory = self.output_dir / relative + directory.mkdir(parents=True, exist_ok=True) + references = [] + assets = [] + for index, url in enumerate(urls): + suffix = Path(urlsplit(url).path).suffix.lower() + if suffix not in (".mp4", ".webm", ".mov"): + suffix = ".mp4" + destination = directory / f"{index:03d}{suffix}" + if not destination.is_file(): + temporary = destination.with_suffix(suffix + ".part") + try: + content_type, size = download_remote(url, str(temporary), 20 * 1024**3) + if not size or (content_type and not content_type.startswith("video/") and content_type != "application/octet-stream"): + raise ValueError("Recovered result is not a video") + os.replace(temporary, destination) + finally: + temporary.unlink(missing_ok=True) + reference = {"filename": destination.name, "subfolder": relative.as_posix(), + "type": "output", "format": "video/" + ("webm" if suffix == ".webm" else "mp4")} + references.append(reference) + assets.append((reference, str(destination))) + self._require_lost(row) + if self.register_assets: + self.register_assets(assets, row["prompt_id"], owner) + return references + + def _publish(self, row, state, message, references=()): + with closing(self.scheduler._connect()) as db: + db.execute("BEGIN IMMEDIATE") + try: + job = db.execute("SELECT status FROM scheduler_jobs WHERE prompt_id=?", (row["prompt_id"],)).fetchone() + published = bool(job and job["status"] == "worker_lost") and self.history.merge_api_recovery(row, state, message, list(references)) + db.execute("COMMIT") + except Exception: + db.execute("ROLLBACK") + raise + if not published: + return False + if self.notify: + self.notify() + return True + + def _require_lost(self, row): + with closing(self.scheduler._connect()) as db: + job = db.execute("SELECT status FROM scheduler_jobs WHERE prompt_id=?", (row["prompt_id"],)).fetchone() + if not job or job["status"] != "worker_lost": + raise RecoverySuperseded() + + def run_once(self): + try: + with self.lock, suspend_api_audit(): + return self._recover_one() + except Timeout: + return False + + def _recover_one(self): + now = time.time() + with closing(self.scheduler._connect()) as db: + db.execute(""" + UPDATE scheduler_api_recovery SET state='ignored', credentials=NULL + WHERE state='pending' AND prompt_id IN + (SELECT prompt_id FROM scheduler_jobs WHERE status IN ('completed','cancelled','failed')) + """) + row = db.execute(""" + SELECT recovery.*, job.owner_id, job.payload, job.submitted_at + FROM scheduler_api_recovery recovery JOIN scheduler_jobs job USING(prompt_id) + WHERE recovery.state='pending' AND recovery.next_check<=? AND job.status='worker_lost' + ORDER BY recovery.next_check, recovery.created_at LIMIT 1 + """, (now,)).fetchone() + if row is None: + return False + row = dict(row) + state, message, references = "pending", "服务器已恢复，正在查询原 API 任务。", [] + failures = row["failures"] + try: + if now - row["created_at"] > 7 * 86400: + raise RecoveryQueryError("自动恢复已超过 7 天，请人工查询原任务 ID。") + outcome, urls = result(row["provider"], self._query(row)) + self._require_lost(row) + if outcome == "succeeded": + references = self._download(row, urls) + state, message = "recovered", "API 生成结果已恢复；原工作流已中断，后续节点未执行。" + elif outcome == "failed": + state, message = "failed", "第三方任务失败、取消或过期；未重新提交生成请求。" + else: + failures = 0 + message = "第三方仍在生成，正在等待原任务完成。" + except RecoverySuperseded: + state, message = "ignored", "原任务状态已更新，停止恢复。" + except RecoveryQueryError as error: + state, message = "failed", str(error) + except Exception: + failures += 1 + message = "查询或下载暂时失败，将自动重试；未重新提交生成请求。" + if failures >= 20: + state, message = "failed", "连续 20 次恢复失败，请人工查询原任务 ID；未重新提交生成请求。" + # Persist history before acknowledging; a crash repeats an idempotent merge. + if state != "ignored" and not self._publish(row, state, message, references): + state = "ignored" + with closing(self.scheduler._connect()) as db: + db.execute(""" + UPDATE scheduler_api_recovery SET state=?, message=?, next_check=?, failures=?, + credentials=CASE WHEN ?='pending' THEN credentials ELSE NULL END WHERE id=? + """, (state, message, now + min(1800, 30 * 2 ** min(failures, 6)), failures, state, row["id"])) + return True diff --git a/utils/history_store.py b/utils/history_store.py index 37df066..493779c 100644 --- a/utils/history_store.py +++ b/utils/history_store.py @@ -1,6 +1,7 @@ import json import os import sqlite3 +import time from contextlib import closing @@ -158,6 +159,42 @@ class HistoryStore: with connection: connection.execute("DELETE FROM history WHERE prompt_id = ?", (prompt_id,)) + def merge_api_recovery(self, task: dict, state: str, message: str, videos: list) -> bool: + """Add recovered output without claiming that downstream nodes completed.""" + with closing(self._connect()) as connection: + with connection: + connection.execute("BEGIN IMMEDIATE") + existing = connection.execute( + "SELECT owner_id,data FROM history WHERE prompt_id=?", (task["prompt_id"],) + ).fetchone() + if existing and existing[0] != task["owner_id"]: + raise ValueError("Recovery history owner mismatch") + item = json.loads(existing[1]) if existing else { + "prompt": json.loads(task["payload"])[:5], "outputs": {}, + "user_id": task["owner_id"], + "status": {"status_str": "error", "completed": False, "messages": [ + ["execution_start", {"timestamp": int(task["submitted_at"] * 1000)}], + ["execution_error", {"node_id": task["node_id"], "node_type": "APIRecovery", + "exception_type": "WorkerInterrupted", "exception_message": "服务器中断，工作流未完整执行。", + "timestamp": int(time.time() * 1000)}], + ]}, + } + if (item.get("status") or {}).get("status_str") == "success": + return False + item.setdefault("api_recovery", {})[task["id"]] = { + "provider": task["provider"], "task_id": task["task_id"], + "node_id": task["node_id"], "state": state, "message": message, + } + output = {"text": [message]} + if videos: + output["video"] = videos + item.setdefault("outputs", {})["api_recovery_" + task["id"]] = output + connection.execute(""" + INSERT INTO history(prompt_id,owner_id,data) VALUES(?,?,?) + ON CONFLICT(prompt_id) DO UPDATE SET data=excluded.data + """, (task["prompt_id"], task["owner_id"], json.dumps(item, ensure_ascii=False))) + return True + def delete_owner(self, owner_id: str) -> None: with closing(self._connect()) as connection: with connection: diff --git a/utils/recovery_providers.py b/utils/recovery_providers.py new file mode 100644 index 0000000..6520f9c --- /dev/null +++ b/utils/recovery_providers.py @@ -0,0 +1,89 @@ +"""Explicit query-only adapters for the async video protocols used by our nodes.""" + +import re +from urllib.parse import urlsplit, urlunsplit + + +KLING_VIDEO_KINDS = { + "text2video", "image2video", "multi-image2video", "omni-video", + "video-extend", "motion-control", "multi-elements", "lip-sync", +} +VAPEUR_VIDEO_KINDS = { + "text-to-video", "image-to-video", "multi-image-to-video", + "video-motion-control", "video-extend", "video-edit", +} + + +def submission(url: str, body: dict): + parts = urlsplit(url) + if parts.scheme != "https" or parts.username or parts.password or parts.query or parts.fragment: + return None + if not isinstance(body, dict) or body.get("error") or body.get("code") not in (None, 0, "0"): + return None + path = parts.path.rstrip("/") + method, query_body = "GET", None + if path.endswith("/contents/generations/tasks"): + provider, task_id = "ark", body.get("id") + query_path = path + "/{}" + elif path.endswith("/v1/videos/generations"): + provider = "wan" + envelope = body.get("data") or body + output = envelope.get("output") or envelope + task_id = envelope.get("task_id") or envelope.get("id") or output.get("task_id") or output.get("id") + query_path = path.removesuffix("/videos/generations") + "/tasks/{}" + elif path.endswith("/qwen/api/v1/services/aigc/video-generation/generation"): + provider, task_id = "wan", (body.get("output") or {}).get("task_id") + query_path = path.removesuffix("/services/aigc/video-generation/generation") + "/tasks/video/{}" + elif path.startswith("/v1/videos/") and path.removeprefix("/v1/videos/") in KLING_VIDEO_KINDS: + provider, task_id = "kling", (body.get("data") or {}).get("task_id") + query_path = path + "/{}" + elif path.startswith("/v1/kling/") and path.endswith("/submit") and path.split("/")[-2] in VAPEUR_VIDEO_KINDS: + provider, task_id = "vapeur_kling", (body.get("Response") or body).get("JobId") + if (body.get("Response") or body).get("Error"): + return None + method, query_path = "POST", path.removesuffix("/submit") + "/describe" + query_body = {"JobId": task_id} + else: + return None + if not isinstance(task_id, str) or not re.fullmatch(r"[A-Za-z0-9_-]{1,256}", task_id): + return None + query_url = urlunsplit((parts.scheme, parts.netloc, query_path.format(task_id), "", "")) + return provider, task_id, method, query_url, query_body + + +def result(provider: str, body: dict) -> tuple[str, list[str]]: + if not isinstance(body, dict): + raise ValueError("Invalid provider response") + if body.get("error") or body.get("code") not in (None, 0, "0"): + raise ValueError("Provider query rejected") + if provider == "ark": + status = body.get("status") + urls = [(body.get("content") or {}).get("video_url")] + elif provider == "wan": + envelope = body.get("data") or body + output = envelope.get("output") or envelope + status = envelope.get("task_status") or envelope.get("status") or output.get("task_status") or output.get("status") + urls = [output.get("video_url") or output.get("url")] + elif provider == "kling": + data = body.get("data") or {} + status = data.get("task_status") + urls = [item.get("url") for item in (data.get("task_result") or {}).get("videos", [])] + elif provider == "vapeur_kling": + data = body.get("Response") or body + if data.get("Error"): + raise ValueError("Provider query rejected") + status = "failed" if data.get("ErrorCode") else data.get("Status") or data.get("TaskStatus") + urls = [data.get("ResultVideoUrl")] + else: + raise ValueError("Unsupported recovery provider") + status = str(status or "").lower() + if status in {"fail", "failed", "error", "aborted", "cancelled", "canceled", "expired"}: + return "failed", [] + if status in {"succeeded", "succeed", "success", "completed", "done", "finish"}: + urls = list(dict.fromkeys(url for url in urls if isinstance(url, str) and url.startswith("https://"))) + if not urls: + raise ValueError("Completed task returned no downloadable media") + return "succeeded", urls + if status in {"submitted", "queued", "pending", "running", "processing", "in_progress", "created", "waiting"}: + return "pending", [] + raise ValueError("Unknown provider task status")