from __future__ import annotations import base64 import json import os import shutil import sqlite3 import tempfile import threading import time from collections.abc import Callable from contextlib import closing from datetime import datetime from pathlib import Path from typing import Any from services.auth_service import DAILY_GIFT_ACCOUNT_POOL_MODEL_ID from services.config import DATA_DIR, config from services.content_filter import request_text from services.error_message_service import error_message_service from services.image_storage_service import image_storage_service from services.log_service import LOG_TYPE_CALL, log_service from services.protocol import openai_v1_image_edit, openai_v1_image_generations from services.protocol.conversation import is_image_quota_exhausted_message from services.sqlite_mirror_service import sqlite_mirror_service from services.time_utils import china_from_timestamp_string, china_now_string from utils.sensitive_data import sanitize_sensitive_text TASK_STATUS_QUEUED = "queued" TASK_STATUS_RUNNING = "running" TASK_STATUS_SUCCESS = "success" TASK_STATUS_ERROR = "error" TERMINAL_STATUSES = {TASK_STATUS_SUCCESS, TASK_STATUS_ERROR} UNFINISHED_STATUSES = {TASK_STATUS_QUEUED, TASK_STATUS_RUNNING} TASK_CLEANUP_INTERVAL_SECONDS = 60 * 60 def _now_iso() -> str: return china_now_string() def _credit_amount(value: object) -> int | float: try: amount = round(max(0.0, float(value or 0)), 2) except (TypeError, ValueError): amount = 0.0 return int(amount) if amount.is_integer() else amount def _timestamp(value: object) -> float: if not isinstance(value, str) or not value.strip(): return 0.0 for fmt in ("%Y-%m-%d %H:%M:%S", "%Y-%m-%dT%H:%M:%S.%f", "%Y-%m-%dT%H:%M:%S"): try: return datetime.strptime(value[:26], fmt).timestamp() except ValueError: continue try: return datetime.fromisoformat(value.replace("Z", "+00:00")).timestamp() except Exception: return 0.0 def _clean(value: object, default: str = "") -> str: return str(value or default).strip() def _database_write_diagnostics(path: Path, exc: BaseException) -> dict[str, Any]: """Capture the filesystem state at the instant a SQLite write fails.""" def stat_payload(candidate: Path) -> dict[str, Any]: try: info = candidate.stat() return { "exists": True, "mode": oct(info.st_mode & 0o7777), "uid": getattr(info, "st_uid", None), "gid": getattr(info, "st_gid", None), "size": info.st_size, } except FileNotFoundError: return {"exists": False} except Exception as stat_exc: return {"stat_error": f"{type(stat_exc).__name__}: {stat_exc}"} parent = path.parent result: dict[str, Any] = { "event": "image_task_database_write_failed", "error": f"{type(exc).__name__}: {exc}", "database": str(path), "database_stat": stat_payload(path), "parent_stat": stat_payload(parent), "journal_stat": stat_payload(Path(f"{path}-journal")), "wal_stat": stat_payload(Path(f"{path}-wal")), "shm_stat": stat_payload(Path(f"{path}-shm")), "process_uid": os.getuid() if hasattr(os, "getuid") else None, "process_gid": os.getgid() if hasattr(os, "getgid") else None, "file_writable": os.access(path, os.W_OK), "parent_writable": os.access(parent, os.W_OK), } try: usage = shutil.disk_usage(parent) result["disk"] = {"total": usage.total, "used": usage.used, "free": usage.free} except Exception as disk_exc: result["disk_error"] = f"{type(disk_exc).__name__}: {disk_exc}" try: fs = os.statvfs(parent) result["filesystem"] = { "blocks": fs.f_blocks, "blocks_free": fs.f_bfree, "blocks_available": fs.f_bavail, "inodes": fs.f_files, "inodes_free": fs.f_ffree, "flags": fs.f_flag, } except Exception as fs_exc: result["filesystem_error"] = f"{type(fs_exc).__name__}: {fs_exc}" try: resolved_parent = str(parent.resolve()) best_mount: dict[str, str] | None = None for line in Path("/proc/mounts").read_text(encoding="utf-8", errors="replace").splitlines(): fields = line.split() if len(fields) < 4: continue mount_point = fields[1].replace("\\040", " ") if resolved_parent == mount_point or resolved_parent.startswith(mount_point.rstrip("/") + "/"): if best_mount is None or len(mount_point) > len(best_mount["mount_point"]): best_mount = { "mount_point": mount_point, "filesystem_type": fields[2], "options": fields[3], } result["mount"] = best_mount except Exception as mount_exc: result["mount_error"] = f"{type(mount_exc).__name__}: {mount_exc}" probe = parent / f".image-task-write-probe-{os.getpid()}-{threading.get_ident()}" try: probe.write_bytes(b"ok") probe.unlink(missing_ok=True) result["directory_write_probe"] = "ok" except Exception as probe_exc: result["directory_write_probe"] = f"{type(probe_exc).__name__}: {probe_exc}" try: probe.unlink(missing_ok=True) except Exception: pass return result def _owner_id(identity: dict[str, object]) -> str: return _clean(identity.get("id")) or "anonymous" def _memory_conversation_id(value: object) -> str: return _clean(value) or "default" def _account_pool_strategy(value: object) -> str: normalized = _clean(value).lower().replace("-", "_") return normalized if normalized in {"own_first", "admin_first", "shared_first", "own_only", "admin_only", "shared_only"} else "own_first" def _task_key(owner_id: str, task_id: str) -> str: return f"{owner_id}:{task_id}" def _collect_image_urls(data: list[Any]) -> list[str]: urls: list[str] = [] for item in data: if isinstance(item, dict): url = item.get("url") if isinstance(url, str) and url: urls.append(url) return urls def _count_image_items(data: object) -> int: if not isinstance(data, list): return 0 count = 0 for item in data: if isinstance(item, dict) and (_clean(item.get("url")) or _clean(item.get("b64_json"))): count += 1 return count def _materialize_inline_task_images(data: object, *, base_url: str = "") -> object: """Move inline task images out of the task payload before it reaches the UI. The image workspace polls task JSON. Keeping a multi-megabyte ``b64_json`` field there makes every successful poll expensive and also causes the browser history cache to duplicate the image. The interactive workspace always consumes a task through a URL, so save inline output into the normal image storage and publish that URL instead. A storage failure must not turn an already-generated image into a failed task; in that rare case retain the original inline result as a backwards-compatible fallback. """ if not isinstance(data, list): return data materialized: list[object] = [] for item in data: if not isinstance(item, dict): materialized.append(item) continue encoded = _clean(item.get("b64_json")) if not encoded: materialized.append(item) continue try: image_bytes = base64.b64decode(encoded, validate=False) if not image_bytes: raise ValueError("empty image data") stored = image_storage_service.save(image_bytes, base_url=base_url or None) normalized = dict(item) normalized.pop("b64_json", None) normalized["url"] = stored.url materialized.append(normalized) except Exception as exc: print(f"[image-task] inline image storage skipped: {exc}") materialized.append(item) return materialized def _collect_third_party_image_api(result: object) -> dict[str, Any]: if isinstance(result, dict) and isinstance(result.get("_third_party_image_api"), dict): return result["_third_party_image_api"] return {} def _public_task(task: dict[str, Any]) -> dict[str, Any]: item = { "id": task.get("id"), "status": task.get("status"), "mode": task.get("mode"), "model": task.get("model"), "size": task.get("size"), "created_at": task.get("created_at"), "started_at": task.get("started_at"), "updated_at": task.get("updated_at"), } if task.get("prompt"): item["prompt"] = task.get("prompt") if task.get("quality"): item["quality"] = task.get("quality") if task.get("resolution"): item["resolution"] = task.get("resolution") if task.get("progress"): item["progress"] = task.get("progress") if task.get("progress_profile"): item["progress_profile"] = task.get("progress_profile") if isinstance(task.get("progress_events"), list): item["progress_events"] = [ event for event in task.get("progress_events", []) if isinstance(event, dict) and _clean(event.get("step")) ] if task.get("duration_ms") is not None: item["duration_ms"] = task.get("duration_ms") if task.get("conversation_id"): item["conversation_id"] = task.get("conversation_id") if task.get("status") in UNFINISHED_STATUSES: base_ts = task.get("started_ts") if task.get("status") == TASK_STATUS_RUNNING else task.get("created_ts") if not base_ts: base_ts = _timestamp(task.get("started_at") if task.get("status") == TASK_STATUS_RUNNING else task.get("created_at")) if base_ts: item["elapsed_secs"] = round(max(0.0, time.time() - float(base_ts)), 1) if task.get("data") is not None: item["data"] = task.get("data") if task.get("error"): # 前台只展示管理员配置过的大白话；未命中规则时保留原文， # 同时由 error_message_service 记下这条待配置报错。 item["error"] = error_message_service.translate( sanitize_sensitive_text(task.get("error")), channel="image_task", ) if task.get("upstream_note"): # 上游 assistant 的原话，不做文案归一（那不是我们的固定话术）， # 但要去掉可能夹带的邮箱/密钥等敏感串。 item["upstream_note"] = sanitize_sensitive_text(task.get("upstream_note")) if "can_continue_waiting" in task: # 显式下发：上游已给终局结论（只回文字没生图）时收掉「继续等待」按钮。 item["can_continue_waiting"] = bool(task.get("can_continue_waiting")) if task.get("can_view_official") is False: # 显式下发：账号生图配额耗尽等「官网没有可核对内容」的失败， # 收掉失败卡片上的「查看官网对话」入口（否则点进去只有同一句提示）。 item["can_view_official"] = False return item class _SharedMemoryConnection: """只读目录 / 磁盘满等环境下所有落盘位置都不可用时的内存兜底连接包装。 ``close()`` 是 no-op，保证 ``contextlib.closing`` 不会真正关掉共享连接； 事务语义（commit / rollback）委托给底层连接。 """ def __init__(self, connection: sqlite3.Connection) -> None: self._connection = connection def close(self) -> None: return None def execute(self, sql: str, parameters: Any = ()) -> sqlite3.Cursor: return self._connection.execute(sql, parameters) def executemany(self, sql: str, seq_of_parameters: Any) -> sqlite3.Cursor: return self._connection.executemany(sql, seq_of_parameters) def __enter__(self) -> sqlite3.Connection: return self._connection.__enter__() def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> bool: return self._connection.__exit__(exc_type, exc, tb) class ImageTaskService: def __init__( self, path: Path, *, generation_handler: Callable[[dict[str, Any]], dict[str, Any]] = openai_v1_image_generations.handle, edit_handler: Callable[[dict[str, Any]], dict[str, Any]] = openai_v1_image_edit.handle, retention_days_getter: Callable[[], int] | None = None, ): self.path = path self.db_path = path.with_suffix(".db") self._memory_database: sqlite3.Connection | None = None self.generation_handler = generation_handler self.edit_handler = edit_handler self.retention_days_getter = retention_days_getter or (lambda: config.image_retention_days) self._lock = threading.RLock() self._tasks: dict[str, dict[str, Any]] = {} self._last_cleanup_at = time.monotonic() try: self.path.parent.mkdir(parents=True, exist_ok=True) except OSError as exc: # 目录本身建不出来（只读挂载/权限）时不要在这里崩： # 数据库位置还有兜底逻辑，JSON 只是 legacy 导入来源。 print(f"[image-tasks] Cannot create data directory {self.path.parent}: {exc}") with self._lock: self._prepare_database_locked() self._tasks = self._load_locked() changed = self._recover_unfinished_locked() changed = self._cleanup_locked() or changed if changed: self._save_locked() def submit_generation( self, identity: dict[str, object], *, client_task_id: str, prompt: str, model: str, size: str | None, base_url: str, memory_enabled: bool = False, memory_conversation_id: str = "", memory_reset: bool = False, account_pool_strategy: str = "own_first", credit_user_id: str = "", credit_amount: int | float = 0, credit_model_id: str = "", credit_pricing: dict[str, Any] | None = None, quality: str = "", resolution: str = "", third_party_image_group_id: str = "", third_party_wallet_user_id: str = "", third_party_wallet_group_id: str = "", third_party_wallet_model_id: str = "", third_party_wallet_amount: int | float = 0, third_party_wallet_cost_per_image: int | float = 1, third_party_wallet_fallback_credit_user_id: str = "", third_party_wallet_fallback_credit_amount: int | float = 0, third_party_wallet_fallback_credit_model_id: str = "", third_party_wallet_fallback_credit_pricing: dict[str, Any] | None = None, internal_model_limit_user_id: str = "", internal_model_limit_model_id: str = "", internal_model_limit_amount: int = 0, ) -> dict[str, Any]: payload = { "prompt": prompt, "model": model, "n": 1, "size": size, "quality": quality, "resolution": resolution, "third_party_image_group_id": third_party_image_group_id, "response_format": "url", "base_url": base_url, "account_owner_id": _owner_id(identity), "account_pool_strategy": _account_pool_strategy(account_pool_strategy), "account_allow_shared": not memory_enabled, } if memory_enabled: payload["memory"] = True payload["memory_owner_id"] = _owner_id(identity) payload["memory_conversation_id"] = _memory_conversation_id(memory_conversation_id) if memory_reset: payload["memory_reset"] = True return self._submit( identity, client_task_id=client_task_id, mode="generate", payload=payload, credit_user_id=credit_user_id, credit_amount=credit_amount, credit_model_id=credit_model_id, credit_pricing=credit_pricing, third_party_wallet_user_id=third_party_wallet_user_id, third_party_wallet_group_id=third_party_wallet_group_id, third_party_wallet_model_id=third_party_wallet_model_id, third_party_wallet_amount=third_party_wallet_amount, third_party_wallet_cost_per_image=third_party_wallet_cost_per_image, third_party_wallet_fallback_credit_user_id=third_party_wallet_fallback_credit_user_id, third_party_wallet_fallback_credit_amount=third_party_wallet_fallback_credit_amount, third_party_wallet_fallback_credit_model_id=third_party_wallet_fallback_credit_model_id, third_party_wallet_fallback_credit_pricing=third_party_wallet_fallback_credit_pricing, internal_model_limit_user_id=internal_model_limit_user_id, internal_model_limit_model_id=internal_model_limit_model_id, internal_model_limit_amount=internal_model_limit_amount, ) def submit_edit( self, identity: dict[str, object], *, client_task_id: str, prompt: str, model: str, size: str | None, base_url: str, images: list[tuple[bytes, str, str]], masks: list[tuple[bytes, str, str]] | None = None, memory_enabled: bool = False, memory_conversation_id: str = "", memory_reset: bool = False, account_pool_strategy: str = "own_first", credit_user_id: str = "", credit_amount: int | float = 0, credit_model_id: str = "", credit_pricing: dict[str, Any] | None = None, quality: str = "", resolution: str = "", third_party_image_group_id: str = "", third_party_wallet_user_id: str = "", third_party_wallet_group_id: str = "", third_party_wallet_model_id: str = "", third_party_wallet_amount: int | float = 0, third_party_wallet_cost_per_image: int | float = 1, third_party_wallet_fallback_credit_user_id: str = "", third_party_wallet_fallback_credit_amount: int | float = 0, third_party_wallet_fallback_credit_model_id: str = "", third_party_wallet_fallback_credit_pricing: dict[str, Any] | None = None, internal_model_limit_user_id: str = "", internal_model_limit_model_id: str = "", internal_model_limit_amount: int = 0, ) -> dict[str, Any]: payload = { "prompt": prompt, "images": images, "mask": masks or [], "model": model, "n": 1, "size": size, "quality": quality, "resolution": resolution, "third_party_image_group_id": third_party_image_group_id, "response_format": "url", "base_url": base_url, "account_owner_id": _owner_id(identity), "account_pool_strategy": _account_pool_strategy(account_pool_strategy), "account_allow_shared": not memory_enabled, } if memory_enabled: payload["memory"] = True payload["memory_owner_id"] = _owner_id(identity) payload["memory_conversation_id"] = _memory_conversation_id(memory_conversation_id) if memory_reset: payload["memory_reset"] = True return self._submit( identity, client_task_id=client_task_id, mode="edit", payload=payload, credit_user_id=credit_user_id, credit_amount=credit_amount, credit_model_id=credit_model_id, credit_pricing=credit_pricing, third_party_wallet_user_id=third_party_wallet_user_id, third_party_wallet_group_id=third_party_wallet_group_id, third_party_wallet_model_id=third_party_wallet_model_id, third_party_wallet_amount=third_party_wallet_amount, third_party_wallet_cost_per_image=third_party_wallet_cost_per_image, third_party_wallet_fallback_credit_user_id=third_party_wallet_fallback_credit_user_id, third_party_wallet_fallback_credit_amount=third_party_wallet_fallback_credit_amount, third_party_wallet_fallback_credit_model_id=third_party_wallet_fallback_credit_model_id, third_party_wallet_fallback_credit_pricing=third_party_wallet_fallback_credit_pricing, internal_model_limit_user_id=internal_model_limit_user_id, internal_model_limit_model_id=internal_model_limit_model_id, internal_model_limit_amount=internal_model_limit_amount, ) def list_tasks(self, identity: dict[str, object], task_ids: list[str]) -> dict[str, Any]: owner = _owner_id(identity) requested_ids = [_clean(task_id) for task_id in task_ids if _clean(task_id)] with self._lock: if self._cleanup_if_due_locked(): self._save_locked() items = [] missing_ids = [] for task_id in requested_ids: task = self._tasks.get(_task_key(owner, task_id)) if task is None: missing_ids.append(task_id) else: items.append(_public_task(task)) if not requested_ids: items = [ _public_task(task) for task in self._tasks.values() if task.get("owner_id") == owner ] items.sort(key=lambda item: str(item.get("updated_at") or ""), reverse=True) missing_ids = [] return {"items": items, "missing_ids": missing_ids} def find_task_by_conversation( self, identity: dict[str, object], conversation_id: str, ) -> dict[str, Any] | None: """按上游 conversation_id 找回当前用户名下的原始任务。 取证接口用它做归属校验：只有这条会话确实是当前用户提交的任务， 才允许把它变成官网分享链接，避免拿到别人的 conversation_id 横向越权。 返回的是内部任务字典（含 account_email），仅供服务端使用，不要直接回给前台。 """ owner = _owner_id(identity) target = _clean(conversation_id) if not target: return None with self._lock: for task in self._tasks.values(): if task.get("owner_id") != owner: continue if _clean(task.get("conversation_id")) == target: return dict(task) return None def find_owned_conversation( self, identity: dict[str, object], conversation_id: str, ) -> dict[str, Any] | None: """按上游 conversation_id 找回当前用户名下的会话：生图任务优先，官网取证台账兜底。 归属校验不能只看生图任务。任务台账按 ``image_retention_days`` 清理，且只服务 「任务状态」这件事；而官网取证台账留着这条会话的 ``owner_id`` 和「已回收」状态 （保留期通常只有几小时，比任务存活得短）。只看任务的话，用户点「查看官网」会拿到 一句像越权警告的 404，真实原因却是"过了保留期被清理"。 """ task = self.find_task_by_conversation(identity, conversation_id) if task is not None: return task from services.share_evidence_service import share_evidence_service record = share_evidence_service.snapshot(conversation_id) if not record: return None if _clean(record.get("owner_id")) != _owner_id(identity): return None return { "id": "", "owner_id": _clean(record.get("owner_id")), "conversation_id": _clean(record.get("conversation_id")), "account_email": _clean(record.get("account_email")), } def _submit( self, identity: dict[str, object], *, client_task_id: str, mode: str, payload: dict[str, Any], credit_user_id: str = "", credit_amount: int | float = 0, credit_model_id: str = "", credit_pricing: dict[str, Any] | None = None, third_party_wallet_user_id: str = "", third_party_wallet_group_id: str = "", third_party_wallet_model_id: str = "", third_party_wallet_amount: int | float = 0, third_party_wallet_cost_per_image: int | float = 1, third_party_wallet_fallback_credit_user_id: str = "", third_party_wallet_fallback_credit_amount: int | float = 0, third_party_wallet_fallback_credit_model_id: str = "", third_party_wallet_fallback_credit_pricing: dict[str, Any] | None = None, internal_model_limit_user_id: str = "", internal_model_limit_model_id: str = "", internal_model_limit_amount: int = 0, ) -> dict[str, Any]: task_id = _clean(client_task_id) if not task_id: raise ValueError("client_task_id is required") owner = _owner_id(identity) key = _task_key(owner, task_id) now = _now_iso() should_start = False with self._lock: cleaned = self._cleanup_if_due_locked() task = self._tasks.get(key) if task is not None: if cleaned: self._save_locked() return _public_task(task) normalized_credit_user_id = _clean(credit_user_id) normalized_credit_amount = _credit_amount(credit_amount) normalized_credit_model_id = _clean(credit_model_id) or DAILY_GIFT_ACCOUNT_POOL_MODEL_ID normalized_fallback_credit_user_id = _clean(third_party_wallet_fallback_credit_user_id) normalized_fallback_credit_amount = _credit_amount(third_party_wallet_fallback_credit_amount) normalized_fallback_credit_model_id = _clean(third_party_wallet_fallback_credit_model_id) or DAILY_GIFT_ACCOUNT_POOL_MODEL_ID daily_reserved = 0 paid_reserved = 0 package_reservations: list[dict[str, object]] = [] grant_reservations: list[dict[str, object]] = [] normalized_wallet_user_id = _clean(third_party_wallet_user_id) normalized_wallet_group_id = _clean(third_party_wallet_group_id) normalized_wallet_model_id = _clean(third_party_wallet_model_id) normalized_wallet_amount = _credit_amount(third_party_wallet_amount) normalized_wallet_cost_per_image = _credit_amount(third_party_wallet_cost_per_image) normalized_internal_limit_user_id = _clean(internal_model_limit_user_id) normalized_internal_limit_model_id = _clean(internal_model_limit_model_id) normalized_internal_limit_amount = max(0, int(internal_model_limit_amount or 0)) normalized_internal_limit_scope = "" if normalized_wallet_user_id and normalized_wallet_group_id and normalized_wallet_model_id and normalized_wallet_amount > 0: from services.auth_service import auth_service try: auth_service.reserve_third_party_image_wallet( normalized_wallet_user_id, group_id=normalized_wallet_group_id, model_id=normalized_wallet_model_id, amount=normalized_wallet_amount, ) except ValueError: if not normalized_fallback_credit_user_id or normalized_fallback_credit_amount <= 0: raise normalized_wallet_user_id = "" normalized_wallet_group_id = "" normalized_wallet_model_id = "" normalized_wallet_amount = 0 normalized_wallet_cost_per_image = 0 normalized_credit_user_id = normalized_fallback_credit_user_id normalized_credit_amount = normalized_fallback_credit_amount normalized_credit_model_id = normalized_fallback_credit_model_id if isinstance(third_party_wallet_fallback_credit_pricing, dict): credit_pricing = dict(third_party_wallet_fallback_credit_pricing) if normalized_credit_user_id and normalized_credit_amount > 0: from services.auth_service import auth_service reservation = auth_service.reserve_credits( normalized_credit_user_id, normalized_credit_amount, model_id=normalized_credit_model_id, ).get("_credit_reservation", {}) if not isinstance(reservation, dict): reservation = {} daily_reserved = _credit_amount(reservation.get("daily")) paid_reserved = _credit_amount(reservation.get("permanent_paid") if "permanent_paid" in reservation else reservation.get("paid")) package_reservations = reservation.get("packages") if isinstance(reservation.get("packages"), list) else [] grant_reservations = reservation.get("grants") if isinstance(reservation.get("grants"), list) else [] if normalized_internal_limit_user_id and normalized_internal_limit_model_id and normalized_internal_limit_amount > 0: from services.auth_service import auth_service reservation = auth_service.reserve_internal_image_model_daily_limit( normalized_internal_limit_user_id, model_id=normalized_internal_limit_model_id, amount=normalized_internal_limit_amount, ).get("_internal_image_model_daily_limit_reservation", {}) reserved_internal_limit_amount = max(0, int(reservation.get("amount") or 0)) if isinstance(reservation, dict) else 0 if reserved_internal_limit_amount <= 0: normalized_internal_limit_user_id = "" normalized_internal_limit_model_id = "" normalized_internal_limit_amount = 0 normalized_internal_limit_scope = "" else: normalized_internal_limit_amount = reserved_internal_limit_amount normalized_internal_limit_scope = _clean(reservation.get("scope")) task = { "id": task_id, "owner_id": owner, "status": TASK_STATUS_QUEUED, "mode": mode, "model": _clean(payload.get("model"), "gpt-image-2"), "size": _clean(payload.get("size")), "progress_profile": self._initial_progress_profile(payload), "created_at": now, "updated_at": now, "created_ts": time.time(), "updated_ts": time.time(), "progress_events": [], "base_url": _clean(payload.get("base_url")), } prompt = _clean(payload.get("prompt")) if prompt: task["prompt"] = prompt quality = _clean(payload.get("quality")) if quality: task["quality"] = quality resolution = _clean(payload.get("resolution")) if resolution: task["resolution"] = resolution if normalized_credit_user_id and normalized_credit_amount > 0: task["credit_user_id"] = normalized_credit_user_id task["credit_reserved"] = normalized_credit_amount task["credit_daily_reserved"] = daily_reserved task["credit_paid_reserved"] = paid_reserved task["credit_package_reservations"] = package_reservations task["credit_grant_reservations"] = grant_reservations task["credit_settled"] = False if isinstance(credit_pricing, dict): task["credit_pricing"] = dict(credit_pricing) if normalized_wallet_user_id and normalized_wallet_amount > 0: task["third_party_wallet_user_id"] = normalized_wallet_user_id task["third_party_wallet_group_id"] = normalized_wallet_group_id task["third_party_wallet_model_id"] = normalized_wallet_model_id task["third_party_wallet_reserved"] = normalized_wallet_amount task["third_party_wallet_cost_per_image"] = normalized_wallet_cost_per_image task["third_party_wallet_settled"] = False if normalized_internal_limit_user_id and normalized_internal_limit_model_id and normalized_internal_limit_amount > 0: task["internal_model_limit_user_id"] = normalized_internal_limit_user_id task["internal_model_limit_model_id"] = normalized_internal_limit_model_id task["internal_model_limit_reserved"] = normalized_internal_limit_amount task["internal_model_limit_scope"] = normalized_internal_limit_scope task["internal_model_limit_settled"] = False self._tasks[key] = task self._save_locked([task]) should_start = True if should_start: thread = threading.Thread( target=self._run_task, args=(key, mode, payload, dict(identity), _clean(payload.get("model"), "gpt-image-2")), name=f"image-task-{task_id[:16]}", daemon=True, ) thread.start() return _public_task(task) def _run_task( self, key: str, mode: str, payload: dict[str, Any], identity: dict[str, object], model: str, ) -> None: started = time.time() self._update_task(key, status=TASK_STATUS_RUNNING, started_at=_now_iso(), started_ts=started, error="") account_email = "" def progress_callback(step: str) -> None: cleaned_step = _clean(step) or "generating" if cleaned_step == "third_party_image_api": self._set_progress_profile(key, "api", reset_events=True) return self._record_progress(key, cleaned_step, set_started_ts=cleaned_step == "image_stream_resolve_start") payload_with_progress = {**payload, "progress_callback": progress_callback} try: handler = self.edit_handler if mode == "edit" else self.generation_handler result = handler(payload_with_progress) if not isinstance(result, dict): raise RuntimeError("image task returned streaming result unexpectedly") data = _materialize_inline_task_images( result.get("data"), base_url=_clean(payload.get("base_url")), ) account_email = _clean(result.get("_account_email") or result.get("account_email")) if not isinstance(data, list) or not data: upstream = _clean(result.get("message")) if upstream: message = upstream else: message = "号池中没有可用账号或所有账号均被限流，请检查号池状态（账号额度、是否被封禁、是否到达生图上限）" raise RuntimeError(message) duration_ms = int((time.time() - started) * 1000) settled_count = _count_image_items(data) self._settle_credit(key, success=True) self._settle_third_party_wallet(key, success_count=settled_count) self._settle_internal_model_daily_limit(key, success_count=settled_count) self._record_usage_ledger(key, success=True, success_count=settled_count) self._log_call( identity, mode, model, started, "调用完成", request_preview=request_text(payload.get("prompt")), urls=_collect_image_urls(data), account_email=account_email, image_api=_collect_third_party_image_api(result), billing=self._billing_log_fields(key), ) # 出图成功同样要留下上游会话 id：官网来源取证要靠它生成分享链接， # 而且 task 记录是前端拿到它的唯一通道（没有单任务详情接口）。 conversation_id = _clean(result.get("_conversation_id")) # Publish terminal success only after reservations and its audit log are settled. self._update_task( key, status=TASK_STATUS_SUCCESS, data=data, error="", duration_ms=duration_ms, **({"conversation_id": conversation_id} if conversation_id else {}), ) except Exception as exc: error_message = str(exc) or "image task failed" account_email = _clean(getattr(exc, "account_email", account_email)) conversation_id = _clean(getattr(exc, "conversation_id", "")) # 超时那一刻上游 assistant 的原话：用来告诉用户「官网到底回没回、回了什么」。 # 上游文案可能很长，截断保存，避免把 task 记录撑爆。 upstream_note = _clean(getattr(exc, "upstream_text", ""))[:2000] duration_ms = int((time.time() - started) * 1000) # 上游已给终局结论（只回文字、本轮没生图）时，继续等待没有任何意义 —— # 前台那个「继续等待」按钮必须收掉，否则用户点它等于让后端白轮询一轮。 terminal_no_image = bool(getattr(exc, "terminal_no_image", False)) # 「查看官网对话」入口只对「官网可能有图/有话可核对」的失败有意义。账号生图配额 # 耗尽（「你已达到 Free 套餐的图像生成请求上限」）时官网这一轮什么都没生成， # 点进去只有同一句提示 → 收掉入口（号池层已标记，这里再按原文兜一次）。 can_view_official = bool(getattr(exc, "can_view_official", True)) if can_view_official and is_image_quota_exhausted_message(error_message): can_view_official = False resumable_timeout = bool(conversation_id) and not terminal_no_image and any( marker in error_message.lower() for marker in ("timeout", "timed out", "超时", "轮询") ) if not resumable_timeout: self._settle_credit(key, success=False) self._settle_third_party_wallet(key, success_count=0) self._settle_internal_model_daily_limit(key, success_count=0) self._record_usage_ledger(key, success=False) self._log_call( identity, mode, model, started, "调用失败", request_preview=request_text(payload.get("prompt")), status="failed", error=error_message, account_email=account_email, image_api=getattr(exc, "image_api", None), billing=self._billing_log_fields(key), ) # Publish terminal failure only after reservations and its audit log are settled. self._update_task( key, status=TASK_STATUS_ERROR, error=error_message, data=[], duration_ms=duration_ms, **({"conversation_id": conversation_id} if conversation_id else {}), **({"upstream_note": upstream_note} if upstream_note else {}), **({"account_email": account_email} if account_email else {}), **({"can_continue_waiting": False} if terminal_no_image else {}), **({"can_view_official": False} if not can_view_official else {}), ) def _log_call( self, identity: dict[str, object], mode: str, model: str, started: float, suffix: str, *, request_preview: str = "", status: str = "success", error: str = "", urls: list[str] | None = None, account_email: str = "", image_api: dict[str, Any] | None = None, billing: dict[str, Any] | None = None, ) -> None: endpoint = "/v1/images/edits" if mode == "edit" else "/v1/images/generations" summary_prefix = "图生图" if mode == "edit" else "文生图" detail = { "key_id": identity.get("id"), "key_name": identity.get("name"), "role": identity.get("role"), "endpoint": endpoint, "model": model, "started_at": china_from_timestamp_string(started), "ended_at": _now_iso(), "duration_ms": int((time.time() - started) * 1000), "status": status, } if request_preview: detail["request_text"] = request_preview if error: detail["error"] = error if account_email: detail["account_email"] = account_email if image_api: detail["image_api_name"] = " / ".join( item for item in [ _clean(image_api.get("group_name") or image_api.get("group_id")), _clean(image_api.get("connection_name") or image_api.get("connection_id")), ] if item ) detail["image_api_group"] = image_api.get("group_name") or image_api.get("group_id") detail["image_api_connection"] = image_api.get("connection_name") or image_api.get("connection_id") detail["image_api_model"] = image_api.get("model_id") detail["image_api_upstream_model"] = image_api.get("upstream_model") if image_api.get("base_url"): detail["image_api_endpoint"] = image_api.get("base_url") detail["image_api_resolution"] = image_api.get("resolution") if image_api.get("response_format"): detail["image_api_response_format"] = image_api.get("response_format") if image_api.get("size"): detail["image_api_size"] = image_api.get("size") if urls: detail["urls"] = list(dict.fromkeys(urls)) if billing: detail.update(billing) try: log_service.add(LOG_TYPE_CALL, f"{summary_prefix}{suffix}", detail) except Exception: pass def _billing_log_fields(self, key: str) -> dict[str, Any]: with self._lock: task = self._tasks.get(key) if not isinstance(task, dict): return {} pricing = task.get("credit_pricing") if isinstance(task.get("credit_pricing"), dict) else {} if not pricing: return {} estimated = _credit_amount(pricing.get("estimated_total", task.get("credit_reserved"))) actual = _credit_amount(task.get("credit_actual", 0)) refund = _credit_amount(task.get("credit_refund", max(0.0, float(estimated) - float(actual)))) return { "billing_base_unit_price": pricing.get("base_unit_price", 1), "billing_effective_unit_price": pricing.get("unit_price", 1), "billing_pay_percentage": pricing.get("pay_percentage", 100), "billing_rule_id": pricing.get("rule_id", ""), "billing_rule_name": pricing.get("rule_name", ""), "billing_estimated": estimated, "billing_actual": actual, "billing_refund": refund, "billing_priced_at": pricing.get("priced_at") or pricing.get("evaluated_at", ""), } def _update_task(self, key: str, **updates: Any) -> None: with self._lock: task = self._tasks.get(key) if task is None: return task.update(updates) task["updated_at"] = _now_iso() task["updated_ts"] = time.time() self._save_locked([task]) def _record_progress(self, key: str, step: str, *, set_started_ts: bool = False) -> None: cleaned_step = _clean(step) or "generating" now_ts = time.time() with self._lock: task = self._tasks.get(key) if task is None: return events = task.get("progress_events") if isinstance(task.get("progress_events"), list) else [] normalized_events = [ event for event in events if isinstance(event, dict) and _clean(event.get("step")) ] if not normalized_events or _clean(normalized_events[-1].get("step")) != cleaned_step: normalized_events.append({ "step": cleaned_step, "at": _now_iso(), "ts": now_ts, }) task["progress_events"] = normalized_events[-20:] task["progress"] = cleaned_step if set_started_ts: task["started_ts"] = now_ts task["updated_at"] = _now_iso() task["updated_ts"] = now_ts # Progress is transient while a task is running. def _set_progress_profile(self, key: str, profile: str, *, reset_events: bool = False) -> None: normalized = profile if profile in {"pool", "api"} else "pool" now_ts = time.time() with self._lock: task = self._tasks.get(key) if task is None: return task["progress_profile"] = normalized if reset_events: task["progress_events"] = [] task.pop("progress", None) task["updated_at"] = _now_iso() task["updated_ts"] = now_ts # Profile changes are transient while a task is running. @staticmethod def _initial_progress_profile(payload: dict[str, Any]) -> str: resolution = _clean(payload.get("resolution") or payload.get("image_resolution")).lower() group_id = _clean(payload.get("third_party_image_group_id")) if group_id or resolution in {"1k", "2k", "4k"} or bool(payload.get("_third_party_image")): return "api" return "pool" def _settle_credit(self, key: str, *, success: bool) -> None: with self._lock: task = self._tasks.get(key) if task is None or bool(task.get("credit_settled")): return user_id = _clean(task.get("credit_user_id")) amount = _credit_amount(task.get("credit_reserved")) daily_amount = _credit_amount(task.get("credit_daily_reserved")) paid_amount = _credit_amount(task.get("credit_paid_reserved")) package_reservations = task.get("credit_package_reservations") if isinstance(task.get("credit_package_reservations"), list) else [] grant_reservations = task.get("credit_grant_reservations") if isinstance(task.get("credit_grant_reservations"), list) else [] package_amount = _credit_amount(sum(float(_credit_amount(item.get("amount"))) for item in package_reservations if isinstance(item, dict))) grant_amount = _credit_amount(sum(float(_credit_amount(item.get("amount"))) for item in grant_reservations if isinstance(item, dict))) if amount > 0 and round(float(daily_amount) + float(paid_amount) + float(package_amount) + float(grant_amount), 2) == 0: paid_amount = amount if not user_id or amount <= 0: task["credit_settled"] = True task["credit_actual"] = 0 task["credit_refund"] = 0 self._save_locked([task]) return try: from services.auth_service import auth_service if success: auth_service.consume_reserved_credits( user_id, amount, daily_amount=daily_amount, paid_amount=paid_amount, package_reservations=package_reservations, grant_reservations=grant_reservations, ) else: auth_service.refund_reserved_credits( user_id, amount, daily_amount=daily_amount, paid_amount=paid_amount, package_reservations=package_reservations, grant_reservations=grant_reservations, ) finally: task["credit_settled"] = True task["credit_actual"] = amount if success else 0 task["credit_refund"] = 0 if success else amount task["updated_at"] = _now_iso() self._save_locked([task]) def _settle_third_party_wallet(self, key: str, *, success_count: int) -> None: with self._lock: task = self._tasks.get(key) if task is None or bool(task.get("third_party_wallet_settled")): return user_id = _clean(task.get("third_party_wallet_user_id")) group_id = _clean(task.get("third_party_wallet_group_id")) model_id = _clean(task.get("third_party_wallet_model_id")) amount = _credit_amount(task.get("third_party_wallet_reserved")) cost_per_image = _credit_amount(task.get("third_party_wallet_cost_per_image")) if not user_id or not group_id or not model_id or amount <= 0: task["third_party_wallet_settled"] = True self._save_locked([task]) return consume_amount = _credit_amount(min(float(amount), max(0, int(success_count or 0)) * float(cost_per_image))) refund_amount = _credit_amount(float(amount) - float(consume_amount)) try: from services.auth_service import auth_service if consume_amount > 0: auth_service.consume_third_party_image_wallet( user_id, group_id=group_id, model_id=model_id, amount=consume_amount, ) if refund_amount > 0: auth_service.refund_third_party_image_wallet( user_id, group_id=group_id, model_id=model_id, amount=refund_amount, ) finally: task["third_party_wallet_settled"] = True task["credit_actual"] = consume_amount task["credit_refund"] = refund_amount task["updated_at"] = _now_iso() self._save_locked([task]) def _settle_internal_model_daily_limit(self, key: str, *, success_count: int) -> None: with self._lock: task = self._tasks.get(key) if task is None or bool(task.get("internal_model_limit_settled")): return user_id = _clean(task.get("internal_model_limit_user_id")) model_id = _clean(task.get("internal_model_limit_model_id")) amount = max(0, int(task.get("internal_model_limit_reserved") or 0)) scope = _clean(task.get("internal_model_limit_scope")) if not user_id or not model_id or amount <= 0: task["internal_model_limit_settled"] = True self._save_locked([task]) return consume_amount = min(amount, max(0, int(success_count or 0))) refund_amount = amount - consume_amount try: from services.auth_service import auth_service if consume_amount > 0: consume_kwargs: dict[str, object] = { "model_id": model_id, "amount": consume_amount, } if scope: consume_kwargs["scope"] = scope auth_service.consume_internal_image_model_daily_limit( user_id, **consume_kwargs, ) if refund_amount > 0: refund_kwargs: dict[str, object] = { "model_id": model_id, "amount": refund_amount, } if scope: refund_kwargs["scope"] = scope auth_service.refund_internal_image_model_daily_limit( user_id, **refund_kwargs, ) finally: task["internal_model_limit_settled"] = True task["updated_at"] = _now_iso() self._save_locked([task]) def _record_usage_ledger(self, key: str, *, success: bool, success_count: int = 0) -> None: """把这次任务的终局落成一行用量流水，供数据看板统计。 必须在 `_settle_credit` 之后调用：实际扣减量 `credit_actual` 要等结算完才准确。 内部用户 / 管理员也要记 —— 他们不产生任何积分账本变动（内部走模型每日额度、 管理员免费），但调用次数是真实用量。看板的「次数」口径要覆盖三类角色， 否则筛到内部就是一片空白。 ⚠️ 模型维度一律取 `task["model"]`：普通用户的 `credit_model_id` 恒为 `__account_pool__`（号池计价用的伪模型），拿它统计会得到一个假模型名。 """ with self._lock: task = self._tasks.get(key) if task is None or bool(task.get("usage_ledger_recorded")): return task["usage_ledger_recorded"] = True snapshot = dict(task) owner_id = _clean(snapshot.get("owner_id")) if not owner_id: return try: from services.auth_service import auth_service from services.credit_ledger_service import credit_ledger_service from services.pool_pricing_service import pool_pricing_service except Exception: # pragma: no cover - 导入失败时静默跳过统计 return role = "" try: role = _clean((auth_service.get_key(owner_id) or {}).get("role")) except Exception: role = "" if not role: # 查不到身份（账号已被删除）时**不要**兜底成 normal：那会把已删用户的用量 # 混进普通用户口径，导致「有消耗的用户」虚高。记成 unknown 只在「全部」视图出现。 role = "unknown" pricing = snapshot.get("credit_pricing") if isinstance(snapshot.get("credit_pricing"), dict) else {} unit_price = _credit_amount(pricing.get("unit_price")) if unit_price <= 0: # 内部 / 管理员任务没有定价快照 → 取当下的普通用户等效单价做折算基准。 try: unit_price = _credit_amount(pool_pricing_service.quote({"role": "normal"}, 1).get("unit_price")) except Exception: unit_price = 0 try: credit_ledger_service.record( identity_id=owner_id, role=role, model=_clean(snapshot.get("model"), "unknown"), status="success" if success else "failed", image_count=int(success_count or 0) if success else 0, credits=_credit_amount(snapshot.get("credit_actual")) if success else 0, unit_price=unit_price, task_id=_clean(snapshot.get("id")), ) except Exception: # 统计写失败绝不能反过来把已经出图成功的任务标成失败。 pass def _connect_database(self) -> sqlite3.Connection | _SharedMemoryConnection: if self._memory_database is not None: return _SharedMemoryConnection(self._memory_database) connection = sqlite3.connect(self.db_path, timeout=30) try: connection.execute("PRAGMA busy_timeout=30000") connection.execute("PRAGMA journal_mode=DELETE") connection.execute("PRAGMA synchronous=FULL") return connection except Exception: connection.close() raise def _prepare_database_locked(self) -> None: try: self._initialize_database_locked() self._check_database_integrity_locked() except sqlite3.OperationalError as exc: # "unable to open database file" / "attempt to write a readonly database" # 等运行期错误是环境问题（目录不可写、磁盘满、路径异常），不是数据损坏： # 隔离没有意义，原地重试必然再次失败。改为尝试备用位置。 if self._relocate_database_locked(exc): return raise except sqlite3.DatabaseError as exc: quarantined = self._quarantine_corrupt_database_locked() print( "[image-tasks] Corrupt SQLite database isolated; rebuilding from legacy JSON: " f"error={exc}, files={quarantined}" ) self._initialize_database_locked() self._check_database_integrity_locked() def _database_fallback_candidates(self) -> list[Path]: """主数据库位置不可用时按顺序尝试的备用位置（测试可覆盖）。""" candidates: list[Path] = [] override = os.getenv("CHATGPT2API_IMAGE_TASKS_DB", "").strip() if override: candidates.append(Path(override)) base_data = Path(__file__).resolve().parents[1] / "data" / "image_tasks.db" candidates.append(base_data) candidates.append(Path(tempfile.gettempdir()) / "chatgpt2api-image-tasks" / "image_tasks.db") return candidates def _relocate_database_locked(self, exc: sqlite3.OperationalError) -> bool: previous_path = self.db_path for candidate in self._database_fallback_candidates(): if candidate == previous_path: continue self.db_path = candidate try: candidate.parent.mkdir(parents=True, exist_ok=True) self._initialize_database_locked() self._check_database_integrity_locked() except Exception: continue print( f"[image-tasks] SQLite database at {previous_path} is unavailable ({exc}); " f"image tasks will be stored at {candidate}." ) return True # 所有落盘位置都不可用（典型：磁盘满）。退化为共享内存库， # 保证服务还能起来、管理员还能登录清理，代价是任务不跨重启保留。 self.db_path = previous_path try: memory_connection = sqlite3.connect(":memory:", timeout=30, check_same_thread=False) memory_connection.execute("PRAGMA busy_timeout=30000") self._memory_database = memory_connection self._initialize_database_locked() self._check_database_integrity_locked() except Exception: self._memory_database = None raise exc print( f"[image-tasks] No writable location for the SQLite database ({exc}); " "falling back to a temporary in-memory database. " "Image tasks will NOT persist across restarts - free disk space and restart." ) return True def _check_database_integrity_locked(self) -> None: with closing(self._connect_database()) as connection: row = connection.execute("PRAGMA quick_check").fetchone() result = str(row[0] if row else "").strip().lower() if result != "ok": raise sqlite3.DatabaseError(f"image task database quick_check failed: {result or 'empty result'}") def _quarantine_corrupt_database_locked(self) -> list[str]: timestamp = datetime.now().strftime("%Y%m%d-%H%M%S") quarantined: list[str] = [] for candidate in ( self.db_path, Path(f"{self.db_path}-wal"), Path(f"{self.db_path}-shm"), ): if not candidate.exists(): continue target = candidate.with_name(f"{candidate.name}.corrupt-{timestamp}") candidate.replace(target) quarantined.append(target.name) return quarantined def _initialize_database_locked(self) -> None: with closing(self._connect_database()) as connection, connection: connection.execute( """ CREATE TABLE IF NOT EXISTS image_tasks ( task_key TEXT PRIMARY KEY, owner_id TEXT NOT NULL, task_id TEXT NOT NULL, updated_at TEXT NOT NULL DEFAULT '', payload TEXT NOT NULL ) """ ) connection.execute( "CREATE INDEX IF NOT EXISTS idx_image_tasks_owner_updated ON image_tasks(owner_id, updated_at DESC)" ) connection.execute( "CREATE TABLE IF NOT EXISTS image_task_metadata (key TEXT PRIMARY KEY, value TEXT NOT NULL)" ) def _read_database_items_locked(self) -> tuple[list[dict[str, Any]], bool]: with closing(self._connect_database()) as connection, connection: rows = connection.execute("SELECT payload FROM image_tasks").fetchall() marker = connection.execute( "SELECT value FROM image_task_metadata WHERE key = 'legacy_json_imported'" ).fetchone() items: list[dict[str, Any]] = [] for (payload,) in rows: try: item = json.loads(payload) except Exception: continue if isinstance(item, dict): items.append(item) return items, marker is not None def _mark_legacy_imported_locked(self) -> None: with closing(self._connect_database()) as connection, connection: connection.execute( "INSERT OR REPLACE INTO image_task_metadata(key, value) VALUES('legacy_json_imported', ?)", (_now_iso(),), ) def _load_locked(self) -> dict[str, dict[str, Any]]: raw_items, legacy_imported = self._read_database_items_locked() should_import_legacy = not raw_items and not legacy_imported if should_import_legacy and self.path.exists(): try: raw = json.loads(self.path.read_text(encoding="utf-8")) raw_items = raw.get("tasks") if isinstance(raw, dict) else raw except Exception: raw_items = [] if not isinstance(raw_items, list): raw_items = [] if not isinstance(raw_items, list): return {} tasks: dict[str, dict[str, Any]] = {} for item in raw_items: if not isinstance(item, dict): continue task_id = _clean(item.get("id")) owner = _clean(item.get("owner_id")) if not task_id or not owner: continue status = _clean(item.get("status")) if status not in {TASK_STATUS_QUEUED, TASK_STATUS_RUNNING, TASK_STATUS_SUCCESS, TASK_STATUS_ERROR}: status = TASK_STATUS_ERROR task = { "id": task_id, "owner_id": owner, "status": status, "mode": "edit" if item.get("mode") == "edit" else "generate", "model": _clean(item.get("model"), "gpt-image-2"), "size": _clean(item.get("size")), "progress_profile": _clean(item.get("progress_profile")) if _clean(item.get("progress_profile")) in {"pool", "api"} else "pool", "created_at": _clean(item.get("created_at"), _now_iso()), "started_at": _clean(item.get("started_at")), "updated_at": _clean(item.get("updated_at"), _clean(item.get("created_at"), _now_iso())), "created_ts": item.get("created_ts"), "updated_ts": item.get("updated_ts"), "started_ts": item.get("started_ts"), "duration_ms": item.get("duration_ms"), "base_url": _clean(item.get("base_url")), } conversation_id = _clean(item.get("conversation_id")) if conversation_id: task["conversation_id"] = conversation_id account_email = _clean(item.get("account_email")) if account_email: task["account_email"] = account_email progress = _clean(item.get("progress")) if progress: task["progress"] = progress if isinstance(item.get("progress_events"), list): task["progress_events"] = [ event for event in item.get("progress_events", [])[-20:] if isinstance(event, dict) and _clean(event.get("step")) ] credit_user_id = _clean(item.get("credit_user_id")) credit_reserved = _credit_amount(item.get("credit_reserved")) if credit_user_id and credit_reserved > 0: task["credit_user_id"] = credit_user_id task["credit_reserved"] = credit_reserved task["credit_daily_reserved"] = _credit_amount(item.get("credit_daily_reserved")) task["credit_paid_reserved"] = _credit_amount(item.get("credit_paid_reserved")) task["credit_package_reservations"] = item.get("credit_package_reservations") if isinstance(item.get("credit_package_reservations"), list) else [] task["credit_grant_reservations"] = item.get("credit_grant_reservations") if isinstance(item.get("credit_grant_reservations"), list) else [] task["credit_settled"] = bool(item.get("credit_settled")) if isinstance(item.get("credit_pricing"), dict): task["credit_pricing"] = dict(item.get("credit_pricing")) if "credit_actual" in item: task["credit_actual"] = _credit_amount(item.get("credit_actual")) if "credit_refund" in item: task["credit_refund"] = _credit_amount(item.get("credit_refund")) wallet_user_id = _clean(item.get("third_party_wallet_user_id")) wallet_reserved = _credit_amount(item.get("third_party_wallet_reserved")) if wallet_user_id and wallet_reserved > 0: task["third_party_wallet_user_id"] = wallet_user_id task["third_party_wallet_group_id"] = _clean(item.get("third_party_wallet_group_id")) task["third_party_wallet_model_id"] = _clean(item.get("third_party_wallet_model_id")) task["third_party_wallet_reserved"] = wallet_reserved task["third_party_wallet_cost_per_image"] = _credit_amount(item.get("third_party_wallet_cost_per_image")) task["third_party_wallet_settled"] = bool(item.get("third_party_wallet_settled")) internal_limit_user_id = _clean(item.get("internal_model_limit_user_id")) internal_limit_reserved = max(0, int(item.get("internal_model_limit_reserved") or 0)) if internal_limit_user_id and internal_limit_reserved > 0: task["internal_model_limit_user_id"] = internal_limit_user_id task["internal_model_limit_model_id"] = _clean(item.get("internal_model_limit_model_id")) task["internal_model_limit_reserved"] = internal_limit_reserved task["internal_model_limit_scope"] = _clean(item.get("internal_model_limit_scope")) task["internal_model_limit_settled"] = bool(item.get("internal_model_limit_settled")) data = item.get("data") if isinstance(data, list): task["data"] = data error = _clean(item.get("error")) if error: task["error"] = error tasks[_task_key(owner, task_id)] = task if should_import_legacy: self._write_database_tasks_locked(tasks.values(), replace_all=True) self._mark_legacy_imported_locked() return tasks def _write_database_tasks_locked( self, tasks: Any, *, replace_all: bool = False, ) -> None: task_list = [task for task in tasks if isinstance(task, dict)] try: self._write_database_task_list_locked(task_list, replace_all=replace_all) except sqlite3.OperationalError as exc: message = str(exc).lower() recoverable = any( marker in message for marker in ( "attempt to write a readonly database", "readonly database", "unable to open database file", "disk i/o error", "database or disk is full", ) ) print(json.dumps(_database_write_diagnostics(self.db_path, exc), ensure_ascii=False, default=str)) if not recoverable or not self._relocate_database_locked(exc): raise # The in-memory task map is authoritative while the process is # running. Repopulate the replacement database so a transient HF # persistent-volume failure does not discard previously submitted # tasks, then retry the write that triggered recovery. replacement_tasks = list(self._tasks.values()) if replacement_tasks: self._write_database_task_list_locked(replacement_tasks, replace_all=True) elif task_list: self._write_database_task_list_locked(task_list, replace_all=replace_all) def _write_database_task_list_locked( self, task_list: list[dict[str, Any]], *, replace_all: bool = False, ) -> None: with closing(self._connect_database()) as connection, connection: if replace_all: current_keys = {_task_key(_clean(task.get("owner_id")), _clean(task.get("id"))) for task in task_list} if current_keys: placeholders = ",".join("?" for _ in current_keys) connection.execute( f"DELETE FROM image_tasks WHERE task_key NOT IN ({placeholders})", tuple(current_keys), ) else: connection.execute("DELETE FROM image_tasks") connection.executemany( """ INSERT INTO image_tasks(task_key, owner_id, task_id, updated_at, payload) VALUES(?, ?, ?, ?, ?) ON CONFLICT(task_key) DO UPDATE SET owner_id = excluded.owner_id, task_id = excluded.task_id, updated_at = excluded.updated_at, payload = excluded.payload """, [ ( _task_key(_clean(task.get("owner_id")), _clean(task.get("id"))), _clean(task.get("owner_id")), _clean(task.get("id")), _clean(task.get("updated_at")), json.dumps(task, ensure_ascii=False, separators=(",", ":")), ) for task in task_list ], ) def _save_locked(self, changed_tasks: Any = None) -> None: if changed_tasks is None: self._write_database_tasks_locked(self._tasks.values(), replace_all=True) return self._write_database_tasks_locked(changed_tasks) def _recover_unfinished_locked(self) -> bool: changed = False for task in self._tasks.values(): if task.get("status") in UNFINISHED_STATUSES: self._refund_task_credit_locked(task) self._refund_task_third_party_wallet_locked(task) self._refund_task_internal_model_daily_limit_locked(task) task["status"] = TASK_STATUS_ERROR task["error"] = "服务已重启，未完成的图片任务已中断" task["updated_at"] = _now_iso() changed = True return changed def _refund_task_credit_locked(self, task: dict[str, Any]) -> None: if bool(task.get("credit_settled")): return user_id = _clean(task.get("credit_user_id")) amount = _credit_amount(task.get("credit_reserved")) daily_amount = _credit_amount(task.get("credit_daily_reserved")) paid_amount = _credit_amount(task.get("credit_paid_reserved")) package_reservations = task.get("credit_package_reservations") if isinstance(task.get("credit_package_reservations"), list) else [] grant_reservations = task.get("credit_grant_reservations") if isinstance(task.get("credit_grant_reservations"), list) else [] package_amount = _credit_amount(sum(float(_credit_amount(item.get("amount"))) for item in package_reservations if isinstance(item, dict))) grant_amount = _credit_amount(sum(float(_credit_amount(item.get("amount"))) for item in grant_reservations if isinstance(item, dict))) if amount > 0 and round(float(daily_amount) + float(paid_amount) + float(package_amount) + float(grant_amount), 2) == 0: paid_amount = amount if not user_id or amount <= 0: task["credit_settled"] = True return try: from services.auth_service import auth_service auth_service.refund_reserved_credits( user_id, amount, daily_amount=daily_amount, paid_amount=paid_amount, package_reservations=package_reservations, grant_reservations=grant_reservations, ) finally: task["credit_settled"] = True task["credit_actual"] = 0 task["credit_refund"] = amount def _refund_task_third_party_wallet_locked(self, task: dict[str, Any]) -> None: if bool(task.get("third_party_wallet_settled")): return user_id = _clean(task.get("third_party_wallet_user_id")) group_id = _clean(task.get("third_party_wallet_group_id")) model_id = _clean(task.get("third_party_wallet_model_id")) amount = _credit_amount(task.get("third_party_wallet_reserved")) if not user_id or not group_id or not model_id or amount <= 0: task["third_party_wallet_settled"] = True return try: from services.auth_service import auth_service auth_service.refund_third_party_image_wallet( user_id, group_id=group_id, model_id=model_id, amount=amount, ) finally: task["third_party_wallet_settled"] = True def _refund_task_internal_model_daily_limit_locked(self, task: dict[str, Any]) -> None: if bool(task.get("internal_model_limit_settled")): return user_id = _clean(task.get("internal_model_limit_user_id")) model_id = _clean(task.get("internal_model_limit_model_id")) amount = max(0, int(task.get("internal_model_limit_reserved") or 0)) scope = _clean(task.get("internal_model_limit_scope")) if not user_id or not model_id or amount <= 0: task["internal_model_limit_settled"] = True return try: from services.auth_service import auth_service refund_kwargs: dict[str, object] = { "model_id": model_id, "amount": amount, } if scope: refund_kwargs["scope"] = scope auth_service.refund_internal_image_model_daily_limit( user_id, **refund_kwargs, ) finally: task["internal_model_limit_settled"] = True def _cleanup_locked(self) -> bool: try: retention_days = max(1, int(self.retention_days_getter())) except Exception: retention_days = 30 cutoff = time.time() - retention_days * 86400 removed_keys: list[str] = [] for key, task in self._tasks.items(): if task.get("status") not in TERMINAL_STATUSES or _timestamp(task.get("updated_at")) >= cutoff: continue try: if not bool(task.get("credit_settled")): self._settle_credit_locked(task, task.get("status") == TASK_STATUS_COMPLETED) if not bool(task.get("internal_model_limit_settled")): self._settle_internal_model_limit_locked(task, task.get("status") == TASK_STATUS_COMPLETED) except Exception: # Keep it so a later cleanup pass can retry the refund/settlement. continue removed_keys.append(key) for key in removed_keys: self._tasks.pop(key, None) return bool(removed_keys) def _cleanup_if_due_locked(self) -> bool: now = time.monotonic() if now - self._last_cleanup_at < TASK_CLEANUP_INTERVAL_SECONDS: return False self._last_cleanup_at = now return self._cleanup_locked() def resume_poll( self, identity: dict[str, object], task_id: str, extra_timeout_secs: float = 30.0, ) -> dict[str, Any]: """Continue polling the upstream conversation created by a timed-out task.""" owner = _owner_id(identity) key = _task_key(owner, _clean(task_id)) with self._lock: task = self._tasks.get(key) if task is None: raise ValueError("image task not found") if task.get("status") != TASK_STATUS_ERROR: raise ValueError("task is not in error state") error_message = _clean(task.get("error")).lower() if not any(marker in error_message for marker in ("timeout", "timed out", "超时", "轮询")): raise ValueError("only timed-out tasks can continue waiting") conversation_id = _clean(task.get("conversation_id")) if not conversation_id: raise ValueError("task has no upstream conversation id") mode = _clean(task.get("mode"), "generate") model = _clean(task.get("model"), "gpt-image-2") task["status"] = TASK_STATUS_RUNNING task["error"] = "" task["started_at"] = _now_iso() task["started_ts"] = time.time() task["updated_at"] = _now_iso() task["updated_ts"] = time.time() self._save_locked([task]) threading.Thread( target=self._run_resume_poll, args=(key, conversation_id, max(5.0, min(float(extra_timeout_secs), 120.0)), dict(identity), mode, model), name=f"image-resume-{_clean(task_id)[:16]}", daemon=True, ).start() return _public_task(task) def _run_resume_poll( self, key: str, conversation_id: str, extra_timeout_secs: float, identity: dict[str, object], mode: str, model: str, ) -> None: started = time.time() backend = None account_email = "" try: from services.account_service import account_service from services.openai_backend_api import OpenAIBackendAPI from services.protocol.conversation import format_image_result with self._lock: task = dict(self._tasks.get(key) or {}) account_email = _clean(task.get("account_email")) candidates = account_service.list_accounts_for_pool_identity(identity) account = next( (item for item in candidates if account_email and _clean(item.get("email")).lower() == account_email.lower()), None, ) access_token = _clean((account or {}).get("access_token")) if not access_token: raise RuntimeError("the account used by this task is no longer available") backend = OpenAIBackendAPI(access_token=access_token) file_ids, sediment_ids, _message_id = backend._poll_image_results( conversation_id, extra_timeout_secs, ) urls = backend.resolve_conversation_image_urls( conversation_id, file_ids, sediment_ids, poll=False, ) if not urls: raise RuntimeError("image URL resolution failed") image_items = [ {"b64_json": __import__("base64").b64encode(image_data).decode("ascii")} for image_data in backend.download_image_bytes(urls) ] data = format_image_result( image_items, "", "url", _clean(task.get("base_url")), int(time.time()), )["data"] if not data: raise RuntimeError("upstream returned no image data") duration_ms = int((time.time() - started) * 1000) success_count = _count_image_items(data) self._settle_credit(key, success=True) self._settle_third_party_wallet(key, success_count=success_count) self._settle_internal_model_daily_limit(key, success_count=success_count) self._record_usage_ledger(key, success=True, success_count=success_count) self._log_call( identity, mode, model, started, "continued polling completed", urls=_collect_image_urls(data), account_email=account_email, billing=self._billing_log_fields(key), ) self._update_task(key, status=TASK_STATUS_SUCCESS, data=data, error="", duration_ms=duration_ms) except Exception as exc: from services.openai_backend_api import ImageTextOnlyAnswer # 上游这一轮其实早就用纯文字收尾了（例：内容策略拒绝），继续等也等不到图片 —— # 和首次生成路径同一套语义：原话当失败原因，并收掉「继续等待」按钮。 terminal_no_image = isinstance(exc, ImageTextOnlyAnswer) error_message = str(exc) or "resume poll failed" # 和首次生成路径同口径：账号生图配额耗尽时官网这轮没生成任何东西， # 不该在卡片上留「查看官网对话」入口。 can_view_official = not is_image_quota_exhausted_message(error_message) resumable_timeout = not terminal_no_image and any( marker in error_message.lower() for marker in ("timeout", "timed out", "超时", "轮询") ) if not resumable_timeout: self._settle_credit(key, success=False) self._settle_third_party_wallet(key, success_count=0) self._settle_internal_model_daily_limit(key, success_count=0) self._record_usage_ledger(key, success=False) self._log_call( identity, mode, model, started, "continued polling failed", status="failed", error=error_message, account_email=account_email, billing=self._billing_log_fields(key), ) self._update_task( key, status=TASK_STATUS_ERROR, error=error_message, data=[], duration_ms=int((time.time() - started) * 1000), **({"can_continue_waiting": False} if terminal_no_image else {}), **({"can_view_official": False} if not can_view_official else {}), ) finally: if backend is not None: backend.close() _image_task_json_path = sqlite_mirror_service.runtime_file(DATA_DIR, "image_tasks.json") image_task_service = ImageTaskService(_image_task_json_path) # ImageTaskService creates/migrates its SQLite database beside the local JSON # path. Register it only afterwards: probing the legacy image_tasks.db directly # on a Hugging Face Bucket can make PRAGMA integrity_check perform minutes of # random network I/O and prevent Uvicorn from ever reaching its startup hook. sqlite_mirror_service.runtime_file(DATA_DIR, "image_tasks.db")