feat: contextual fixes

Author: algebraic-dev

src/Std/Internal/Http/Data/Body.lean

@@ -6,6 +6,7 @@ Authors: Sofia Rodrigues
 module
 
 prelude
+public import Std.Internal.Async.Context
 public import Std.Internal.Http.Data.Body.Length
 public import Std.Internal.Http.Data.Body.ByteStream
 
@@ -74,21 +75,19 @@ instance : Coe ByteArray Body where
 instance : Coe Body.ByteStream Body where
   coe := .stream
 
-/--
-Iterate over the body content in chunks, processing each ByteArray chunk with the given step function.
--/
-@[inline]
-protected partial def forIn
-  {β : Type} (body : Body) (acc : β)
-  (step : Chunk → β → Async (ForInStep β)) :
-  Async β := do
+instance : ForIn Async Body Chunk where
+  forIn body acc step :=
     match body with
     | .zero => pure acc
     | .bytes data => return (← step (Chunk.mk data #[]) acc).value
     | .stream stream' => ByteStream.forIn stream' acc step
 
-instance : ForIn Async Body Chunk where
-  forIn := Body.forIn
+instance : ForIn ContextAsync Body Chunk where
+  forIn body acc step :=
+    match body with
+    | .zero => pure acc
+    | .bytes data => return (← step (Chunk.mk data #[]) acc).value
+    | .stream stream' => ByteStream.forIn' stream' acc step
 
 /--
 Collect all data from the body into a single `ByteArray`. This reads the entire body content into memory,

src/Std/Internal/Http/Data/Body/ByteStream.lean

@@ -161,6 +161,30 @@ protected partial def forIn
 instance : ForIn Async ByteStream Chunk where
   forIn := Std.Http.Body.ByteStream.forIn
 
+/--
+Iterate over the body content in chunks, processing each chunk with the given step function.
+-/
+@[inline]
+protected partial def forIn'
+    {β : Type} (stream : ByteStream) (acc : β)
+    (step : Chunk → β → ContextAsync (ForInStep β)) : ContextAsync β := do
+
+  let rec @[specialize] loop (stream : ByteStream) (acc : β) : ContextAsync β := do
+    if let some chunk ← stream.recv then
+      match ← step chunk acc with
+      | .done res => return res
+      | .yield res => loop stream res
+    else
+      return acc
+
+  loop stream acc
+
+instance : ForIn Async ByteStream Chunk where
+  forIn := Std.Http.Body.ByteStream.forIn
+
+instance : ForIn ContextAsync ByteStream Chunk where
+  forIn := Std.Http.Body.ByteStream.forIn'
+
 instance : IO.AsyncRead ByteStream (Option Chunk) where
   read stream := stream.recv
 

src/Std/Internal/Http/Data/Response.lean

@@ -136,7 +136,7 @@ def body (builder : Builder) (body : t) : Response t :=
 /--
 Builds and returns the final HTTP Response with a stream builder
 -/
-def stream (builder : Builder) (body : Body.ByteStream → Async Unit) : Async (Response Body) := do
+def stream (builder : Builder) (body : Body.ByteStream → ContextAsync Unit) : ContextAsync (Response Body) := do
   let stream ← Body.ByteStream.empty
   background (body stream)
   return { head := builder.head, body := stream }

tests/lean/run/async_http_connection.lean

@@ -1,7 +1,8 @@
 import Std.Internal.Http
 import Std.Internal.Async
+import Std.Internal.Async.Timer
 
-open Std.Internal.IO.Async
+open Std.Internal.IO Async
 open Std Http
 
 structure TestCase where
@@ -10,7 +11,7 @@ structure TestCase where
   /-- The HTTP request to send -/
   request : Request (Array Chunk)
   /-- Handler function to process the request -/
-  handler : Request Body → Async (Response Body)
+  handler : Request Body → ContextAsync (Response Body)
   /-- Expected response string -/
   expected : String
   /-- Whether to use chunked encoding -/
@@ -27,20 +28,21 @@ def toByteArray (req : Request (Array Chunk)) (chunked := false) : IO ByteArray
 
 /-- Send multiple requests through a mock connection and return the response data. -/
 def sendRequests (client : Mock.Client) (server : Mock.Server) (reqs : Array (Request (Array Chunk)))
-    (onRequest : Request Body → Async (Response Body))
+    (onRequest : Request Body → ContextAsync (Response Body))
     (chunked : Bool := false) : IO ByteArray := Async.block do
   let mut data := .empty
   for req in reqs do data := data ++ (← toByteArray req chunked)
 
   client.send data
   Std.Http.Server.serveConnection server onRequest (config := { lingeringTimeout := 3000 })
+    |>.run (← Context.new)
 
   let res ← client.recv?
   pure <| res.getD .empty
 
 /-- Run a single test case, comparing actual response against expected response. -/
 def runTest (name : String) (client : Mock.Client) (server : Mock.Server) (req : Request (Array Chunk))
-    (handler : Request Body → Async (Response Body)) (expected : String) (chunked : Bool := false) :
+    (handler : Request Body → ContextAsync (Response Body)) (expected : String) (chunked : Bool := false) :
     IO Unit := do
   let response ← sendRequests client server #[req] handler chunked
   let responseData := String.fromUTF8! response
@@ -270,10 +272,10 @@ def hasUri (req : Request Body) (uri : String) : Bool :=
     |>.body #[]
 
   handler := fun _ => do
-    let largeBody := String.mk (List.replicate 1000 'X')
+    let largeBody := String.ofList (List.replicate 1000 'X')
     return Response.ok largeBody
 
-  expected := s!"HTTP/1.1 200 OK\x0d\nContent-Length: 1000\x0d\nConnection: close\x0d\nServer: LeanHTTP/1.1\x0d\n\x0d\n{String.mk (List.replicate 1000 'X')}"
+  expected := s!"HTTP/1.1 200 OK\x0d\nContent-Length: 1000\x0d\nConnection: close\x0d\nServer: LeanHTTP/1.1\x0d\n\x0d\n{String.ofList (List.replicate 1000 'X')}"
 }
 
 #eval runTestCase {

tests/lean/run/async_http_connection_limits.lean

@@ -10,7 +10,7 @@ structure TestCase where
   /-- The HTTP request to send -/
   request : Request (Array Chunk)
   /-- Handler function to process the request -/
-  handler : Request Body → Async (Response Body)
+  handler : Request Body → ContextAsync (Response Body)
   /-- Expected response string -/
   expected : String
   /-- Whether to use chunked encoding -/
@@ -27,20 +27,21 @@ def toByteArray (req : Request (Array Chunk)) (chunked := false) : IO ByteArray
 
 /-- Send multiple requests through a mock connection and return the response data. -/
 def sendRequests (client : Mock.Client) (server : Mock.Server) (reqs : Array (Request (Array Chunk)))
-    (onRequest : Request Body → Async (Response Body))
+    (onRequest : Request Body → ContextAsync (Response Body))
     (chunked : Bool := false) : IO ByteArray := Async.block do
   let mut data := .empty
   for req in reqs do data := data ++ (← toByteArray req chunked)
 
   client.send data
   Std.Http.Server.serveConnection server onRequest (config := { lingeringTimeout := 3000 })
+  |>.run (← Context.new)
 
   let res ← client.recv?
   pure <| res.getD .empty
 
 /-- Run a single test case, comparing actual response against expected response. -/
 def runTest (name : String) (client : Mock.Client) (server : Mock.Server) (req : Request (Array Chunk))
-    (handler : Request Body → Async (Response Body)) (expected : String) (chunked : Bool := false) :
+    (handler : Request Body → ContextAsync (Response Body)) (expected : String) (chunked : Bool := false) :
     IO Unit := do
   let response ← sendRequests client server #[req] handler chunked
   let responseData := String.fromUTF8! response
@@ -116,7 +117,7 @@ def hasUri (req : Request Body) (uri : String) : Bool :=
   request :=
     Request.new
     |>.method .get
-    |>.uri (.originForm (.mk #[String.mk (List.replicate 2000 'a')] true) none none)
+    |>.uri (.originForm (.mk #[String.ofList (List.replicate 2000 'a')] true) none none)
     |>.header! "Host" "api.example.com"
     |>.header! "Connection" "close"
     |>.body #[]
@@ -134,8 +135,8 @@ def hasUri (req : Request Body) (uri : String) : Bool :=
   request :=
     Request.new
     |>.method .get
-    |>.uri (.originForm (.mk #[String.mk (List.replicate 200 'a')] true) none none)
-    |>.header! "Host" (String.mk (List.replicate 8230 'a'))
+    |>.uri (.originForm (.mk #[String.ofList (List.replicate 200 'a')] true) none none)
+    |>.header! "Host" (String.ofList (List.replicate 8230 'a'))
     |>.header! "Connection" "close"
     |>.body #[]
 
@@ -173,7 +174,7 @@ def hasUri (req : Request Body) (uri : String) : Bool :=
     |>.method .get
     |>.uri! "/api/test"
     |>.header! "Host" "api.example.com"
-    |>.header! "X-Long-Value" (String.mk (List.replicate 9000 'x'))
+    |>.header! "X-Long-Value" (String.ofList (List.replicate 9000 'x'))
       |>.header! "Connection" "close"
     |>.body #[]
 
@@ -194,7 +195,7 @@ def hasUri (req : Request Body) (uri : String) : Bool :=
       |>.header! "Connection" "close"
 
     for i in [0:200] do
-      req := req |>.header! s!"X-Header-{i}" (String.mk (List.replicate 200 'a'))
+      req := req |>.header! s!"X-Header-{i}" (String.ofList (List.replicate 200 'a'))
     return req |>.body #[]
 
   handler := fun _ => do

tests/lean/run/async_http_connection_streaming.lean

@@ -10,7 +10,7 @@ structure TestCase where
   /-- The HTTP request to send -/
   request : Request (Array Chunk)
   /-- Handler function to process the request -/
-  handler : Request Body → Async (Response Body)
+  handler : Request Body → ContextAsync (Response Body)
   /-- Expected response string -/
   expected : String
   /-- Whether to use chunked encoding -/
@@ -27,20 +27,21 @@ def toByteArray (req : Request (Array Chunk)) (chunked := false) : IO ByteArray
 
 /-- Send multiple requests through a mock connection and return the response data. -/
 def sendRequests (client : Mock.Client) (server : Mock.Server) (reqs : Array (Request (Array Chunk)))
-    (onRequest : Request Body → Async (Response Body))
+    (onRequest : Request Body → ContextAsync (Response Body))
     (chunked : Bool := false) : IO ByteArray := Async.block do
   let mut data := .empty
   for req in reqs do data := data ++ (← toByteArray req chunked)
 
   client.send data
   Std.Http.Server.serveConnection server onRequest (config := { lingeringTimeout := 3000 })
+  |>.run (← Context.new)
 
   let res ← client.recv?
   pure <| res.getD .empty
 
 /-- Run a single test case, comparing actual response against expected response. -/
 def runTest (name : String) (client : Mock.Client) (server : Mock.Server) (req : Request (Array Chunk))
-    (handler : Request Body → Async (Response Body)) (expected : String) (chunked : Bool := false) :
+    (handler : Request Body → ContextAsync (Response Body)) (expected : String) (chunked : Bool := false) :
     IO Unit := do
   let response ← sendRequests client server #[req] handler chunked
   let responseData := String.fromUTF8! response