fix: clean up the async conversion's output file when cancelled

Offloading the JSON-to-NDJSON conversion to a thread last commit removed the event
loop block but introduced a leak: a thread cannot be cancelled, so on cancellation
the conversion still runs to completion and creates its file, while the awaiting
coroutine has already raised CancelledError and discarded the path. Nothing was left
that could delete it. Reproduced -- cancel mid-conversion and an
unst_elements_*.ndjson survives.

The conversion is now awaited through asyncio.shield, which keeps a handle on the
thread's result after the caller stops waiting, and a done callback unlinks the
finished file. A conversion that raised needs no callback, since it already removes
its own partial file.

This is the same shape as the split hook's cancellation race: work that outlives the
operation that requested it, publishing a path nobody will read. It did not exist on
the inline version, which had no await point to cancel at.

Regression test asserts against the path production actually created, so it cannot
pass by observing that no file was ever made. Confirmed to fail against the plain
awaited to_thread.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: badGarnet

_test_unstructured_client/unit/test_ndjson_elements_file.py

@@ -9,11 +9,13 @@
   - leave no temp files behind other than the combined file the caller owns
 """
 
+import asyncio
 import errno
 import json
 import os
 import tempfile
 import threading
+import time
 from pathlib import Path
 from unittest import mock
 
@@ -752,3 +754,41 @@ def _spy(http_res):
         assert _read_ndjson(res.elements_file) == elements
     finally:
         Path(res.elements_file).unlink(missing_ok=True)
+
+
+@pytest.mark.asyncio
+async def test_async_conversion_cleans_up_when_cancelled(tmp_path, monkeypatch):
+    """A cancelled conversion must not orphan the file its thread went on to create.
+
+    Regression guard for the cost of the offload: a thread cannot be cancelled, so the
+    conversion runs to completion regardless, and a plain `await asyncio.to_thread(...)`
+    discards the path it returned -- leaving nothing that could delete the file.
+    """
+    monkeypatch.setattr(tempfile, "tempdir", str(tmp_path))
+    created = _record_created_paths(monkeypatch, general, "_new_elements_file")
+    real_convert = general._json_body_to_elements_file
+
+    def _slow_convert(http_res):
+        time.sleep(0.3)
+        return real_convert(http_res)
+
+    monkeypatch.setattr(general, "_json_body_to_elements_file", _slow_convert)
+    response = httpx.Response(
+        200, headers={"Content-Type": "application/json"}, json=[{"type": "Table"}]
+    )
+
+    task = asyncio.ensure_future(general._json_body_to_elements_file_async(response))
+    await asyncio.sleep(0.05)
+    task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await task
+
+    # Let the shielded thread finish and its cleanup callback run.
+    for _ in range(200):
+        await asyncio.sleep(0.02)
+        if created and not os.path.exists(created[0]):
+            break
+
+    # Non-vacuous: a file really was created, and it is now gone.
+    assert len(created) == 1
+    assert not os.path.exists(created[0])

src/unstructured_client/general.py

@@ -83,6 +83,33 @@ def _json_body_to_elements_file(http_res: httpx.Response) -> str:
     return out.name
 
 
+async def _json_body_to_elements_file_async(http_res: httpx.Response) -> str:
+    """Run the conversion off the event loop, cleaning up if the caller gives up.
+
+    The work is offloaded so it does not block the loop, but a thread cannot be cancelled:
+    on cancellation the conversion still runs to completion and creates its file, and a
+    plain `await asyncio.to_thread(...)` discards the returned path, so nothing is left
+    that could delete it. Shielding keeps a handle on the thread's result so the finished
+    file can be removed once the caller has stopped waiting for it.
+    """
+    task = asyncio.ensure_future(asyncio.to_thread(_json_body_to_elements_file, http_res))
+    try:
+        return await asyncio.shield(task)
+    except BaseException:
+        task.add_done_callback(_discard_abandoned_elements_file)
+        raise
+
+
+def _discard_abandoned_elements_file(task: "asyncio.Future[str]") -> None:
+    """Delete the file a conversion produced after its caller stopped waiting."""
+    if task.cancelled():
+        return
+    if task.exception() is not None:
+        # The conversion raised, and it already removed its own partial file.
+        return
+    _discard_elements_file(task.result())
+
+
 def _discard_elements_file(path: str) -> None:
     """Remove a partially written elements file, ignoring a failure to do so.
 
@@ -383,9 +410,7 @@ async def partition_async(
         if utils.match_response(http_res, "200", "application/json"):
             if _ndjson_requested(req):
                 return operations.PartitionResponse(
-                    elements_file=await asyncio.to_thread(
-                        _json_body_to_elements_file, http_res
-                    ),
+                    elements_file=await _json_body_to_elements_file_async(http_res),
                     status_code=http_res.status_code,
                     content_type=http_res.headers.get("Content-Type") or "",
                     raw_response=http_res,