final pass

Author: PastelStorm

CHANGELOG.md

@@ -6,6 +6,7 @@
 
 ### Fixes
 * Route split-PDF `partition_async()` result collection through awaited async hook dispatch instead of creating a nested event loop in a worker thread.
+  Sync-only hooks on the async path now run on a worker thread, so hook code that depends on event-loop-thread `contextvars` or thread-local state should pass that state explicitly.
 * Add cancellation cleanup for in-flight split-PDF chunk tasks and preserve existing sync `partition()` split-PDF behavior with lazy executor creation.
 
 ## 0.43.2

…gration/test_split_pdf_async_behavior.py …nt/unit/test_split_pdf_async_behavior.py_test_unstructured_client/integration/test_split_pdf_async_behavior.py renamed to _test_unstructured_client/unit/test_split_pdf_async_behavior.py _test_unstructured_client/integration/test_split_pdf_async_behavior.py renamed to _test_unstructured_client/unit/test_split_pdf_async_behavior.py

@@ -110,7 +110,7 @@ def _make_sdk_split_task(
 
 @pytest.mark.asyncio
 async def test_partition_async_split_collects_chunks_in_order_without_executor():
-    operation_id = "integration-happy"
+    operation_id = "unit-happy"
     split_hook = SplitPdfHook()
     active_chunks = 0
     max_active_chunks = 0
@@ -191,7 +191,7 @@ def _chunk_client_factory(*args, **kwargs):
 
 @pytest.mark.asyncio
 async def test_partition_async_cancellation_cleans_split_state_and_tempdir():
-    operation_id = "integration-cancel"
+    operation_id = "unit-cancel"
     split_hook = SplitPdfHook()
     started = asyncio.Event()
     cancelled_counter = Counter()
@@ -236,7 +236,7 @@ async def _hanging_chunk(
 
 @pytest.mark.asyncio
 async def test_partition_async_strict_failure_drains_sibling_chunks_before_close():
-    operation_id = "integration-strict-failure"
+    operation_id = "unit-strict-failure"
     split_hook = SplitPdfHook()
     sibling_started = asyncio.Event()
     sibling_cancelled = asyncio.Event()
@@ -370,7 +370,7 @@ async def _release_sync_hook_from_loop() -> None:
 
 @pytest.mark.asyncio
 async def test_partition_async_reassembly_does_not_block_event_loop():
-    operation_id = "integration-offload"
+    operation_id = "unit-offload"
     split_hook = SplitPdfHook()
     reassembly_started = threading.Event()
     reassembly_released = threading.Event()

_test_unstructured_client/unit/test_split_pdf_hook.py

@@ -4,17 +4,17 @@
 import io
 import logging
 import threading
-from asyncio import Task
 from collections import Counter
 from concurrent import futures
 from functools import partial
 from pathlib import Path
+from typing import Any, Coroutine
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import httpx
 import pytest
-import requests
-from requests_toolbelt import MultipartDecoder
+import requests  # type: ignore[import-untyped]
+from requests_toolbelt import MultipartDecoder  # type: ignore[import-untyped]
 
 from unstructured_client._hooks.custom import form_utils, pdf_utils, request_utils
 from unstructured_client._hooks.custom.form_utils import (
@@ -438,7 +438,7 @@ async def _request_mock(
 @pytest.mark.asyncio
 async def test_unit_disallow_failed_coroutines(
         allow_failed: bool,
-        tasks: list[Task],
+        tasks: list[partial[Coroutine[Any, Any, httpx.Response]]],
         expected_responses: list[str],
 ):
     """Test disallow failed coroutines method properly sets the flag to False."""
@@ -824,6 +824,70 @@ def after_error(self, hook_ctx, response, error):  # pragma: no cover - dispatch
     assert "Cancellation cleanup cancelled" in caplog.text
 
 
+@pytest.mark.asyncio
+async def test_unit_do_request_async_secondary_cancellation_waits_for_cleanup():
+    cleanup_started = asyncio.Event()
+    release_cleanup = asyncio.Event()
+    cleanup_finished = asyncio.Event()
+
+    class PreparedRequestHook:
+        def before_request(self, hook_ctx, request):
+            del hook_ctx, request
+            return httpx.Request(
+                "GET",
+                "http://localhost:8888/general/docs",
+                headers={"operation_id": "secondary-cancel-cleanup"},
+                extensions={"split_pdf_operation_id": "secondary-cancel-cleanup"},
+            )
+
+    class SlowCleanupHook:
+        async def after_error_async(self, hook_ctx, response, error):
+            del hook_ctx, response, error
+            cleanup_started.set()
+            await release_cleanup.wait()
+            cleanup_finished.set()
+            return None, None
+
+        def after_error(self, hook_ctx, response, error):  # pragma: no cover - dispatch guard
+            raise AssertionError("async hook should be awaited")
+
+    hooks = SDKHooks()
+    hooks.before_request_hooks = [PreparedRequestHook()]  # type: ignore[list-item]
+    hooks.after_error_hooks = [SlowCleanupHook()]  # type: ignore[list-item]
+
+    client = _BlockingAsyncClient()
+    config = SDKConfiguration(
+        client=None,
+        client_supplied=False,
+        async_client=client,  # type: ignore[arg-type]
+        async_client_supplied=True,
+        debug_logger=logging.getLogger("test"),
+    )
+    config.__dict__["_hooks"] = hooks
+    sdk = BaseSDK(config)
+    task = asyncio.create_task(
+        sdk.do_request_async(
+            _make_sdk_hook_context(),
+            httpx.Request("POST", "http://localhost:8888/general/v0/general"),
+            error_status_codes=[],
+        )
+    )
+
+    await client.started.wait()
+    task.cancel()
+    await cleanup_started.wait()
+    task.cancel()
+    await asyncio.sleep(0)
+
+    assert not task.done()
+
+    release_cleanup.set()
+    with pytest.raises(asyncio.CancelledError):
+        await task
+
+    assert cleanup_finished.is_set()
+
+
 def test_before_request_returns_dummy_with_timeout_and_operation_id():
     hook, mock_hook_ctx, result = _make_hook_with_split_request()
     operation_id = result.headers["operation_id"]
@@ -836,6 +900,17 @@ def test_before_request_returns_dummy_with_timeout_and_operation_id():
     assert operation_id in hook.pending_operation_ids
 
 
+def test_before_request_rejects_reused_operation_id():
+    hook = SplitPdfHook()
+    hook.coroutines_to_execute["reused-operation-id"] = []
+
+    with patch(
+        "unstructured_client._hooks.custom.split_pdf_hook.uuid.uuid4",
+        return_value="reused-operation-id",
+    ), pytest.raises(RuntimeError, match="Split PDF operation ID already in use"):
+        _make_hook_with_split_request(hook=hook)
+
+
 def test_before_request_logs_split_plan(caplog: pytest.LogCaptureFixture):
     caplog.set_level(logging.INFO, logger="unstructured-client")
 
@@ -1784,15 +1859,16 @@ def test_unit_allow_failed_partial_results(caplog: pytest.LogCaptureFixture):
     hook.concurrency_level[operation_id] = 3
     hook.allow_failed[operation_id] = True
     hook.cache_tmp_data_feature[operation_id] = False
-    hook.executors[operation_id] = MagicMock()
+    executor = MagicMock()
+    hook.executors[operation_id] = executor
 
     fake_future = MagicMock()
     fake_future.result.return_value = [
         (1, _httpx_json_response([{"page_number": 1}])),
         (2, _httpx_response("boom", status_code=500)),
         (3, _httpx_json_response([{"page_number": 3}])),
     ]
-    hook.executors[operation_id].submit.return_value = fake_future
+    executor.submit.return_value = fake_future
 
     elements = hook._await_elements(operation_id)
 

src/unstructured_client/_hooks/custom/split_pdf_hook.py

@@ -10,6 +10,7 @@
 import tempfile
 import threading
 import uuid
+from collections import deque
 from collections.abc import Awaitable, Iterable
 from concurrent import futures
 from functools import partial
@@ -72,6 +73,67 @@ def __init__(self, index: int, inner: BaseException):
         self.inner = inner
 
 
+class _AsyncThreadSafeBoundedSemaphore:
+    """A cancellable async gate that can be shared across event loops."""
+
+    def __init__(self, value: int) -> None:
+        if value <= 0:
+            raise ValueError("Semaphore value must be greater than zero")
+        self._initial_value = value
+        self._value = value
+        self._waiters: deque[asyncio.Future[None]] = deque()
+        self._lock = threading.Lock()
+
+    async def acquire(self) -> None:
+        loop = asyncio.get_running_loop()
+        waiter: Optional[asyncio.Future[None]] = None
+        with self._lock:
+            if self._value > 0:
+                self._value -= 1
+                return
+            waiter = loop.create_future()
+            self._waiters.append(waiter)
+
+        try:
+            await waiter
+        except asyncio.CancelledError:
+            release_transferred_slot = False
+            with self._lock:
+                try:
+                    self._waiters.remove(waiter)
+                except ValueError:
+                    release_transferred_slot = waiter.done() and not waiter.cancelled()
+            if release_transferred_slot:
+                self.release()
+            raise
+
+    def release(self) -> None:
+        while True:
+            with self._lock:
+                while self._waiters:
+                    waiter = self._waiters.popleft()
+                    if not waiter.cancelled():
+                        break
+                else:
+                    if self._value >= self._initial_value:
+                        raise ValueError("Semaphore released too many times")
+                    self._value += 1
+                    return
+
+            def _wake_waiter() -> None:
+                if waiter.cancelled():
+                    self.release()
+                else:
+                    waiter.set_result(None)
+
+            try:
+                waiter.get_loop().call_soon_threadsafe(_wake_waiter)
+                return
+            except RuntimeError:
+                # The waiting loop closed before it could receive the slot.
+                continue
+
+
 def _get_request_timeout_seconds(request: httpx.Request) -> Optional[float]:
     timeout = request.extensions.get("timeout")
     if timeout is None:
@@ -352,8 +414,8 @@ class SplitPdfHook(SDKInitHook, BeforeRequestHook, AfterSuccessHook, AfterErrorH
         max_workers=1,
         thread_name_prefix="split-pdf-setup",
     )
-    _split_pdf_setup_gate = threading.BoundedSemaphore(value=1)
-    _split_pdf_setup_poll_interval_seconds = 0.01
+    _split_pdf_setup_gate = _AsyncThreadSafeBoundedSemaphore(value=1)
+    # Thread-local flag shared by all hook instances to prove pypdfium work is under the setup lock.
     _split_pdf_setup_state = threading.local()
 
     def __init__(self) -> None:
@@ -635,6 +697,9 @@ def _before_request_unlocked(
         if split_size >= page_count and page_count == len(pdf.pages):
             return request
 
+        if operation_id in self.coroutines_to_execute:
+            raise RuntimeError(f"Split PDF operation ID already in use: {operation_id}")
+
         self.allow_failed[operation_id] = allow_failed
         self.cache_tmp_data_feature[operation_id] = cache_tmp_data_feature
         self.cache_tmp_data_dir[operation_id] = cache_tmp_data_dir
@@ -778,13 +843,17 @@ async def before_request_async(
 
     @classmethod
     async def _acquire_split_pdf_setup_slot(cls) -> None:
-        while not cls._split_pdf_setup_gate.acquire(blocking=False):
-            await asyncio.sleep(cls._split_pdf_setup_poll_interval_seconds)
+        await cls._split_pdf_setup_gate.acquire()
 
     async def _finish_cancelled_split_setup(
         self,
         setup_future: asyncio.Future[Union[httpx.Request, Exception]],
     ) -> Optional[Union[httpx.Request, Exception]]:
+        """Finish setup after caller cancellation so prepared state can be cleaned.
+
+        Non-cancellation failures mean setup failed before returning a prepared
+        request, so there is no operation ID to clear here.
+        """
         while True:
             try:
                 return await asyncio.shield(setup_future)

src/unstructured_client/basesdk.py

@@ -21,6 +21,8 @@
 
 
 class _RequestBoundCancelledError(asyncio.CancelledError):
+    """Cancellation wrapper that exposes the request to after-error hooks."""
+
     def __init__(self, request: httpx.Request, cancellation: asyncio.CancelledError):
         super().__init__(str(cancellation) or "Request cancelled")
         self.request = request
@@ -317,16 +319,25 @@ async def cleanup_cancelled_request(
             if cleanup_request is None and response is not None:
                 cleanup_request = response.request
             assert cleanup_request is not None
-            try:
-                await hooks.after_error_async(
+            cleanup_task = asyncio.create_task(
+                hooks.after_error_async(
                     AfterErrorContext(hook_ctx),
                     response,
                     _RequestBoundCancelledError(cleanup_request, cancellation),
                 )
-            except asyncio.CancelledError:
-                logger.debug("Cancellation cleanup cancelled", exc_info=True)
-            except Exception:
-                logger.debug("Cancellation cleanup failed", exc_info=True)
+            )
+            while True:
+                try:
+                    await asyncio.shield(cleanup_task)
+                    return
+                except asyncio.CancelledError:
+                    if cleanup_task.done():
+                        logger.debug("Cancellation cleanup cancelled", exc_info=True)
+                        return
+                    logger.debug("Cancellation cleanup still running after cancellation")
+                except BaseException:
+                    logger.debug("Cancellation cleanup failed", exc_info=True)
+                    return
 
         async def do():
             http_res = None