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).
"""
import asyncio
import logging
import os
import time
import uuid
from typing import Annotated

from fastapi import Cookie, FastAPI, HTTPException, Response
from fastapi.responses import JSONResponse
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 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())

app = FastAPI()

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) -> str:
	"""Зарегистрировать задачу и поставить её в очередь: в персональную очередь воркера, где открыто
	окно сессии (#83), иначе — в общую."""
	task_id = str(uuid.uuid4())
	await store.create(task_id, session_id, payload.url)
	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,
	))
	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())


@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,
	session_id: Annotated[str | None, Cookie()] = None,
):
	# TODO: добавить поддержку api ключей (#4)
	session_id = use_session(response, session_id)
	window = await store.live_window(session_id)
	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)
	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, 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):
		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(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:
		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,
	session_id: Annotated[str | None, Cookie()] = None,
):
	"""Синхронный контракт v1: задача ставится в ту же очередь, ответ — когда воркер её выполнит.
	Если свободных слотов нет или воркер не взял задачу за V1_QUEUE_WAIT — 503, как и раньше."""
	# TODO: добавить поддержку api ключей (#4)
	started = time.monotonic()
	session_id = use_session(response, session_id)
	window = await store.live_window(session_id)
	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)
	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