feat: Add configuration option for apikey security (#322)

Signed-off-by: Michele Dolfi <dol@zurich.ibm.com>

docling_serve/app.py

@@ -18,6 +18,7 @@ from fastapi import (
     UploadFile,
     WebSocket,
     WebSocketDisconnect,
+    status,
 )
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.openapi.docs import (
@@ -48,6 +49,7 @@ from docling_jobkit.orchestrators.base_orchestrator import (
     TaskNotFoundError,
 )
 
+from docling_serve.auth import APIKeyAuth, AuthenticationResult
 from docling_serve.datamodel.convert import ConvertDocumentsRequestOptions
 from docling_serve.datamodel.requests import (
     ConvertDocumentsRequest,
@@ -156,6 +158,7 @@ def create_app():  # noqa: C901
         offline_docs_assets = True
         _log.info("Found static assets.")
 
+    require_auth = APIKeyAuth(docling_serve_settings.api_key)
     app = FastAPI(
         title="Docling Serve",
         docs_url=None if offline_docs_assets else "/swagger",
@@ -400,6 +403,7 @@ def create_app():  # noqa: C901
     )
     async def process_url(
         background_tasks: BackgroundTasks,
+        auth: Annotated[AuthenticationResult, Depends(require_auth)],
         orchestrator: Annotated[BaseOrchestrator, Depends(get_async_orchestrator)],
         conversion_request: ConvertDocumentsRequest,
     ):
@@ -443,6 +447,7 @@ def create_app():  # noqa: C901
     )
     async def process_file(
         background_tasks: BackgroundTasks,
+        auth: Annotated[AuthenticationResult, Depends(require_auth)],
         orchestrator: Annotated[BaseOrchestrator, Depends(get_async_orchestrator)],
         files: list[UploadFile],
         options: Annotated[
@@ -485,6 +490,7 @@ def create_app():  # noqa: C901
         response_model=TaskStatusResponse,
     )
     async def process_url_async(
+        auth: Annotated[AuthenticationResult, Depends(require_auth)],
         orchestrator: Annotated[BaseOrchestrator, Depends(get_async_orchestrator)],
         conversion_request: ConvertDocumentsRequest,
     ):
@@ -507,6 +513,7 @@ def create_app():  # noqa: C901
         response_model=TaskStatusResponse,
     )
     async def process_file_async(
+        auth: Annotated[AuthenticationResult, Depends(require_auth)],
         orchestrator: Annotated[BaseOrchestrator, Depends(get_async_orchestrator)],
         background_tasks: BackgroundTasks,
         files: list[UploadFile],
@@ -535,6 +542,7 @@ def create_app():  # noqa: C901
         response_model=TaskStatusResponse,
     )
     async def task_status_poll(
+        auth: Annotated[AuthenticationResult, Depends(require_auth)],
         orchestrator: Annotated[BaseOrchestrator, Depends(get_async_orchestrator)],
         task_id: str,
         wait: Annotated[
@@ -562,7 +570,15 @@ def create_app():  # noqa: C901
         websocket: WebSocket,
         orchestrator: Annotated[BaseOrchestrator, Depends(get_async_orchestrator)],
         task_id: str,
+        api_key: Annotated[str, Query()] = "",
     ):
+        if docling_serve_settings.api_key:
+            if api_key != docling_serve_settings.api_key:
+                raise HTTPException(
+                    status_code=status.HTTP_401_UNAUTHORIZED,
+                    detail="Api key is required as ?api_key=SECRET.",
+                )
+
         assert isinstance(orchestrator.notifier, WebsocketNotifier)
         await websocket.accept()
 
@@ -629,6 +645,7 @@ def create_app():  # noqa: C901
         },
     )
     async def task_result(
+        auth: Annotated[AuthenticationResult, Depends(require_auth)],
         orchestrator: Annotated[BaseOrchestrator, Depends(get_async_orchestrator)],
         background_tasks: BackgroundTasks,
         task_id: str,
@@ -656,6 +673,7 @@ def create_app():  # noqa: C901
         response_model=ProgressCallbackResponse,
     )
     async def callback_task_progress(
+        auth: Annotated[AuthenticationResult, Depends(require_auth)],
         orchestrator: Annotated[BaseOrchestrator, Depends(get_async_orchestrator)],
         request: ProgressCallbackRequest,
     ):
@@ -677,6 +695,7 @@ def create_app():  # noqa: C901
         response_model=ClearResponse,
     )
     async def clear_converters(
+        auth: Annotated[AuthenticationResult, Depends(require_auth)],
         orchestrator: Annotated[BaseOrchestrator, Depends(get_async_orchestrator)],
     ):
         await orchestrator.clear_converters()
@@ -688,6 +707,7 @@ def create_app():  # noqa: C901
         response_model=ClearResponse,
     )
     async def clear_results(
+        auth: Annotated[AuthenticationResult, Depends(require_auth)],
         orchestrator: Annotated[BaseOrchestrator, Depends(get_async_orchestrator)],
         older_then: float = 3600,
     ):

docling_serve/auth.py

@@ -0,0 +1,56 @@
+from typing import Any
+
+from fastapi import HTTPException, Request, status
+from fastapi.security import APIKeyHeader
+from pydantic import BaseModel
+
+
+class AuthenticationResult(BaseModel):
+    valid: bool
+    errors: list[str] = []
+    detail: Any | None = None
+
+
+class APIKeyAuth(APIKeyHeader):
+    """
+    FastAPI dependency which evaluates a status API Key.
+    """
+
+    def __init__(
+        self,
+        api_key: str,
+        header_name: str = "X-Api-Key",
+        fail_on_unauthorized: bool = True,
+    ) -> None:
+        self.api_key = api_key
+        self.header_name = header_name
+        super().__init__(name=self.header_name, auto_error=False)
+
+    async def _validate_api_key(self, header_api_key: str | None):
+        if header_api_key is None:
+            return AuthenticationResult(
+                valid=False, errors=[f"Missing header {self.header_name}."]
+            )
+
+        header_api_key = header_api_key.strip()
+
+        # Otherwise check the apikey
+        if header_api_key == self.api_key or self.api_key == "":
+            return AuthenticationResult(
+                valid=True,
+                detail=header_api_key,
+            )
+        else:
+            return AuthenticationResult(
+                valid=False,
+                errors=["The provided API Key is invalid."],
+            )
+
+    async def __call__(self, request: Request) -> AuthenticationResult:  # type: ignore
+        header_api_key = await super().__call__(request=request)
+        result = await self._validate_api_key(header_api_key)
+        if self.api_key and not result.valid:
+            raise HTTPException(
+                status_code=status.HTTP_401_UNAUTHORIZED, detail=result.detail
+            )
+        return result

docling_serve/gradio_ui.py

@@ -233,15 +233,21 @@ def change_ocr_lang(ocr_engine):
         return "english,chinese"
 
 
-def wait_task_finish(task_id: str, return_as_file: bool):
+def wait_task_finish(auth: str, task_id: str, return_as_file: bool):
     conversion_sucess = False
     task_finished = False
     task_status = ""
+
+    headers = {}
+    if docling_serve_settings.api_key:
+        headers["X-Api-Key"] = str(auth)
+
     ssl_ctx = get_ssl_context()
     while not task_finished:
         try:
             response = httpx.get(
                 f"{get_api_endpoint()}/v1/status/poll/{task_id}?wait=5",
+                headers=headers,
                 verify=ssl_ctx,
                 timeout=15,
             )
@@ -265,6 +271,7 @@ def wait_task_finish(task_id: str, return_as_file: bool):
         try:
             response = httpx.get(
                 f"{get_api_endpoint()}/v1/result/{task_id}",
+                headers=headers,
                 timeout=15,
                 verify=ssl_ctx,
             )
@@ -279,6 +286,7 @@ def wait_task_finish(task_id: str, return_as_file: bool):
 
 
 def process_url(
+    auth,
     input_sources,
     to_formats,
     image_export_mode,
@@ -326,11 +334,18 @@ def process_url(
     ):
         logger.error("No input sources provided.")
         raise gr.Error("No input sources provided.", print_exception=False)
+
+    headers = {}
+    if docling_serve_settings.api_key:
+        headers["X-Api-Key"] = str(auth)
+
+    print(f"{headers=}")
     try:
         ssl_ctx = get_ssl_context()
         response = httpx.post(
             f"{get_api_endpoint()}/v1/convert/source/async",
             json=parameters,
+            headers=headers,
             verify=ssl_ctx,
             timeout=60,
         )
@@ -354,6 +369,7 @@ def file_to_base64(file):
 
 
 def process_file(
+    auth,
     files,
     to_formats,
     image_export_mode,
@@ -402,11 +418,16 @@ def process_file(
         "target": target,
     }
 
+    headers = {}
+    if docling_serve_settings.api_key:
+        headers["X-Api-Key"] = str(auth)
+
     try:
         ssl_ctx = get_ssl_context()
         response = httpx.post(
             f"{get_api_endpoint()}/v1/convert/source/async",
             json=parameters,
+            headers=headers,
             verify=ssl_ctx,
             timeout=60,
         )
@@ -565,6 +586,15 @@ with gr.Blocks(
                 file_process_btn = gr.Button("Process File", scale=1)
                 file_reset_btn = gr.Button("Reset", scale=1)
 
+    # Auth
+    with gr.Row(visible=bool(docling_serve_settings.api_key)):
+        with gr.Column():
+            auth = gr.Textbox(
+                label="Authentication",
+                placeholder="API Key",
+                type="password",
+            )
+
     # Options
     with gr.Accordion("Options") as options:
         with gr.Row():
@@ -590,6 +620,7 @@ with gr.Blocks(
                     label="Image Export Mode",
                     value="embedded",
                 )
+
         with gr.Row():
             with gr.Column(scale=1, min_width=200):
                 pipeline = gr.Radio(
@@ -724,6 +755,7 @@ with gr.Blocks(
     ).then(
         process_url,
         inputs=[
+            auth,
             url_input,
             to_formats,
             image_export_mode,
@@ -750,7 +782,7 @@ with gr.Blocks(
         outputs=[content_output, file_output],
     ).then(
         wait_task_finish,
-        inputs=[task_id_rendered, return_as_file],
+        inputs=[auth, task_id_rendered, return_as_file],
         outputs=[
             output_markdown,
             output_markdown_rendered,
@@ -811,6 +843,7 @@ with gr.Blocks(
     ).then(
         process_file,
         inputs=[
+            auth,
             file_input,
             to_formats,
             image_export_mode,
@@ -837,7 +870,7 @@ with gr.Blocks(
         outputs=[content_output, file_output],
     ).then(
         wait_task_finish,
-        inputs=[task_id_rendered, return_as_file],
+        inputs=[auth, task_id_rendered, return_as_file],
         outputs=[
             output_markdown,
             output_markdown_rendered,

docling_serve/settings.py

@@ -51,6 +51,8 @@ class DoclingServeSettings(BaseSettings):
     enable_remote_services: bool = False
     allow_external_plugins: bool = False
 
+    api_key: str = ""
+
     max_document_timeout: float = 3_600 * 24 * 7  # 7 days
     max_num_pages: int = sys.maxsize
     max_file_size: int = sys.maxsize

docs/configuration.md

@@ -52,6 +52,7 @@ THe following table describes the options to configure the Docling Serve app.
 |  | `DOCLING_SERVE_CORS_ORIGINS` | `["*"]` | A list of origins that should be permitted to make cross-origin requests. |
 |  | `DOCLING_SERVE_CORS_METHODS` | `["*"]` | A list of HTTP methods that should be allowed for cross-origin requests. |
 |  | `DOCLING_SERVE_CORS_HEADERS` | `["*"]` | A list of HTTP request headers that should be supported for cross-origin requests. |
+|  | `DOCLING_SERVE_ENG_API_KEY` | | If specified, all the API requests must contain the header `X-Api-Key` with this value. |
 |  | `DOCLING_SERVE_ENG_KIND` | `local` | The compute engine to use for the async tasks. Possible values are `local`, `rq` and `kfp`. See below for more configurations of the engines. |
 
 ### Compute engine

docs/usage.md

@@ -30,6 +30,10 @@ On top of the source of file (see below), both endpoints support the same parame
 - `include_images` (bool): If enabled, images will be extracted from the document. Defaults to false.
 - `images_scale` (float): Scale factor for images. Defaults to 2.0.
 
+### Authentication
+
+When authentication is activated (see the parameter `DOCLING_SERVE_ENG_API_KEY` in [configuration.md](./configuration.md)), all the API requests **must** provide the header `X-Api-Key` with the correct secret key.
+
 ## Convert endpoints
 
 ### Source endpoint

tests/test_1-file-all-outputs.py

@@ -6,10 +6,15 @@ import pytest
 import pytest_asyncio
 from pytest_check import check
 
+from docling_serve.settings import docling_serve_settings
+
 
 @pytest_asyncio.fixture
 async def async_client():
-    async with httpx.AsyncClient(timeout=60.0) as client:
+    headers = {}
+    if docling_serve_settings.api_key:
+        headers["X-Api-Key"] = docling_serve_settings.api_key
+    async with httpx.AsyncClient(timeout=60.0, headers=headers) as client:
         yield client
 
 

tests/test_1-file-async.py

@@ -6,10 +6,15 @@ import httpx
 import pytest
 import pytest_asyncio
 
+from docling_serve.settings import docling_serve_settings
+
 
 @pytest_asyncio.fixture
 async def async_client():
-    async with httpx.AsyncClient(timeout=60.0) as client:
+    headers = {}
+    if docling_serve_settings.api_key:
+        headers["X-Api-Key"] = docling_serve_settings.api_key
+    async with httpx.AsyncClient(timeout=60.0, headers=headers) as client:
         yield client
 
 

tests/test_1-url-all-outputs.py

@@ -5,10 +5,15 @@ import pytest
 import pytest_asyncio
 from pytest_check import check
 
+from docling_serve.settings import docling_serve_settings
+
 
 @pytest_asyncio.fixture
 async def async_client():
-    async with httpx.AsyncClient(timeout=60.0) as client:
+    headers = {}
+    if docling_serve_settings.api_key:
+        headers["X-Api-Key"] = docling_serve_settings.api_key
+    async with httpx.AsyncClient(timeout=60.0, headers=headers) as client:
         yield client
 
 

tests/test_1-url-async-ws.py

@@ -6,16 +6,24 @@ import pytest
 import pytest_asyncio
 from websockets.sync.client import connect
 
+from docling_serve.settings import docling_serve_settings
+
 
 @pytest_asyncio.fixture
 async def async_client():
-    async with httpx.AsyncClient(timeout=60.0) as client:
+    headers = {}
+    if docling_serve_settings.api_key:
+        headers["X-Api-Key"] = docling_serve_settings.api_key
+    async with httpx.AsyncClient(timeout=60.0, headers=headers) as client:
         yield client
 
 
 @pytest.mark.asyncio
 async def test_convert_url(async_client: httpx.AsyncClient):
     """Test convert URL to all outputs"""
+    headers = {}
+    if docling_serve_settings.api_key:
+        headers["X-Api-Key"] = docling_serve_settings.api_key
 
     doc_filename = Path("tests/2408.09869v5.pdf")
     encoded_doc = base64.b64encode(doc_filename.read_bytes()).decode()
@@ -57,7 +65,7 @@ async def test_convert_url(async_client: httpx.AsyncClient):
 
     task = response.json()
 
-    uri = f"ws://localhost:5001/v1/status/ws/{task['task_id']}"
+    uri = f"ws://localhost:5001/v1/status/ws/{task['task_id']}?api_key={docling_serve_settings.api_key}"
     with connect(uri) as websocket:
         for message in websocket:
             print(message)

tests/test_1-url-async.py

@@ -6,10 +6,15 @@ import httpx
 import pytest
 import pytest_asyncio
 
+from docling_serve.settings import docling_serve_settings
+
 
 @pytest_asyncio.fixture
 async def async_client():
-    async with httpx.AsyncClient(timeout=60.0) as client:
+    headers = {}
+    if docling_serve_settings.api_key:
+        headers["X-Api-Key"] = docling_serve_settings.api_key
+    async with httpx.AsyncClient(timeout=60.0, headers=headers) as client:
         yield client
 
 

tests/test_2-files-all-outputs.py

@@ -5,10 +5,15 @@ import pytest
 import pytest_asyncio
 from pytest_check import check
 
+from docling_serve.settings import docling_serve_settings
+
 
 @pytest_asyncio.fixture
 async def async_client():
-    async with httpx.AsyncClient(timeout=60.0) as client:
+    headers = {}
+    if docling_serve_settings.api_key:
+        headers["X-Api-Key"] = docling_serve_settings.api_key
+    async with httpx.AsyncClient(timeout=60.0, headers=headers) as client:
         yield client
 
 

tests/test_2-urls-all-outputs.py

@@ -3,10 +3,15 @@ import pytest
 import pytest_asyncio
 from pytest_check import check
 
+from docling_serve.settings import docling_serve_settings
+
 
 @pytest_asyncio.fixture
 async def async_client():
-    async with httpx.AsyncClient(timeout=60.0) as client:
+    headers = {}
+    if docling_serve_settings.api_key:
+        headers["X-Api-Key"] = docling_serve_settings.api_key
+    async with httpx.AsyncClient(timeout=60.0, headers=headers) as client:
         yield client
 
 

tests/test_2-urls-async-all-outputs.py

@@ -6,10 +6,15 @@ import pytest
 import pytest_asyncio
 from pytest_check import check
 
+from docling_serve.settings import docling_serve_settings
+
 
 @pytest_asyncio.fixture
 async def async_client():
-    async with httpx.AsyncClient(timeout=60.0) as client:
+    headers = {}
+    if docling_serve_settings.api_key:
+        headers["X-Api-Key"] = docling_serve_settings.api_key
+    async with httpx.AsyncClient(timeout=60.0, headers=headers) as client:
         yield client
 
 

tests/test_fastapi_endpoints.py

@@ -13,6 +13,7 @@ from pytest_check import check
 from docling_core.types.doc import DoclingDocument, PictureItem
 
 from docling_serve.app import create_app
+from docling_serve.settings import docling_serve_settings
 
 
 @pytest.fixture(scope="session")
@@ -20,6 +21,14 @@ def event_loop():
     return asyncio.get_event_loop()
 
 
+@pytest.fixture(scope="session")
+def auth_headers():
+    headers = {}
+    if docling_serve_settings.api_key:
+        headers["X-Api-Key"] = docling_serve_settings.api_key
+    return headers
+
+
 @pytest_asyncio.fixture(scope="session")
 async def app():
     app = create_app()
@@ -46,7 +55,7 @@ async def test_health(client: AsyncClient):
 
 
 @pytest.mark.asyncio
-async def test_convert_file(client: AsyncClient):
+async def test_convert_file(client: AsyncClient, auth_headers: dict):
     """Test convert single file to all outputs"""
 
     endpoint = "/v1/convert/file"
@@ -79,7 +88,9 @@ async def test_convert_file(client: AsyncClient):
         "files": ("2206.01062v1.pdf", open(file_path, "rb"), "application/pdf"),
     }
 
-    response = await client.post(endpoint, files=files, data=options)
+    response = await client.post(
+        endpoint, files=files, data=options, headers=auth_headers
+    )
     assert response.status_code == 200, "Response should be 200 OK"
 
     data = response.json()
@@ -160,7 +171,7 @@ async def test_convert_file(client: AsyncClient):
 
 
 @pytest.mark.asyncio
-async def test_referenced_artifacts(client: AsyncClient):
+async def test_referenced_artifacts(client: AsyncClient, auth_headers: dict):
     """Test that paths in the zip file are relative to the zip file root."""
 
     endpoint = "/v1/convert/file"
@@ -178,7 +189,9 @@ async def test_referenced_artifacts(client: AsyncClient):
         "files": ("2206.01062v1.pdf", open(file_path, "rb"), "application/pdf"),
     }
 
-    response = await client.post(endpoint, files=files, data=options)
+    response = await client.post(
+        endpoint, files=files, data=options, headers=auth_headers
+    )
     assert response.status_code == 200, "Response should be 200 OK"
 
     with zipfile.ZipFile(io.BytesIO(response.content)) as zip_file:

tests/test_file_opts.py

@@ -11,6 +11,7 @@ from docling_core.types import DoclingDocument
 from docling_core.types.doc.document import PictureDescriptionData
 
 from docling_serve.app import create_app
+from docling_serve.settings import docling_serve_settings
 
 
 @pytest.fixture(scope="session")
@@ -18,6 +19,14 @@ def event_loop():
     return asyncio.get_event_loop()
 
 
+@pytest.fixture(scope="session")
+def auth_headers():
+    headers = {}
+    if docling_serve_settings.api_key:
+        headers["X-Api-Key"] = docling_serve_settings.api_key
+    return headers
+
+
 @pytest_asyncio.fixture(scope="session")
 async def app():
     app = create_app()
@@ -37,7 +46,7 @@ async def client(app):
 
 
 @pytest.mark.asyncio
-async def test_convert_file(client: AsyncClient):
+async def test_convert_file(client: AsyncClient, auth_headers: dict):
     """Test convert single file to all outputs"""
 
     endpoint = "/v1/convert/file"
@@ -63,7 +72,9 @@ async def test_convert_file(client: AsyncClient):
         "files": ("2206.01062v1.pdf", open(file_path, "rb"), "application/pdf"),
     }
 
-    response = await client.post(endpoint, files=files, data=options)
+    response = await client.post(
+        endpoint, files=files, data=options, headers=auth_headers
+    )
     assert response.status_code == 200, "Response should be 200 OK"
 
     data = response.json()

tests/test_results_clear.py

@@ -17,6 +17,14 @@ def event_loop():
     return asyncio.get_event_loop()
 
 
+@pytest.fixture(scope="session")
+def auth_headers():
+    headers = {}
+    if docling_serve_settings.api_key:
+        headers["X-Api-Key"] = docling_serve_settings.api_key
+    return headers
+
+
 @pytest_asyncio.fixture(scope="session")
 async def app():
     app = create_app()
@@ -35,7 +43,7 @@ async def client(app):
         yield client
 
 
-async def convert_file(client: AsyncClient):
+async def convert_file(client: AsyncClient, auth_headers: dict):
     doc_filename = Path("tests/2408.09869v5.pdf")
     encoded_doc = base64.b64encode(doc_filename.read_bytes()).decode()
 
@@ -52,7 +60,9 @@ async def convert_file(client: AsyncClient):
         ],
     }
 
-    response = await client.post("/v1/convert/source/async", json=payload)
+    response = await client.post(
+        "/v1/convert/source/async", json=payload, headers=auth_headers
+    )
     assert response.status_code == 200, "Response should be 200 OK"
 
     task = response.json()
@@ -60,7 +70,9 @@ async def convert_file(client: AsyncClient):
     print(json.dumps(task, indent=2))
 
     while task["task_status"] not in ("success", "failure"):
-        response = await client.get(f"/v1/status/poll/{task['task_id']}")
+        response = await client.get(
+            f"/v1/status/poll/{task['task_id']}", headers=auth_headers
+        )
         assert response.status_code == 200, "Response should be 200 OK"
         task = response.json()
         print(f"{task['task_status']=}")
@@ -74,52 +86,62 @@ async def convert_file(client: AsyncClient):
 
 
 @pytest.mark.asyncio
-async def test_clear_results(client: AsyncClient):
+async def test_clear_results(client: AsyncClient, auth_headers: dict):
     """Test removal of task."""
 
     # Set long delay deletion
     docling_serve_settings.result_removal_delay = 100
 
     # Convert and wait for completion
-    task = await convert_file(client)
+    task = await convert_file(client, auth_headers=auth_headers)
 
     # Get result once
-    result_response = await client.get(f"/v1/result/{task['task_id']}")
+    result_response = await client.get(
+        f"/v1/result/{task['task_id']}", headers=auth_headers
+    )
     assert result_response.status_code == 200, "Response should be 200 OK"
     print("Result 1 ok.")
     result = result_response.json()
     assert result["document"]["json_content"]["schema_name"] == "DoclingDocument"
 
     # Get result twice
-    result_response = await client.get(f"/v1/result/{task['task_id']}")
+    result_response = await client.get(
+        f"/v1/result/{task['task_id']}", headers=auth_headers
+    )
     assert result_response.status_code == 200, "Response should be 200 OK"
     print("Result 2 ok.")
     result = result_response.json()
     assert result["document"]["json_content"]["schema_name"] == "DoclingDocument"
 
     # Clear
-    clear_response = await client.get("/v1/clear/results?older_then=0")
+    clear_response = await client.get(
+        "/v1/clear/results?older_then=0", headers=auth_headers
+    )
     assert clear_response.status_code == 200, "Response should be 200 OK"
     print("Clear ok.")
 
     # Get deleted result
-    result_response = await client.get(f"/v1/result/{task['task_id']}")
+    result_response = await client.get(
+        f"/v1/result/{task['task_id']}", headers=auth_headers
+    )
     assert result_response.status_code == 404, "Response should be removed"
     print("Result was no longer found.")
 
 
 @pytest.mark.asyncio
-async def test_delay_remove(client: AsyncClient):
+async def test_delay_remove(client: AsyncClient, auth_headers: dict):
     """Test automatic removal of task with delay."""
 
     # Set short delay deletion
     docling_serve_settings.result_removal_delay = 5
 
     # Convert and wait for completion
-    task = await convert_file(client)
+    task = await convert_file(client, auth_headers=auth_headers)
 
     # Get result once
-    result_response = await client.get(f"/v1/result/{task['task_id']}")
+    result_response = await client.get(
+        f"/v1/result/{task['task_id']}", headers=auth_headers
+    )
     assert result_response.status_code == 200, "Response should be 200 OK"
     print("Result ok.")
     result = result_response.json()
@@ -129,5 +151,7 @@ async def test_delay_remove(client: AsyncClient):
     await asyncio.sleep(10)
 
     # Get deleted result
-    result_response = await client.get(f"/v1/result/{task['task_id']}")
+    result_response = await client.get(
+        f"/v1/result/{task['task_id']}", headers=auth_headers
+    )
     assert result_response.status_code == 404, "Response should be removed"