| 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"): |
| |
| |
| item["error"] = error_message_service.translate( |
| sanitize_sensitive_text(task.get("error")), |
| channel="image_task", |
| ) |
| if task.get("upstream_note"): |
| |
| |
| 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: |
| |
| |
| 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), |
| ) |
| |
| |
| conversation_id = _clean(result.get("_conversation_id")) |
| |
| 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", "")) |
| |
| |
| 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)) |
| |
| |
| |
| 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), |
| ) |
| |
| 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 |
| |
|
|
| 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 |
| |
|
|
| @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: |
| return |
| role = "" |
| try: |
| role = _clean((auth_service.get_key(owner_id) or {}).get("role")) |
| except Exception: |
| role = "" |
| if not role: |
| |
| |
| 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: |
| |
| |
| |
| 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 |
|
|
| |
| |
| |
| |
| 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: |
| |
| 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) |
| |
| |
| |
| |
| sqlite_mirror_service.runtime_file(DATA_DIR, "image_tasks.db") |
|
|
