Merge pull request #1549 from corbitsdev/cl-10975-stop-a-finished-session-operation-from-rejecting-the-next

fix(tui): stop a finished session op from rejecting the next one

Author: TheGreatAxios

src/tui/delivery-queue.ts

@@ -461,27 +461,55 @@ export function createSessionOperationQueue(): SessionOperationQueue {
   // the await.
   let nextId = 0;
   let activeId = 0;
+  // Body promise of the op in `activeId`. Unset during that op's first
+  // synchronous turn, before `run` returns.
+  let activeResult: Promise<void> | undefined;
 
   // True only when this continuation is the op that is running right now.
   // A detached continuation still carries its parent's id; that must not
   // match a later op's activeId.
   const inThisOp = (): boolean => runningSessionOp.getStore() === activeId;
 
+  // The body promise can be fulfilled while `activeId` is still set: `finally`
+  // clears the id on a later microtask. A continuation in that gap is not the
+  // op waiting on itself.
+  const bodyStillRunning = (): boolean => {
+    if (activeResult === undefined) return true;
+    try {
+      return Bun.peek(activeResult) === activeResult;
+    } catch {
+      return false;
+    }
+  };
+
   const execute = (op: () => Promise<void>): Promise<void> => {
     const id = ++nextId;
     activeId = id;
     let result: Promise<void>;
     try {
       result = runningSessionOp.run(id, op);
     } catch (err) {
-      if (activeId === id) activeId = 0;
+      if (activeId === id) {
+        activeId = 0;
+        activeResult = undefined;
+      }
       throw err;
     }
+    activeResult = result;
     return Promise.resolve(result).finally(() => {
-      if (activeId === id) activeId = 0;
+      if (activeId === id) {
+        activeId = 0;
+        activeResult = undefined;
+      }
     });
   };
 
+  // `observed === activeResult` is the adoption of a thenable returned from
+  // the op. That `.then` runs after AsyncLocalStorage.run returns.
+  const isSelfWait = (during: number, observed: object): boolean =>
+    activeId === during &&
+    (observed === activeResult || (inThisOp() && bodyStillRunning()));
+
   const enqueue = (op: () => Promise<void>): Promise<void> => {
     // `tail.then(op, op)` — a failed op must not wedge the chain. The
     // returned promise is not that tail: a self-wait rejection has to
@@ -493,7 +521,9 @@ export function createSessionOperationQueue(): SessionOperationQueue {
     tail = scheduled;
     if (!inThisOp()) return scheduled;
     const enqueuedDuring = activeId;
-    return observeUnlessSelfWait(scheduled, () => activeId === enqueuedDuring);
+    return observeUnlessSelfWait(scheduled, (observed) =>
+      isSelfWait(enqueuedDuring, observed),
+    );
   };
 
   const awaitTail = (): Promise<void> => {
@@ -503,7 +533,9 @@ export function createSessionOperationQueue(): SessionOperationQueue {
     // deadlocks the same way as awaiting enqueue.
     if (!inThisOp()) return drained;
     const awaitedDuring = activeId;
-    return observeUnlessSelfWait(drained, () => activeId === awaitedDuring);
+    return observeUnlessSelfWait(drained, (observed) =>
+      isSelfWait(awaitedDuring, observed),
+    );
   };
 
   return { enqueue, awaitTail };
@@ -517,7 +549,7 @@ export function createSessionOperationQueue(): SessionOperationQueue {
  */
 function observeUnlessSelfWait(
   forward: Promise<void>,
-  stillRunning: () => boolean,
+  stillRunning: (observed: object) => boolean,
 ): Promise<void> {
   let reported: SessionOperationSelfWaitError | undefined;
   let guarded = false;
@@ -555,7 +587,7 @@ function observeUnlessSelfWait(
     ): Promise<undefined>;
   } = {
     then(onFulfilled, onRejected) {
-      if (stillRunning()) {
+      if (stillRunning(thenable)) {
         guardUnobserved();
         return reject().then(onFulfilled, onRejected);
       }

src/tui/session-operation-queue.test.ts

@@ -437,3 +437,128 @@ test("a self-wait does not leave the queued op's rejection unhandled", async ()
     process.off("unhandledRejection", onRejection);
   }
 });
+
+test("a microtask detached from a finished op body is not a self-wait", async () => {
+  const log: string[] = [];
+  const { enqueue, awaitTail } = createSessionOperationQueue();
+  let detached: Promise<void> = Promise.resolve();
+
+  const outer = enqueue(async () => {
+    detached = (async () => {
+      await Promise.resolve();
+      await enqueue(async () => {
+        log.push("inner");
+      });
+      log.push("detached:ok");
+    })();
+    log.push("outer:end");
+  });
+
+  await outer;
+  await detached;
+  await awaitTail();
+  expect(log).toEqual(["outer:end", "inner", "detached:ok"]);
+});
+
+test("queueMicrotask after the op body returns is not a self-wait", async () => {
+  const log: string[] = [];
+  const { enqueue, awaitTail } = createSessionOperationQueue();
+  let done: Promise<void> = Promise.resolve();
+
+  const outer = enqueue(async () => {
+    done = new Promise((resolve, reject) => {
+      queueMicrotask(() => {
+        void enqueue(async () => {
+          log.push("inner");
+        }).then(resolve, reject);
+      });
+    });
+    log.push("outer:end");
+  });
+
+  await outer;
+  await done;
+  await awaitTail();
+  expect(log).toEqual(["outer:end", "inner"]);
+});
+
+test("an outside waiter is not rejected while the enqueuing op is blocked", async () => {
+  const log: string[] = [];
+  const { enqueue, awaitTail } = createSessionOperationQueue();
+  let release: () => void = () => undefined;
+  const gate = new Promise<void>((resolve) => {
+    release = resolve;
+  });
+  let inner: Promise<void> | undefined;
+
+  const outer = enqueue(async () => {
+    inner = enqueue(async () => {
+      log.push("inner");
+    });
+    await gate;
+    log.push("outer:end");
+  });
+
+  await Promise.resolve();
+  await Promise.resolve();
+  if (inner === undefined) throw new Error("inner missing");
+  let observed = "pending";
+  const waiting = inner.then(
+    () => {
+      observed = "resolved";
+    },
+    (err: unknown) => {
+      observed =
+        err instanceof SessionOperationSelfWaitError ? "self-wait" : "error";
+    },
+  );
+  await Promise.resolve();
+  expect(observed).toBe("pending");
+  release();
+  await outer;
+  await waiting;
+  await awaitTail();
+  expect(observed).toBe("resolved");
+  expect(log).toEqual(["outer:end", "inner"]);
+});
+
+async function settleOrTimeout(promise: Promise<unknown>): Promise<string> {
+  return Promise.race([
+    promise.then(
+      () => "resolved",
+      (err: unknown) =>
+        err instanceof SessionOperationSelfWaitError ? "self-wait" : "error",
+    ),
+    new Promise<string>((resolve) => setTimeout(() => resolve("timeout"), 200)),
+  ]);
+}
+
+test("a non-async callback that returns enqueue rejects and the queue drains", async () => {
+  const log: string[] = [];
+  const { enqueue, awaitTail } = createSessionOperationQueue();
+
+  const outer = enqueue(() => {
+    log.push("outer");
+    return enqueue(async () => {
+      log.push("inner");
+    });
+  });
+
+  expect(await settleOrTimeout(outer)).toBe("self-wait");
+  expect(await settleOrTimeout(awaitTail())).toBe("resolved");
+  expect(log).toEqual(["outer", "inner"]);
+});
+
+test("a non-async callback that returns awaitTail rejects and the queue drains", async () => {
+  const log: string[] = [];
+  const { enqueue, awaitTail } = createSessionOperationQueue();
+
+  const outer = enqueue(() => {
+    log.push("outer");
+    return awaitTail();
+  });
+
+  expect(await settleOrTimeout(outer)).toBe("self-wait");
+  expect(await settleOrTimeout(awaitTail())).toBe("resolved");
+  expect(log).toEqual(["outer"]);
+});