streams: invoke underlyingSource.cancel() when a type:"direct" ReadableStream is cancelled

src/jsc/bindings/webcore/streams/ReadableStreamOperations.cpp

@@ -380,6 +380,37 @@ void readableStreamError(JSGlobalObject* globalObject, JSReadableStream* stream,
     RELEASE_AND_RETURN(scope, readableStreamBYOBReaderErrorReadIntoRequests(globalObject, static_cast<JSReadableStreamBYOBReader*>(reader), error));
 }
 
+// The Bun type:"direct" [[cancelAlgorithm]]: invoke underlyingSource.cancel(reason) under
+// the stream's construction-time async context, wrapping the completion in a promise.
+static JSPromise* directUnderlyingSourceCancel(JSC::VM& vm, JSGlobalObject* globalObject, JSReadableStream* stream, JSObject* underlyingSource, JSValue reason)
+{
+    auto scope = DECLARE_THROW_SCOPE(vm);
+    if (!underlyingSource)
+        RELEASE_AND_RETURN(scope, promiseFulfilledWith(globalObject, JSC::jsUndefined()));
+    StreamAsyncContextScope asyncContextScope(globalObject, stream);
+    JSValue cancelFunction = underlyingSource->get(globalObject, builtinNames(vm).cancelPublicName());
+    RETURN_IF_EXCEPTION(scope, nullptr);
+    auto callData = JSC::getCallData(cancelFunction);
+    if (callData.type == CallData::Type::None)
+        RELEASE_AND_RETURN(scope, promiseFulfilledWith(globalObject, JSC::jsUndefined()));
+    JSValue result;
+    JSValue thrown;
+    {
+        auto catchScope = DECLARE_TOP_EXCEPTION_SCOPE(vm);
+        MarkedArgumentBuffer args;
+        args.append(reason);
+        ASSERT(!args.hasOverflowed());
+        result = JSC::call(globalObject, cancelFunction, callData, underlyingSource, args);
+        if (catchScope.exception()) [[unlikely]]
+            thrown = takeAbruptCompletion(globalObject, catchScope);
+    }
+    if (!thrown.isEmpty())
+        RELEASE_AND_RETURN(scope, promiseRejectedWith(globalObject, thrown));
+    if (result.isEmpty())
+        return nullptr;
+    RELEASE_AND_RETURN(scope, promiseResolvedWith(globalObject, result));
+}
+
 // ReadableStreamCancel(stream, reason)
 JSPromise* readableStreamCancel(JSGlobalObject* globalObject, JSReadableStream* stream, JSValue reason)
 {
@@ -412,7 +443,13 @@ JSPromise* readableStreamCancel(JSGlobalObject* globalObject, JSReadableStream*
     JSPromise* sourceCancelPromise = nullptr;
     switch (stream->m_controllerKind) {
     case ControllerKind::None:
-        sourceCancelPromise = promiseFulfilledWith(globalObject, JSC::jsUndefined());
+        if (JSObject* underlyingSource = stream->m_directUnderlyingSource.get()) {
+            stream->m_directUnderlyingSource.clear();
+            stream->m_bunMode = BunStreamMode::Default;
+            sourceCancelPromise = directUnderlyingSourceCancel(vm, globalObject, stream, underlyingSource, reason);
+        } else {
+            sourceCancelPromise = promiseFulfilledWith(globalObject, JSC::jsUndefined());
+        }
         break;
     case ControllerKind::Default:
         sourceCancelPromise = defaultControllerOf(stream)->cancelSteps(globalObject, reason);
@@ -422,19 +459,18 @@ JSPromise* readableStreamCancel(JSGlobalObject* globalObject, JSReadableStream*
         break;
     case ControllerKind::Direct: {
         auto* controller = uncheckedDowncast<WebCore::JSDirectStreamController>(stream->m_controller.get());
-        controller->onClose(globalObject, reason);
-        RETURN_IF_EXCEPTION(scope, nullptr);
         // readableStreamClose above already moved the stream out of Readable, so onClose
-        // early-returned; a direct read still pending on the controller settles as done here
-        // (a canceled read resolves with { value: undefined, done: true }).
+        // would early-return; a direct read still pending on the controller settles as done
+        // here (a canceled read resolves with { value: undefined, done: true }).
+        controller->m_closed = true;
         if (auto* pendingRead = controller->m_pendingRead.get()) {
             controller->m_pendingRead.clear();
             JSObject* doneResult = createIteratorResultObject(globalObject, jsUndefined(), true);
             RETURN_IF_EXCEPTION(scope, nullptr);
             pendingRead->fulfill(vm, doneResult);
             RETURN_IF_EXCEPTION(scope, nullptr);
         }
-        sourceCancelPromise = promiseFulfilledWith(globalObject, JSC::jsUndefined());
+        sourceCancelPromise = directUnderlyingSourceCancel(vm, globalObject, stream, controller->m_underlyingSource.get(), reason);
         break;
     }
     case ControllerKind::NativeSink: {

test/js/node/async_hooks/AsyncLocalStorage.test.ts

@@ -601,8 +601,7 @@ describe("async context passes through", () => {
     await promise;
     expect(value).toBe("value");
   });
-  // blocked by a bug with .cancel
-  test.todo("readable stream direct .cancel", async () => {
+  test("readable stream direct .cancel", async () => {
     const s = new AsyncLocalStorage<string>();
     let stream!: ReadableStream;
     let value: string | undefined;
@@ -617,7 +616,6 @@ describe("async context passes through", () => {
           controller.write("hello");
         },
         cancel(reason) {
-          console.log("1");
           value2 = s.getStore();
           resolve();
         },
@@ -627,7 +625,6 @@ describe("async context passes through", () => {
     const reader = stream.getReader();
     await reader.read();
     await reader.cancel();
-    await stream.cancel();
     await promise;
     expect(value).toBe("value");
     expect(value2).toBe("value");

test/js/web/streams/streams.test.js

@@ -684,6 +684,73 @@ describe("multi-chunk consumers produce exactly the concatenated bytes", () => {
     expect(result.done).toBe(true);
   });
 
+  it("canceling a direct stream invokes the source's cancel() callback", async () => {
+    // reader.cancel(): controller materialized
+    {
+      let reason;
+      const rs = new ReadableStream({
+        type: "direct",
+        pull(c) {
+          c.write("hello");
+          c.flush();
+        },
+        cancel(r) {
+          reason = r;
+        },
+      });
+      const reader = rs.getReader();
+      await reader.read();
+      await reader.cancel("bye");
+      expect(reason).toBe("bye");
+    }
+    // stream.cancel() before any reader: controller not yet materialized
+    {
+      let reason;
+      const rs = new ReadableStream({
+        type: "direct",
+        pull() {},
+        cancel(r) {
+          reason = r;
+        },
+      });
+      await rs.cancel("early");
+      expect(reason).toBe("early");
+    }
+    // the cancel promise chains onto the source's returned promise
+    {
+      const order = [];
+      const rs = new ReadableStream({
+        type: "direct",
+        pull() {},
+        async cancel() {
+          await Promise.resolve();
+          order.push("source");
+        },
+      });
+      await rs.cancel();
+      order.push("awaited");
+      expect(order).toEqual(["source", "awaited"]);
+    }
+  });
+
+  it("a direct stream's controller.write() throws after reader.cancel()", async () => {
+    let capturedController;
+    const rs = new ReadableStream({
+      type: "direct",
+      pull(c) {
+        capturedController = c;
+        c.write("a");
+        c.flush();
+      },
+    });
+    const reader = rs.getReader();
+    await reader.read();
+    await reader.cancel();
+    expect(() => capturedController.write("b")).toThrow(
+      expect.objectContaining({ name: "TypeError", message: "ReadableStreamDirectController is now closed" }),
+    );
+  });
+
   it("releasing a direct stream's reader during an async pull does not crash close", async () => {
     const rs = new ReadableStream({
       type: "direct",