fix(ipc): unblock thread executor aclose() when initialize is cancelled

[BlueX888]

fix(ipc): unblock thread executor aclose() when initialize is cancelled

What breaks

ThreadJobExecutor.aclose() hangs forever when initialize() never completed — either because it
was cancelled mid-warm-up, or because it was never called at all. The worker process then never
finishes shutting down.

This is not an exotic path: ThreadJobExecutor is the default job executor on Windows
(worker.py falls back to JobExecutorType.THREAD there) and is forced by the devmode CLI
(cli.py). During worker shutdown, ProcPool._main_task teardown cancels the proc-spawn tasks; a
spawn task sitting in await proc.initialize() gets a CancelledError, its
except asyncio.CancelledError: pass branch runs, and then await proc.aclose() never returns —
so ProcPool.aclose() / Worker.aclose() hang indefinitely. It is also reproducible on the
public API by cancelling an in-flight initialize() and then closing the executor.

Root cause

• ThreadJobExecutor.initialize() resolves self._initialize_fut only under
  except asyncio.TimeoutError, except Exception and the success else branch. A
  CancelledError is a BaseException, so cancelling an in-flight initialize() propagates
  without ever resolving the future.
• _main_task() awaits self._initialize_fut before creating _monitor_task, which is the
  only code that reads the IPC channel and resolves _shutdown_ack_fut / _shutting_down_fut
  (the runner thread's ack of the queued ShutdownRequest sits unread in the socket buffer).
• aclose() then waits bounded on _shutdown_ack_fut (logging "job did not ack shutdown in
  time") and blocks forever on await self._shutting_down_fut — even though the runner thread
  exits cleanly on its own.

The process-based sibling SupervisedProc already guards against exactly this: aclose() checks
if not self._initialize_fut.done(): await self.kill(), and kill() resolves the future
("unblock the supervise task waiting on this future"). ThreadJobExecutor was simply missing the
equivalent.

The fix

Mirror the sibling guard, adapted for a runner that cannot be killed: when the initialize future
is unresolved at aclose() time, fail it so _main_task() proceeds, starts the monitor task
(the sole channel reader), and the graceful shutdown sequence can complete — the runner acks the
queued ShutdownRequest, sends ShuttingDown, exits, and aclose() returns.

if not self._initialize_fut.done():
    # never initialized (a cancelled initialize() also leaves this future
    # unresolved, its CancelledError isn't caught): unblock the main task
    # waiting on this future so it starts the monitor task, the sole reader
    # of the channel, and the runner can ack the queued shutdown request
    self._initialize_fut.set_exception(
        RuntimeError("runner closed before initialization completed")
    )

This also covers a runner that was started but never initialized: previously the same unbounded
wait applied there (the main task never started the reader, so the shutdown ack could never be
observed).

Test evidence

New regression test test_thread_aclose_after_initialize_cancelled in tests/test_ipc.py. It
starts a thread runner whose warm-up is held on an event, cancels the in-flight initialize()
task (so the cancel cannot race with a warm-up that already finished), releases the warm-up, then
calls aclose() under an outer timeout and asserts a clean SUCCESS shutdown.

• Before the fix: aclose() never returns — the test fails with TimeoutError after waiting 5s
  (the runner thread had exited on its own within a fraction of a second), plus a leaked
  ThreadJobExecutor._main_task at teardown.
• After the fix: the test passes in well under a second, and the full tests/test_ipc.py suite
  passes (18 passed), as do the related test_ipc_preload.py, test_supervised_proc_memory.py,
  test_drain_timeout.py and test_worker_connection.py.

[BlueX888]

Both findings reproduced and are fixed in 6b21603 (on top of c08335d), with deterministic regression tests for each:

• Concurrent close breaks initialization reader: initialize() now runs its response read as a tracked task and guards every _initialize_fut write with done() checks; aclose() cancel-and-waits the in-flight read before failing the future, so _monitor_task is never a second duplex reader and the double-set InvalidStateError is impossible. Empirically confirmed: pre-fix the interleaving produced monitor RuntimeError → aclose() raise → initialize() InvalidStateError; post-fix aclose() returns, initialize() aborts with CancelledError, status SUCCESS.
• Uninitialized shutdown raises in runner thread: _ProcClient.initialize() now accepts a leading ShutdownRequest (logs, closes the duplex, returns False) and thread_main exits cleanly without starting the job loop; the process-side proc_main entrypoints honor the same contract, and _main_task no longer records SUCCESS over a crashed runner.

New tests test_thread_aclose_during_pending_initialize and test_thread_aclose_without_initialize (both red at base and red at the previous revision with the exact reported failure modes; green at 6b21603). Full surrounding suite at the new head: 50 passed, 1 platform skip.

[BlueX888]

Follow-up on the two residual interrupt windows in 45e9056 — both reproduced, and a third one found next to them; all three are closed in d9ab64c3.

The windows. The 1-second grace in aclose() only covers a response frame that finishes in time. Deterministic reproductions (a real ThreadJobExecutor whose runner thread is stalled between the two sendalls of the InitializeResponse):

1. Payload slower than the grace — aclose() waits 1s, then falls back to cancel_and_wait(init_read_task) while the read is parked mid-frame. The pending payload is then read by the monitor task as a new frame header.
2. Externally cancelled initialize() — its finally cancelled the read unconditionally, mid-frame included; the monitor that aclose() starts afterwards inherits the desynchronized stream.
3. initialize() timeout — asyncio.wait_for(read_task, ...) cancels the read it wraps when the timeout fires, mid-frame included (this one predates the PR; upstream's original wait_for around the bare arecv_message had the same behavior).

In all three, the outcome was the same failure the first revision fixed for the sub-second case: the monitor misparses the leftover payload (struct.error: unpack requires a buffer of 4 bytes), the error escapes _monitor_task → _main_task → aclose(), the executor never reaches a terminal status, and the ShutdownRequestAck/ShuttingDown frames are lost. Verified at 45e9056 with system-level probes (real ThreadJobExecutor, instrumented runner send): aclose() raised the decoding error in every window.

The fix — close, don't cancel. Retaining the full read would reintroduce the original hang for a runner that never finishes the frame (nothing bounds the wait). Closing the duplex instead is the bounded interrupt: the read ends with DuplexClosed (frame-aligned), and the monitor task that takes the channel over immediately sees a closed channel. Concretely:

• initialize() shields the read from its wait_for timeout (the timeout cancels the shield, not the read), and its cleanup closes the duplex when the read has not finished instead of cancelling it — a bare cancel could land between a frame header and its payload. The close comes first so a read parked on the socket read is woken by the teardown; the cancel_and_wait afterwards only cleans up.
• aclose() does the same past its one-second grace: close the duplex, then clean up the read. Within the grace nothing changes — a frame that finishes in time keeps the graceful shutdown path.
• The monitor task treats a frame that fails to decode as a closed channel and breaks (resolving the shutdown futures, as its finally already does): bytes can already sit in the reader buffer when a read is interrupted, and no teardown can scrub those — so the belt-and-suspenders guarantees the close never fails with a decoding error even then. The decode-failure set covers what a garbage frame can raise from the reader (struct.error, ValueError, KeyError for an unknown message id, protobuf DecodeError).
• The thread runner exits cleanly when the channel closes during its warm-up: thread_main now mirrors proc_main's guard, so the abrupt close surfaces as a debug log instead of an unhandled thread exception.

One existing test's expectation changed with this: after aclose() aborts an in-flight initialize(), the coroutine now ends with DuplexClosed (the channel was closed under its read) instead of CancelledError — both are accepted (test_thread_aclose_during_pending_initialize).

Tests. Three regression tests added (tests/test_ipc.py), each holding the response frame mid-flight through the exact window: test_thread_aclose_slow_inflight_response_frame (payload released past the grace), test_thread_initialize_cancelled_midframe (external cancel), test_thread_initialize_timeout_midframe (short initialize_timeout). All three fail with the struct.error escape at 45e9056 (and the two interrupt tests also fail on current main); all pass at d9ab64c3. Full local verification: test_ipc.py 24 passed, and with the neighboring suites (test_ipc_preload.py, test_supervised_proc_memory.py, test_drain_timeout.py, test_worker_connection.py) 54 passed / 1 skipped; ruff clean.

[devin-ai-integration]

Devin Review found 1 new potential issue.

[devin-ai-integration]

🟡 Warm-up failures lose their traceback

When warm-up raises, thread_main suppresses its exception and exits without reporting the failure. ThreadJobExecutor.initialize ignores the response error and marks the exited runner ready, so jobs fail without a warm-up traceback.

Learn more

The runner's warm-up callback runs inside _ProcClient.initialize. If it raises, the client sends InitializeResponse(error=str(e)) and re-raises. The new handler catches that failure as though the channel had closed. ThreadJobExecutor.initialize does not inspect init_res.error, so it can announce success while the runner exits. Suppressing the exception removes the traceback that previously identified the warm-up failure.

Example: A warm-up callback raises RuntimeError("model load failed"). The client sends that text to the parent and exits through this handler. The parent marks the runner ready; a job can then be dispatched to the stopped thread, with no traceback for the model-load failure.

Recommended fix: Distinguish a closed channel from failures in the warm-up callback. Check InitializeResponse.error in ThreadJobExecutor.initialize() before resolving _initialize_fut, and preserve diagnostics for unexpected failures in thread_main.

---

Was this helpful? React with 👍 or 👎 to provide feedback.