Blame for mcp_tools/run.py

"""
MCP Tool: run
Unified image generation task submission and execution interface.
"""

import time
import uuid
from core.task_scheduler import QueueFullError, submit_background
from core.runtime_config import CONFIG
from .common import (
    _load_yaml,
    _MODEL_LIST_PATH,
    _CHAIN_FEATURES_PATH,
    _TASK_DEFINITIONS,
    _TASKS_DB,
    _TASKS_LOCK,
    _get_task_snapshot,
    _execute_imagegen_pipeline,
    _get_model_feature_context,
    _canonicalize_chain_type,
    _resolve_feature_name,
)
from .error_schema import make_error, make_validation_error, make_not_found_error
from .image_requests import normalize_image_request, validate_image_plan
from core.reference_inputs import ImageInputError


def handle_run(params: dict) -> dict:
    """Unified image generation task execution interface."""
    if not isinstance(params, dict):
        return make_validation_error("Request params must be an object.")
    params = params.copy()
    if isinstance(params.get("chain"), list):
        params["chain"] = [
            item.copy() if isinstance(item, dict) else item
            for item in params["chain"]
        ]

    missing = []
    for req_field in ["task_type", "model", "prompt"]:
        if req_field not in params or not params[req_field]:
            missing.append(req_field)
    if missing:
        return make_validation_error(
            f"Missing required parameter(s): {', '.join(missing)}",
            missing_fields=missing,
        )

    task_type = params["task_type"]
    if (
        not isinstance(task_type, str)
        or not isinstance(params["model"], str)
        or not isinstance(params["prompt"], str)
        or not task_type.strip()
        or not params["model"].strip()
        or not params["prompt"].strip()
    ):
        return make_validation_error(
            "task_type, model, and prompt must be strings.",
            invalid_fields={
                key: "Expected a non-empty string"
                for key in ("task_type", "model", "prompt")
                if not isinstance(params.get(key), str) or not params.get(key, "").strip()
            },
        )
    task_type = task_type.strip()
    params["task_type"] = task_type
    params["model"] = params["model"].strip()
    valid_tasks = [t["task_type"] for t in _TASK_DEFINITIONS]
    if task_type not in valid_tasks:
        return make_validation_error(
            f"Invalid task_type '{task_type}'. Must be one of {valid_tasks}.",
            invalid_fields={"task_type": f"Must be in {valid_tasks}"},
        )

    model_list = _load_yaml(_MODEL_LIST_PATH)
    checkpoints = model_list.get("Checkpoint", {})
    all_models = set()
    for arch_name, arch_data in checkpoints.items():
        if isinstance(arch_data, dict):
            for m in arch_data.get("models", []):
                all_models.add(m.get("display_name"))

    if params["model"] not in all_models:
        return make_not_found_error("model", params["model"])

    try:
        normalize_image_request(params)
    except ImageInputError as exc:
        return make_validation_error(str(exc), invalid_fields={"images": str(exc)})

    task_requirements = {
        "txt2img": ("width", "height"),
        "reference": ("width", "height", "images"),
        "img2img": ("image",),
        "inpaint": ("image",),
        "outpaint": ("image", "pad_left", "pad_right", "pad_top", "pad_bottom"),
        "hires_fix": ("image", "upscale_by"),
    }
    missing_task_fields = [
        field
        for field in task_requirements.get(task_type, ())
        if not (field == "image" and params.get("images"))
        if field not in params or params[field] is None or params[field] == ""
    ]
    if missing_task_fields:
        return make_validation_error(
            f"Missing required parameter(s) for {task_type}: {', '.join(missing_task_fields)}",
            missing_fields=missing_task_fields,
        )

    invalid_fields = {}

    chain_features = _load_yaml(_CHAIN_FEATURES_PATH)

    def _number(name, *, integer=False, minimum=None, maximum=None):
        if name not in params:
            return
        raw_value = params[name]
        if raw_value is None or (isinstance(raw_value, str) and not raw_value.strip()):
            if name in task_requirements.get(task_type, ()):
                invalid_fields[name] = "Expected a non-empty number"
            else:
                # 可选数值的空值等同于未提供，统一沿用执行层的任务 / 模型默认值。
                params.pop(name)
            return
        if isinstance(raw_value, bool):
            invalid_fields[name] = "Expected a number" if not integer else "Expected an integer"
            return
        try:
            numeric_value = float(raw_value)
        except (TypeError, ValueError):
            invalid_fields[name] = "Expected a number" if not integer else "Expected an integer"
            return
        if numeric_value != numeric_value or numeric_value in (float("inf"), float("-inf")):
            invalid_fields[name] = "Must be finite"
            return
        if integer and not numeric_value.is_integer():
            invalid_fields[name] = "Expected an integer"
            return
        value = int(numeric_value) if integer else numeric_value
        if minimum is not None and value < minimum:
            invalid_fields[name] = f"Must be >= {minimum}"
        elif maximum is not None and value > maximum:
            invalid_fields[name] = f"Must be <= {maximum}"
        else:
            params[name] = value

    for field in ("width", "height"):
        if task_type == "txt2img" or field in params:
            _number(field, integer=True, minimum=1, maximum=8192)
    for field in ("pad_left", "pad_right", "pad_top", "pad_bottom"):
        _number(field, integer=True, minimum=0, maximum=4096)
    _number("upscale_by", minimum=0.01, maximum=4.0)
    _number("denoise", minimum=0.0, maximum=1.0)
    _number("steps", integer=True, minimum=1, maximum=100)
    _number("cfg", minimum=0.0, maximum=100.0)
    _number("batch_size", integer=True, minimum=1, maximum=CONFIG.max_batch_size)
    _number("seed", integer=True, minimum=-1)

    for field in ("sampler", "scheduler", "negative_prompt"):
        if field in params and params[field] is not None and not isinstance(params[field], str):
            invalid_fields[field] = "Expected a string"

    if "async_execution" in params:
        async_execution = params["async_execution"]
        if isinstance(async_execution, str):
            normalized = async_execution.strip().casefold()
            if normalized in {"true", "1"}:
                params["async_execution"] = True
            elif normalized in {"false", "0"}:
                params["async_execution"] = False
            else:
                invalid_fields["async_execution"] = "Expected true/false or 1/0"
        elif not isinstance(async_execution, bool):
            invalid_fields["async_execution"] = "Expected a boolean"

    if "chain" in params:
        chain = params["chain"]
        if not isinstance(chain, list):
            invalid_fields["chain"] = "Expected a list of objects"
        else:
            context = _get_model_feature_context(params["model"])
            for index, item in enumerate(chain):
                field_name = f"chain[{index}]"
                if not isinstance(item, dict):
                    invalid_fields[field_name] = "Expected an object"
                    continue
                injector_type = item.get("injector_type")
                if not isinstance(injector_type, str) or not injector_type.strip():
                    invalid_fields[field_name] = "injector_type must be a non-empty string"
                    continue
                injector_type = injector_type.strip()
                injector_type = _canonicalize_chain_type(injector_type)
                item["injector_type"] = injector_type
                feature_name = _resolve_feature_name(injector_type)
                if feature_name is None:
                    return make_validation_error(
                        f"Unknown chain/injector type '{injector_type}'.",
                        invalid_fields={field_name: "Unknown injector_type"},
                    )
                feature_data = chain_features.get(feature_name, {})
                aliases = feature_data.get("chains", feature_name)
                aliases = aliases if isinstance(aliases, list) else [aliases]
                enabled = context and context["enabled_chains"]
                supports = context and context["supports_chain"]
                if not enabled or not any(alias in enabled for alias in [feature_name, *aliases]):
                    return make_error(
                        "FEATURE_NOT_SUPPORTED",
                        f"Model '{params['model']}' does not support chain '{injector_type}'.",
                        {"model": params["model"], "chain": injector_type},
                    )
                if supports and not any(supports(alias) for alias in [injector_type, feature_name, *aliases]):
                    return make_error(
                        "FEATURE_NOT_SUPPORTED",
                        f"Model '{params['model']}' does not support chain '{injector_type}'.",
                        {"model": params["model"], "chain": injector_type},
                    )

    if invalid_fields:
        return make_validation_error("Request contains invalid parameters.", invalid_fields=invalid_fields)

    try:
        validate_image_plan(params)
    except ImageInputError as exc:
        return make_validation_error(str(exc), invalid_fields={"image_plan": str(exc)})

    batch_size = params.get("batch_size", 1)
    if not 1 <= batch_size <= CONFIG.max_batch_size:
        return make_validation_error(
            f"batch_size must be between 1 and {CONFIG.max_batch_size}.",
            invalid_fields={"batch_size": f"Expected 1..{CONFIG.max_batch_size}"},
        )

    with _TASKS_LOCK:
        if len(_TASKS_DB) >= CONFIG.mcp_task_retention:
            finished = sorted(
                (
                    (old_task_id, task)
                    for old_task_id, task in _TASKS_DB.items()
                    if task.get("status") in {"completed", "failed"}
                ),
                key=lambda pair: pair[1].get("completed_at", pair[1].get("failed_at", 0)),
            )
            remove_count = max(1, len(_TASKS_DB) - CONFIG.mcp_task_retention + 1)
            for old_task_id, _ in finished[:remove_count]:
                _TASKS_DB.pop(old_task_id, None)
            if len(_TASKS_DB) >= CONFIG.mcp_task_retention:
                return make_error(
                    "QUEUE_FULL",
                    "任务记录已满且当前任务均未结束，请稍后再试。",
                )

        task_id = f"img_task_{uuid.uuid4().hex[:10]}"
        created_at = int(time.time())
        _TASKS_DB[task_id] = {
            "task_id": task_id,
            "status": "queued",
            "progress": 0,
            "created_at": created_at,
        }

    async_exec = params.get("async_execution", False)

    if async_exec:
        try:
            submit_background(_execute_imagegen_pipeline, task_id, params)
        except QueueFullError as exc:
            with _TASKS_LOCK:
                _TASKS_DB.pop(task_id, None)
            return make_error("QUEUE_FULL", str(exc))
        return {
            "status": "queued",
            "task_id": task_id,
            "poll_interval_ms": 2000,
            "message": "Task queued successfully. Poll get_task_status for results.",
        }
    else:
        _execute_imagegen_pipeline(task_id, params)
        return _get_task_snapshot(task_id)