app/routers/take.py

@@ -15,6 +15,9 @@ class Task(pydantic.BaseModel):
 
 
 class Response(pydantic.BaseModel):
+    id: str
+    attempt: int
+    payload: dict
     task: Task|None
 
 
@@ -22,5 +25,5 @@ class Response(pydantic.BaseModel):
 async def execute(queue: typing.Annotated[str, fastapi.Header()]) -> Response:
     task = await tasks.take_task(queue)
     if not task:
-        return Response(task=None)
+        raise fastapi.HTTPException(404)
-    return Response(task=Task(id=str(task._id), attempt=task.attempts, payload=task.payload))
+    return Response(id=str(task._id), attempt=task.attempts, payload=task.payload, task=Task(id=str(task._id), attempt=task.attempts, payload=task.payload))

app/storage/mongo/tasks.py

@@ -1,6 +1,7 @@
 import bson
 import datetime
 import pydantic
+import typing
 
 from app.storage.mongo import database
 from app.utils import time
@@ -26,7 +27,7 @@ async def add_task(task: Task) -> str:
     return result.inserted_id
 
 
-async def take_task(queue: str) -> Task|None:
+async def take_task(queue: str) -> typing.Optional[Task]:
     now = time.now()
     async for raw_task in collection.find({'queue': queue, 'available_from': {'$lte': now}}):
         task = Task.model_validate(raw_task)