From 7b32bd8854b27ddfaf1efa67ed29392a72689b94 Mon Sep 17 00:00:00 2001
From: Devin Ivy <devinivy@gmail.com>
Date: Mon, 20 Jul 2026 10:48:28 -0700
Subject: [PATCH] =?UTF-8?q?fix:=20harden=20lazy=20refMode=20=E2=80=94=20st?=
 =?UTF-8?q?ream=20rejection=20arm,=20refCount=20coverage?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/worker-pool.ts                |  5 +++--
 test/fixtures/ref-mode-main.ts    |  2 +-
 test/fixtures/ref-mode-overlap.ts | 12 ++++++++++++
 test/pool-ref.test.ts             | 11 +++++++++--
 4 files changed, 25 insertions(+), 5 deletions(-)
 create mode 100644 test/fixtures/ref-mode-overlap.ts

diff --git a/src/worker-pool.ts b/src/worker-pool.ts
index e28f475..24c23d0 100644
--- a/src/worker-pool.ts
+++ b/src/worker-pool.ts
@@ -89,11 +89,12 @@ export function workers(sizeOrOpts?: number | WorkerOptions, opts?: WorkerOption
     inflight.add(done);
     counts.inc(index);
     refAll();
-    done.then(() => {
+    const cleanup = () => {
       inflight.delete(done);
       counts.dec(index);
       unrefAll();
-    });
+    };
+    done.then(cleanup, cleanup);
   }
 
   function terminateAll(): void {
diff --git a/test/fixtures/ref-mode-main.ts b/test/fixtures/ref-mode-main.ts
index d3181c6..56c5fbb 100644
--- a/test/fixtures/ref-mode-main.ts
+++ b/test/fixtures/ref-mode-main.ts
@@ -4,7 +4,7 @@ import { busy } from './load-balancing.ts';
 // A refMode:'lazy' pool must not keep the process alive once work is done,
 // even WITHOUT explicit disposal. If the workers stay ref'd, this process
 // hangs and the test times out.
-const run = workers(1, { refMode: 'lazy' } as any);
+const run = workers(1, { refMode: 'lazy' });
 const result = await run(busy(20));
 console.log('DONE ' + result);
 // No dispose — natural exit is the assertion.
diff --git a/test/fixtures/ref-mode-overlap.ts b/test/fixtures/ref-mode-overlap.ts
new file mode 100644
index 0000000..234673f
--- /dev/null
+++ b/test/fixtures/ref-mode-overlap.ts
@@ -0,0 +1,12 @@
+import { workers } from 'moroutine';
+import { busy } from './load-balancing.ts';
+
+// Overlapping tasks on a lazy pool: if completing the short task unref'd
+// the pool while the long one was still in flight, the process would exit
+// (code 13) before LONG completes and this would never print.
+const run = workers(2, { refMode: 'lazy' });
+const short = run(busy(20));
+const long = run(busy(120));
+await short;
+const r = await long;
+console.log('DONE ' + r);
diff --git a/test/pool-ref.test.ts b/test/pool-ref.test.ts
index 54eb545..110a3e1 100644
--- a/test/pool-ref.test.ts
+++ b/test/pool-ref.test.ts
@@ -24,11 +24,18 @@ describe('pool worker ref behavior', () => {
   });
 
   it('refMode lazy pool keeps the process alive while work is in flight', async () => {
-    // Reuses the same fixture: if workers were unref'd DURING the busy() call,
-    // the process would exit before printing DONE and stdout would be empty.
+    // Reuses the same fixture: the exec call rejects (exit 13, unsettled
+    // top-level await) if workers are unref'd during the call.
     const { stdout } = await exec(process.execPath, ['--no-warnings', join(fixturesDir, 'ref-mode-main.ts')], {
       timeout: 15000,
     });
     assert.match(stdout, /DONE done/);
   });
+
+  it("refCount keeps the pool ref'd while any overlapping task is in flight", async () => {
+    const { stdout } = await exec(process.execPath, ['--no-warnings', join(fixturesDir, 'ref-mode-overlap.ts')], {
+      timeout: 15000,
+    });
+    assert.match(stdout, /DONE done/);
+  });
 });
-- 
2.51.2