112 lines
3.8 KiB
Python
112 lines
3.8 KiB
Python
import asyncio
|
|
from contextlib import asynccontextmanager
|
|
|
|
from fastapi import FastAPI, HTTPException
|
|
from fastapi.responses import JSONResponse
|
|
|
|
from api.browser_rpc_client import BrowserRpcError
|
|
from api.config import settings
|
|
from api.schemas import BrowserTaskRequest, BrowserTaskAcceptedResponse, BrowserTaskStatusResponse, \
|
|
BrowserTaskResultResponse, TaskStatus
|
|
from api.task_store import TaskStore
|
|
from api.browser_rpc_client import run_browser_task
|
|
|
|
store = TaskStore()
|
|
_semaphore = asyncio.Semaphore(settings.max_concurrency)
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI):
|
|
yield
|
|
|
|
|
|
app = FastAPI(title="Browser API", version="1.0.0", lifespan=lifespan)
|
|
|
|
|
|
async def _worker(task_id: str) -> None:
|
|
rec = await store.set_running(task_id)
|
|
if rec is None:
|
|
return None
|
|
|
|
async with _semaphore:
|
|
try:
|
|
raw = await asyncio.wait_for(
|
|
run_browser_task(
|
|
rpc_url=settings.browser_rpc_url,
|
|
task=rec.task,
|
|
timeout_sec=float(rec.timeout),
|
|
),
|
|
timeout=float(rec.timeout) + 5,
|
|
)
|
|
success = bool(raw.get("success"))
|
|
await store.set_done(task_id=task_id, success=success, raw_response=raw, error=None)
|
|
except asyncio.TimeoutError:
|
|
await store.set_done(
|
|
task_id=task_id,
|
|
success=False,
|
|
raw_response=None,
|
|
error="Timeout exceeded",
|
|
)
|
|
except BrowserRpcError as exc:
|
|
await store.set_done(
|
|
task_id=task_id,
|
|
success=False,
|
|
raw_response=None,
|
|
error=str(exc),
|
|
)
|
|
except Exception as exc:
|
|
await store.set_done(
|
|
task_id=task_id,
|
|
success=False,
|
|
raw_response=None,
|
|
error=f"Internal error: {exc}",
|
|
)
|
|
|
|
|
|
@app.get("/health")
|
|
async def health() -> dict:
|
|
return {"ok": True}
|
|
|
|
|
|
@app.post("/api/browser/tasks", response_model=BrowserTaskAcceptedResponse, status_code=202)
|
|
async def create_task(payload: BrowserTaskRequest) -> BrowserTaskAcceptedResponse:
|
|
rec = await store.create(task=payload.task.strip(), timeout=payload.timeout, metadata=payload.metadata)
|
|
asyncio.create_task(_worker(rec.task_id))
|
|
return BrowserTaskAcceptedResponse(task_id=rec.task_id, status=rec.status)
|
|
|
|
|
|
@app.get("/api/browser/tasks/{task_id}", response_model=BrowserTaskStatusResponse)
|
|
async def get_task_status(task_id: str) -> BrowserTaskStatusResponse:
|
|
rec = await store.get(task_id=task_id)
|
|
if rec is None:
|
|
raise HTTPException(status_code=404, detail="Task not found")
|
|
return BrowserTaskStatusResponse(task_id=rec.task_id, status=rec.status, create_at=rec.create_at,
|
|
started_at=rec.started_at, finished_at=rec.finished_at, error=rec.error)
|
|
|
|
|
|
@app.get("/api/browser/tasks/{task_id}/result", response_model=BrowserTaskResultResponse)
|
|
async def get_task_result(task_id: str) -> BrowserTaskResultResponse:
|
|
rec = await store.get(task_id=task_id)
|
|
if rec is None:
|
|
raise HTTPException(status_code=404, detail="Task not found")
|
|
|
|
if rec.status in (TaskStatus.queued, TaskStatus.running):
|
|
return JSONResponse(status_code=202, content={
|
|
"task_id": rec.task_id,
|
|
"status": rec.status,
|
|
"success": False,
|
|
"execution_time": rec.execution_time,
|
|
"result": None,
|
|
"error": None,
|
|
"raw_response": None,
|
|
})
|
|
|
|
return BrowserTaskResultResponse(
|
|
task_id=rec.task_id,
|
|
status=rec.status,
|
|
success=(rec.status == "succeeded"),
|
|
execution_time=rec.execution_time,
|
|
result=rec.result,
|
|
error=rec.error,
|
|
raw_response=rec.raw_response,
|
|
)
|