""" Сервис api (#78): регистрирует задачи в очереди RabbitMQ и отдаёт результаты из Redis. Chrome здесь нет — задачи выполняет сервис worker (src/worker). * POST /v2/tasks, GET /v2/tasks/{id}, DELETE /v2/sessions/current — асинхронный API (#80, #82). * POST /v1/solve — прежний синхронный контракт поверх той же очереди (#84). Все эндпоинты, кроме /health, — только с ключом клиента в X-API-Key (#4); ключи — в Postgres (db.cli). Задача и окно сессии видны только клиенту, который их создал. """ import asyncio import logging import os import time import uuid from contextlib import asynccontextmanager from typing import Annotated from fastapi import Cookie, Depends, FastAPI, HTTPException, Response, Security from fastapi.responses import JSONResponse from fastapi.security import APIKeyHeader from starlette.concurrency import run_in_threadpool logging.basicConfig( level=os.getenv("LOG_LEVEL", "INFO").upper(), format="%(asctime)s %(levelname)s %(name)s: %(message)s", ) from api.schemas import SolveRequest, SolveResponse, TaskCreated, TaskRequest, TaskResponse # noqa: E402 from db import clients # noqa: E402 from db.engine import connect as connect_db # noqa: E402 from db.models import Client # noqa: E402 from tasks import broker as queue # noqa: E402 from tasks.sessions import new_session_id, parse_session_id # noqa: E402 from tasks.store import Status, TaskStore, connect # noqa: E402 logger = logging.getLogger(__name__) RETRY_AFTER = os.getenv("NO_SLOTS_RETRY_AFTER", "5") # сколько /v1/solve ждёт, пока воркер возьмёт задачу, прежде чем ответить 503 no_slots V1_QUEUE_WAIT = float(os.getenv("V1_QUEUE_WAIT", "30")) POLL_INTERVAL = 0.2 store = TaskStore(connect()) database = connect_db() @asynccontextmanager async def lifespan(app: FastAPI): await database.init() yield await database.close() app = FastAPI(lifespan=lifespan) api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False) async def current_client(key: Annotated[str | None, Security(api_key_header)]) -> Client: client = await clients.authenticate(database, key) if client is None: raise HTTPException(401, "invalid or missing X-API-Key") return client ClientDep = Annotated[Client, Depends(current_client)] NO_SLOTS = {"status": "err", "code": "no_slots", "error": "No slots available"} def no_slots_response() -> JSONResponse: return JSONResponse(status_code=503, headers={"Retry-After": RETRY_AFTER}, content=NO_SLOTS) def use_session(response: Response, cookie: str | None) -> str: """session_id из cookie или новый; cookie ставится в ответ.""" session_id = parse_session_id(cookie) if cookie and not session_id: logger.warning("invalid session_id cookie, issuing a new session") session_id = session_id or new_session_id() response.set_cookie("session_id", session_id) return session_id async def dispatch(message) -> None: # pika блокирующий — отправка в пуле потоков, чтобы не останавливать event loop await run_in_threadpool(queue.broker.enqueue, message) async def submit(payload: TaskRequest, session_id: str, window: dict | None, client: Client, api: str) -> str: """Зарегистрировать задачу и поставить её в очередь: в персональную очередь воркера, где открыто окно сессии (#83), иначе — в общую.""" task_id = str(uuid.uuid4()) await store.create(task_id, session_id, payload.url, client.id) target = queue.personal_queue(window["worker_id"]) if window else queue.COMMON_QUEUE await dispatch(queue.task_message( task_id, session_id, payload.model_dump(mode="json"), payload.browser_idle_timeout, target, meta={"client_id": client.id, "api": api}, )) logger.info("task %s [%s] -> %s: %s", task_id, session_id, target, payload.url or "") return task_id async def free_slots() -> int: return sum(w["free_slots"] for w in await store.workers()) async def client_window(session_id: str, client: Client) -> dict | None: """Открытое окно сессии, если оно этого клиента. Чужое окно — как отсутствующее.""" window = await store.live_window(session_id) return window if window and window.get("client_id") == str(client.id) else None @app.get("/") @app.get("/health") async def health(): workers = await store.workers() return { "status": "ok", "browser_slots": sum(w["browser_slots"] for w in workers), "free_slots": sum(w["free_slots"] for w in workers), "workers": workers, } @app.post("/v2/tasks", status_code=202, response_model=TaskCreated) async def create_task( payload: TaskRequest, response: Response, client: ClientDep, session_id: Annotated[str | None, Cookie()] = None, ): session_id = use_session(response, session_id) window = await client_window(session_id, client) if payload.url is None and window is None: raise HTTPException(422, "url is required: the session has no open browser window") task_id = await submit(payload, session_id, window, client, "v2") return {"task_id": task_id, "status": Status.QUEUED, "url": f"/v2/tasks/{task_id}"} @app.get("/v2/tasks/{task_id}", response_model=TaskResponse, response_model_exclude_none=True) async def get_task(task_id: str, client: ClientDep, session_id: Annotated[str | None, Cookie()] = None): # задача видна только своему клиенту и своей сессии; чужая и несуществующая неразличимы task = await store.get(task_id) if parse_session_id(task_id) else None if not task or task["session_id"] != parse_session_id(session_id) or task.get("client_id") != str(client.id): raise HTTPException(404, "task not found") return {"task_id": task_id, "status": task["status"], "error": task.get("error"), "result": task.get("result")} @app.delete("/v2/sessions/current") async def close_session(client: ClientDep, session_id: Annotated[str | None, Cookie()] = None): """Закрыть окно сессии досрочно и освободить слот.""" session_id = parse_session_id(session_id) window = await store.get_session(session_id) if session_id else None if not window or window.get("client_id") != str(client.id): return {"closed": False} await store.delete_session(session_id) if await store.worker_alive(window["worker_id"]): await dispatch(queue.close_message(session_id, queue.personal_queue(window["worker_id"]))) return {"closed": True} @app.post("/v1/solve", response_model=SolveResponse, response_model_exclude_none=True) async def solve( payload: SolveRequest, response: Response, client: ClientDep, session_id: Annotated[str | None, Cookie()] = None, ): """Синхронный контракт v1: задача ставится в ту же очередь, ответ — когда воркер её выполнит. Если свободных слотов нет или воркер не взял задачу за V1_QUEUE_WAIT — 503, как и раньше.""" started = time.monotonic() session_id = use_session(response, session_id) window = await client_window(session_id, client) if window is None and await free_slots() <= 0: logger.warning("no free browser slots, rejecting %s", payload.url) return no_slots_response() task_id = await submit(TaskRequest.model_validate(payload.model_dump()), session_id, window, client, "v1") while True: task = await store.get(task_id) if task is None: result = {"status": "err", "code": "request_failed", "error": "task expired"} break if task["status"] == Status.DONE: result = task["result"] break if task["status"] in (Status.FAILED, Status.CANCELLED): result = {"status": "err", "code": "request_failed", "error": task.get("error") or task["status"]} break if ( task["status"] == Status.QUEUED and time.monotonic() - started > V1_QUEUE_WAIT and await store.cancel(task_id) ): logger.warning("task %s was not picked up in %ss, rejecting", task_id, V1_QUEUE_WAIT) return no_slots_response() await asyncio.sleep(POLL_INTERVAL) result.setdefault("actions", []) result["duration_ms"] = int((time.monotonic() - started) * 1000) return result