robotex/src/main.py

"""
Сервис 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 "<current page>")
	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