Aether/_deprecated_py_src/services/task/service.py

from __future__ import annotations

from collections.abc import Awaitable, Callable
from types import SimpleNamespace
from typing import Any

from sqlalchemy.orm import Session

from src.models.database import ApiKey
from src.services.candidate.recorder import CandidateRecorder
from src.services.task.core.context import TaskMode
from src.services.task.core.protocol import AttemptKind, AttemptResult
from src.services.task.core.schema import ExecutionResult, TaskStatusResult
from src.services.task.execute.error_handler import TaskErrorOperationsService
from src.services.task.execute.failure import TaskFailureOperationsService
from src.services.task.execute.pool import TaskPoolOperationsService
from src.services.task.execute.sync_execute import SyncTaskExecutionService
from src.services.task.request_state import RequestBodyState
from src.services.task.submit.submit_service import AsyncTaskSubmitService
from src.services.task.video.facade import TaskVideoFacadeService
from src.services.task.video.operations import VideoTaskOperationsService


async def pool_on_error(
    provider: Any,
    key: Any,
    status_code: int,
    cause: Any,
) -> None:
    """Notify the pool manager about an upstream error (health policy)."""
    try:
        from src.services.provider.pool.config import parse_pool_config
        from src.services.provider.pool.health_policy import apply_health_policy

        pool_cfg = parse_pool_config(getattr(provider, "config", None))
        if pool_cfg is None:
            return

        error_text = ""
        resp_headers: dict[str, str] = {}
        if getattr(cause, "response", None) is not None:
            try:
                error_text = (cause.response.text or "")[:4000]
            except Exception:
                pass
            try:
                resp_headers = dict(cause.response.headers)
            except Exception:
                pass
        elif isinstance(getattr(cause, "error_message", None), str):
            error_text = str(getattr(cause, "error_message", "") or "")[:4000]

        await apply_health_policy(
            provider_id=str(provider.id),
            key_id=str(key.id),
            status_code=status_code,
            error_body=error_text,
            response_headers=resp_headers,
            config=pool_cfg,
        )
    except Exception:
        pass


class TaskService:
    """
    Unified task service facade (Phase 3).

    Phase 3.1 scope:
    - Provide a single entrypoint for SYNC tasks.
    - Keep behavior consistent with the pre-Phase-3 implementation.
    - Return a structured `ExecutionResult` for downstream compatibility.
    """

    def __init__(self, db: Session, redis_client: Any | None = None) -> None:
        self.db = db
        self.redis = redis_client
        self._candidate_recorder = CandidateRecorder(db)
        # 兼容历史注入点：_execute_facade_ops/_submit_facade_ops
        # 不再依赖独立门面类，默认直接绑定 TaskService 内部实现。
        self._execute_facade_ops = SimpleNamespace(
            execute=self._execute_internal,
            _get_candidate_keys=self._candidate_recorder.get_candidate_keys,
        )

        pool_ops = TaskPoolOperationsService()
        error_ops = TaskErrorOperationsService(db, pool_ops=pool_ops)
        failure_ops = TaskFailureOperationsService()

        self._sync_ops = SyncTaskExecutionService(
            db,
            redis_client,
            recorder=self._candidate_recorder,
            pool_ops=pool_ops,
            error_ops=error_ops,
            failure_ops=failure_ops,
        )
        self._submit_ops = AsyncTaskSubmitService(
            db,
            redis_client,
            apply_pool_reorder=pool_ops.apply_pool_reorder,
            expand_pool_candidates_for_async_submit=pool_ops.expand_pool_candidates_for_async_submit,
        )
        self._video_ops = VideoTaskOperationsService(db, redis_client)
        self._submit_facade_ops = self._submit_ops
        self._video_facade_ops = TaskVideoFacadeService(self._video_ops)

    async def execute(
        self,
        *,
        task_type: str,  # chat/cli/video/image
        task_mode: TaskMode,
        api_format: str,
        model_name: str,
        user_api_key: ApiKey,
        request_func: Callable[..., Any],
        request_id: str | None = None,
        is_stream: bool = False,
        capability_requirements: dict[str, bool] | None = None,
        preferred_key_ids: list[str] | None = None,
        request_body_state: RequestBodyState | None = None,
        request_headers: dict[str, Any] | None = None,
        request_body: dict[str, Any] | None = None,
        # ASYNC-only (video submit)
        extract_external_task_id: Any | None = None,
        supported_auth_types: set[str] | None = None,
        allow_format_conversion: bool = False,
        max_candidates: int | None = None,
        create_pending_usage: bool = True,
    ) -> ExecutionResult:
        """兼容入口：默认绑定到 TaskService 内部执行路由。"""
        return await self._execute_facade_ops.execute(
            task_type=task_type,
            task_mode=task_mode,
            api_format=api_format,
            model_name=model_name,
            user_api_key=user_api_key,
            request_func=request_func,
            request_id=request_id,
            is_stream=is_stream,
            capability_requirements=capability_requirements,
            preferred_key_ids=preferred_key_ids,
            request_body_state=request_body_state,
            request_headers=request_headers,
            request_body=request_body,
            extract_external_task_id=extract_external_task_id,
            supported_auth_types=supported_auth_types,
            allow_format_conversion=allow_format_conversion,
            max_candidates=max_candidates,
            create_pending_usage=create_pending_usage,
        )

    async def _execute_internal(
        self,
        *,
        task_type: str,  # chat/cli/video/image
        task_mode: TaskMode,
        api_format: str,
        model_name: str,
        user_api_key: ApiKey,
        request_func: Callable[..., Any],
        request_id: str | None = None,
        is_stream: bool = False,
        capability_requirements: dict[str, bool] | None = None,
        preferred_key_ids: list[str] | None = None,
        request_body_state: RequestBodyState | None = None,
        request_headers: dict[str, Any] | None = None,
        request_body: dict[str, Any] | None = None,
        extract_external_task_id: Any | None = None,
        supported_auth_types: set[str] | None = None,
        allow_format_conversion: bool = False,
        max_candidates: int | None = None,
        create_pending_usage: bool = True,
    ) -> ExecutionResult:
        if task_mode == TaskMode.ASYNC:
            if extract_external_task_id is None:
                raise ValueError("extract_external_task_id is required for task_mode=ASYNC")

            outcome = await self.submit_with_failover(
                api_format=api_format,
                model_name=model_name,
                affinity_key=str(user_api_key.id),
                user_api_key=user_api_key,
                request_id=request_id,
                task_type=task_type,
                submit_func=request_func,
                extract_external_task_id=extract_external_task_id,
                supported_auth_types=supported_auth_types,
                allow_format_conversion=allow_format_conversion,
                capability_requirements=capability_requirements,
                max_candidates=max_candidates,
                request_body=request_body,
            )

            candidate_keys = []
            if request_id:
                try:
                    candidate_keys = self._execute_facade_ops._get_candidate_keys(request_id)
                except Exception:
                    candidate_keys = []

            selected_idx = -1
            if candidate_keys:
                for ck in candidate_keys:
                    if str(getattr(ck, "status", "")) == "success":
                        idx_val = getattr(ck, "candidate_index", -1)
                        selected_idx = int(idx_val) if idx_val is not None else -1
                        break

            attempt_count = 0
            if candidate_keys:
                attempt_count = sum(
                    1
                    for ck in candidate_keys
                    if str(getattr(ck, "status", ""))
                    in {"pending", "success", "failed", "cancelled"}
                )

            attempt_result = AttemptResult(
                kind=AttemptKind.ASYNC_SUBMIT,
                http_status=int(outcome.upstream_status_code or 200),
                http_headers=dict(outcome.upstream_headers or {}),
                provider_task_id=str(outcome.external_task_id),
                response_body=outcome.upstream_payload,
            )

            return ExecutionResult(
                success=True,
                attempt_result=attempt_result,
                candidate=outcome.candidate,
                candidate_index=selected_idx,
                retry_index=0,
                provider_id=str(outcome.candidate.provider.id),
                provider_name=str(outcome.candidate.provider.name),
                endpoint_id=str(outcome.candidate.endpoint.id),
                key_id=str(outcome.candidate.key.id),
                candidate_keys=candidate_keys,
                attempt_count=attempt_count,
                request_candidate_id=None,
            )

        _ = task_type  # reserved for future routing (chat/cli/video/image)

        return await self._sync_ops.execute_sync_unified(
            api_format=api_format,
            model_name=model_name,
            user_api_key=user_api_key,
            request_func=request_func,
            request_id=request_id,
            is_stream=is_stream,
            capability_requirements=capability_requirements,
            preferred_key_ids=preferred_key_ids,
            request_body_state=request_body_state,
            request_headers=request_headers,
            request_body=request_body,
            create_pending_usage=create_pending_usage,
        )

    async def execute_sync_candidates(
        self,
        *,
        api_format: str,
        model_name: str,
        candidates: list[Any],
        request_func: Callable[..., Any],
        request_id: str | None = None,
        current_user: Any | None = None,
        user_api_key: ApiKey | None = None,
        is_stream: bool = False,
        capability_requirements: dict[str, bool] | None = None,
        request_body_state: RequestBodyState | None = None,
        request_headers: dict[str, Any] | None = None,
        request_body: dict[str, Any] | None = None,
        affinity_key: str | None = None,
        create_pending_usage: bool = False,
        enable_cache_affinity: bool = False,
        is_cancelled: Callable[[], Awaitable[bool]] | None = None,
    ) -> ExecutionResult:
        """Execute a pre-built candidate set through the unified SYNC runtime."""
        from uuid import uuid4

        from src.models.database import User
        from src.services.candidate.failover import FailoverEngine
        from src.services.candidate.policy import RetryPolicy, SkipPolicy
        from src.services.candidate.resolver import CandidateResolver
        from src.services.orchestration.error_classifier import ErrorClassifier
        from src.services.orchestration.request_dispatcher import RequestDispatcher
        from src.services.provider.format import normalize_endpoint_signature
        from src.services.rate_limit.adaptive_rpm import get_adaptive_rpm_manager
        from src.services.rate_limit.concurrency_manager import get_concurrency_manager
        from src.services.request.candidate import RequestCandidateService
        from src.services.request.executor import RequestExecutor
        from src.services.scheduling.aware_scheduler import (
            CacheAwareScheduler,
            get_cache_aware_scheduler,
        )
        from src.services.system.config import SystemConfigService
        from src.services.task.execute.exception_classification import (
            CandidateErrorAction,
            classify_candidate_error_action,
        )
        from src.services.task.execute.state_transition import (
            SyncExecutionState,
            resolve_execution_error_transition,
        )
        from src.services.usage.service import UsageService

        if not request_id:
            request_id = str(uuid4())

        api_format_norm = normalize_endpoint_signature(api_format)

        priority_mode = SystemConfigService.get_config(
            self.db,
            "provider_priority_mode",
            CacheAwareScheduler.PRIORITY_MODE_PROVIDER,
        )
        scheduling_mode = SystemConfigService.get_config(
            self.db,
            "scheduling_mode",
            CacheAwareScheduler.SCHEDULING_MODE_CACHE_AFFINITY,
        )
        cache_scheduler = await get_cache_aware_scheduler(
            self.redis,
            priority_mode=priority_mode,
            scheduling_mode=scheduling_mode,
        )
        await cache_scheduler._ensure_initialized()

        concurrency_manager = await get_concurrency_manager()
        adaptive_manager = get_adaptive_rpm_manager()
        request_executor = RequestExecutor(
            db=self.db,
            concurrency_manager=concurrency_manager,
            adaptive_manager=adaptive_manager,
        )
        candidate_resolver = CandidateResolver(
            db=self.db,
            cache_scheduler=cache_scheduler,
        )
        error_classifier = ErrorClassifier(
            db=self.db,
            cache_scheduler=cache_scheduler,
            adaptive_manager=adaptive_manager,
        )
        request_dispatcher = RequestDispatcher(
            db=self.db,
            request_executor=request_executor,
            cache_scheduler=cache_scheduler if enable_cache_affinity else None,
        )

        pool_ops = self._sync_ops._pool_ops
        error_ops = self._sync_ops._error_ops
        failure_ops = self._sync_ops._failure_ops

        resolved_user = current_user
        if resolved_user is None and user_api_key is not None:
            try:
                resolved_user = user_api_key.user if hasattr(user_api_key, "user") else None
            except Exception:
                resolved_user = None
            if resolved_user is None and getattr(user_api_key, "user_id", None):
                resolved_user = self.db.query(User).filter(User.id == user_api_key.user_id).first()

        user_id: str | None = None
        if resolved_user is not None and getattr(resolved_user, "id", None):
            user_id = str(resolved_user.id)
        elif user_api_key is not None and getattr(user_api_key, "user_id", None):
            user_id = str(user_api_key.user_id)

        username_snapshot = getattr(resolved_user, "username", None) if resolved_user else None
        api_key_name_snapshot = getattr(user_api_key, "name", None) if user_api_key else None

        resolved_affinity_key = affinity_key
        if not resolved_affinity_key:
            api_key_id = getattr(user_api_key, "id", None) if user_api_key is not None else None
            resolved_affinity_key = str(api_key_id) if api_key_id else f"internal-test:{request_id}"

        if create_pending_usage:
            try:
                UsageService.create_pending_usage(
                    db=self.db,
                    request_id=request_id,
                    user=resolved_user,
                    api_key=user_api_key,
                    model=model_name,
                    is_stream=is_stream,
                    api_format=api_format_norm,
                    request_headers=request_headers,
                    request_body=request_body,
                )
            except Exception as exc:
                from src.core.logger import logger as _logger

                _logger.warning("创建 pending 使用记录失败: {}", str(exc))

        all_candidates = list(candidates)

        # 号池排序涉及大量 Redis 操作，提前释放 DB 连接避免连接池压力
        from src.services.scheduling.utils import release_db_connection_before_await

        release_db_connection_before_await(self.db)

        all_candidates, pool_traces = await pool_ops.apply_pool_reorder(
            all_candidates, request_body=request_body
        )

        candidate_record_map = await candidate_resolver.create_candidate_records_async(
            all_candidates=all_candidates,
            request_id=request_id,
            user_id=user_id,
            user_api_key=user_api_key,
            required_capabilities=capability_requirements,
        )

        max_attempts = candidate_resolver.count_total_attempts(all_candidates)
        execution_state = SyncExecutionState(
            candidate_record_map=candidate_record_map,
            request_body_state=request_body_state,
            last_candidate=all_candidates[-1] if all_candidates else None,
        )

        async def _attempt(candidate: Any) -> AttemptResult:
            execution_state.touch_candidate(candidate)

            candidate_index = int(getattr(candidate, "_utf_candidate_index", -1))
            retry_index = int(getattr(candidate, "_utf_retry_index", 0))
            candidate_record_id = str(getattr(candidate, "_utf_candidate_record_id", "") or "")
            attempt_counter = int(getattr(candidate, "_utf_attempt_count", 0))
            max_attempts_local = int(getattr(candidate, "_utf_max_attempts", max_attempts))

            if not candidate_record_id:
                from src.services.scheduling.schemas import PoolCandidate

                pool_extra = (
                    getattr(candidate.key, "_pool_extra_data", None)
                    if isinstance(getattr(candidate.key, "_pool_extra_data", None), dict)
                    else {}
                )
                extra_data: dict[str, Any] = {
                    "needs_conversion": bool(getattr(candidate, "needs_conversion", False)),
                    "provider_api_format": getattr(candidate, "provider_api_format", None) or None,
                    "mapping_matched_model": getattr(candidate, "mapping_matched_model", None)
                    or None,
                    **pool_extra,
                }
                if isinstance(candidate, PoolCandidate):
                    extra_data["pool_group_id"] = str(candidate.provider.id)
                    extra_data["pool_key_index"] = int(
                        getattr(candidate, "_pool_key_index", 0) or 0
                    )

                candidate_record = RequestCandidateService.create_candidate(
                    db=self.db,
                    request_id=request_id,
                    candidate_index=candidate_index,
                    retry_index=retry_index,
                    user_id=user_id,
                    api_key_id=(getattr(user_api_key, "id", None) if user_api_key else None),
                    username=username_snapshot,
                    api_key_name=api_key_name_snapshot,
                    provider_id=str(candidate.provider.id),
                    endpoint_id=str(candidate.endpoint.id),
                    key_id=str(candidate.key.id),
                    status="available",
                    is_cached=bool(getattr(candidate, "is_cached", False)),
                    extra_data=extra_data,
                )
                self.db.flush()
                candidate_record_id = str(candidate_record.id)
                execution_state.candidate_record_map[(candidate_index, retry_index)] = (
                    candidate_record_id
                )
                setattr(candidate, "_utf_candidate_record_id", candidate_record_id)

            (
                response,
                _provider_name,
                attempt_id,
                _provider_id,
                _endpoint_id,
                _key_id,
                _first_byte_time_ms,
            ) = await request_dispatcher.dispatch(
                candidate=candidate,
                candidate_index=candidate_index,
                retry_index=retry_index,
                candidate_record_id=candidate_record_id,
                user_api_key=user_api_key,
                user_id=user_id,
                request_func=request_func,
                request_id=request_id,
                api_format=api_format_norm,
                model_name=model_name,
                affinity_key=resolved_affinity_key,
                global_model_id=model_name,
                attempt_counter=attempt_counter,
                max_attempts=max_attempts_local,
                is_stream=is_stream,
            )
            _ = (attempt_id, _provider_name, _provider_id, _endpoint_id, _key_id)

            await pool_ops.pool_on_success(candidate, request_body)

            if is_stream:
                return AttemptResult(
                    kind=AttemptKind.STREAM,
                    http_status=200,
                    http_headers={},
                    stream_iterator=response,
                )
            return AttemptResult(
                kind=AttemptKind.SYNC_RESPONSE,
                http_status=200,
                http_headers={},
                response_body=response,
            )

        async def _handle_exec_err(
            *,
            exec_err: Any,
            candidate: Any,
            candidate_index: int,
            retry_index: int,
            max_retries_for_candidate: int,
            record_id: str | None,
            attempt_count: int,
            max_attempts: int | None,
        ) -> tuple[Any, int | None]:
            execution_state.track_execution_error(exec_err=exec_err, candidate=candidate)
            candidate_record_id = execution_state.resolve_candidate_record_id(
                candidate_index=candidate_index,
                record_id=record_id,
            )

            raw_action = await error_ops.handle_candidate_error(
                exec_err=exec_err,
                candidate=candidate,
                candidate_record_id=candidate_record_id,
                retry_index=retry_index,
                max_retries_for_candidate=max_retries_for_candidate,
                affinity_key=resolved_affinity_key,
                api_format=api_format_norm,
                global_model_id=model_name,
                request_id=request_id,
                attempt=attempt_count,
                max_attempts=int(max_attempts or 0),
                request_body_state=request_body_state,
                error_classifier=error_classifier,
            )
            action = classify_candidate_error_action(raw_action)

            if action == CandidateErrorAction.RAISE_ERROR:
                execution_state.raise_classified_error(
                    fallback_error=exec_err,
                    failure_ops=failure_ops,
                    model_name=model_name,
                    api_format=api_format_norm,
                )

            return resolve_execution_error_transition(
                action=action,
                state=execution_state,
                max_retries_for_candidate=max_retries_for_candidate,
                retry_index=retry_index,
            ).as_failover_tuple()

        engine = FailoverEngine(
            self.db,
            error_classifier=error_classifier,
            recorder=self._candidate_recorder,
        )
        result = await engine.execute(
            candidates=all_candidates,
            attempt_func=_attempt,
            retry_policy=RetryPolicy.for_sync_task(),
            skip_policy=SkipPolicy(),
            request_id=request_id,
            user_id=user_id,
            api_key_id=(
                str(user_api_key.id) if user_api_key and getattr(user_api_key, "id", None) else None
            ),
            username=username_snapshot,
            api_key_name=api_key_name_snapshot,
            candidate_record_map=candidate_record_map,
            max_attempts=max_attempts,
            execution_error_handler=_handle_exec_err,
            is_cancelled=is_cancelled,
        )

        if result.success:
            if pool_traces and result.key_id:
                try:
                    attempted_key_ids: set[str] = set()
                    for ck in result.candidate_keys or []:
                        status = str(getattr(ck, "status", "") or "").strip().lower()
                        if status in {"", "available", "pending", "skipped", "unused"}:
                            continue
                        kid = getattr(ck, "key_id", None)
                        if isinstance(kid, str) and kid:
                            attempted_key_ids.add(kid)
                    if not attempted_key_ids:
                        attempted_key_ids.add(str(result.key_id))

                    for pt in pool_traces:
                        summary = pt.build_summary(
                            result.key_id,
                            attempted_key_ids=attempted_key_ids,
                        )
                        if summary:
                            result.pool_summary = summary
                            break
                except Exception:
                    pass
            return result

        failure_ops.raise_all_failed_exception(
            request_id,
            max_attempts,
            execution_state.last_candidate,
            model_name,
            api_format_norm,
            execution_state.last_error,
        )

    async def submit_with_failover(
        self,
        *,
        api_format: str,
        model_name: str,
        affinity_key: str,
        user_api_key: ApiKey,
        request_id: str | None,
        task_type: str,
        submit_func: Any,
        extract_external_task_id: Any,
        supported_auth_types: set[str] | None = None,
        allow_format_conversion: bool = False,
        capability_requirements: dict[str, bool] | None = None,
        max_candidates: int | None = None,
        request_body: dict[str, Any] | None = None,
    ) -> Any:
        """
        Unified ASYNC submit entrypoint (Phase 3.2).

        兼容入口：默认直接绑定 AsyncTaskSubmitService。
        """
        return await self._submit_facade_ops.submit_with_failover(
            api_format=api_format,
            model_name=model_name,
            affinity_key=affinity_key,
            user_api_key=user_api_key,
            request_id=request_id,
            task_type=task_type,
            submit_func=submit_func,
            extract_external_task_id=extract_external_task_id,
            supported_auth_types=supported_auth_types,
            allow_format_conversion=allow_format_conversion,
            capability_requirements=capability_requirements,
            max_candidates=max_candidates,
            request_body=request_body,
        )

    # ====================
    # Phase 3.1: Async task helpers (poll/finalize)
    # ====================

    async def poll(self, task_id: str, *, user_id: str) -> TaskStatusResult:
        return await self._video_facade_ops.poll(task_id, user_id=user_id)

    async def poll_now(self, task_id: str, *, user_id: str) -> TaskStatusResult:
        return await self._video_facade_ops.poll_now(task_id, user_id=user_id)

    async def cancel(
        self,
        task_id: str,
        *,
        user_id: str,
        original_headers: dict[str, str] | None = None,
    ) -> Any:
        return await self._video_facade_ops.cancel(
            task_id,
            user_id=user_id,
            original_headers=original_headers,
        )

    async def finalize_video_task(self, task: Any) -> bool:
        return await self._video_facade_ops.finalize_video_task(task)

    async def finalize(self, task_id: str) -> bool:
        return await self._video_facade_ops.finalize(task_id)