More fixes

softwaremill · Feb 3, 2025 · 24b8ceb · 24b8ceb
1 parent 58420f1
commit 24b8ceb
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 3 deletions.
diff --git a/effects/zio/src/main/scalajvm/sttp/client4/httpclient/zio/HttpClientZioBackend.scala b/effects/zio/src/main/scalajvm/sttp/client4/httpclient/zio/HttpClientZioBackend.scala
@@ -115,7 +115,7 @@ class HttpClientZioBackend private (
       b: ZioStreams.BinaryStream,
       callback: () => Unit
   ): ZioStreams.BinaryStream =
-    b.ensuring(ZIO.attempt(callback()).orDie)
+    b.ensuringWith(exit => if (exit.isSuccess) ZIO.attempt(callback()).orDie else ZIO.unit)
 }
 
 object HttpClientZioBackend {

diff --git a/effects/zio1/src/main/scalajvm/sttp/client4/httpclient/zio/HttpClientZioBackend.scala b/effects/zio1/src/main/scalajvm/sttp/client4/httpclient/zio/HttpClientZioBackend.scala
@@ -122,8 +122,7 @@ class HttpClientZioBackend private (
   override protected def addOnEndCallbackToBody(
       b: ZioStreams.BinaryStream,
       callback: () => Unit
-  ): ZioStreams.BinaryStream =
-    b.ensuring(ZIO.effect(callback()).orDie)
+  ): ZioStreams.BinaryStream = b.++(ZStream.execute(ZIO.effect(callback())))
 }
 
 object HttpClientZioBackend {