Add SSE middleware

Author: teamon

Diff for: README.md

@@ -243,7 +243,11 @@ Tesla.get(client, "/", opts: [adapter: [recv_timeout: 30_000]])
 
 ## Streaming
 
-If adapter supports it, you can pass a [Stream](https://elixir-lang.org/docs/stable/elixir/Stream.html) as body, e.g.:
+### Streaming Request Body
+
+If adapter supports it, you can pass a
+[Stream](https://elixir-lang.org/docs/stable/elixir/Stream.html) as request
+body, e.g.:
 
 ```elixir
 defmodule ElasticSearch do
@@ -259,7 +263,31 @@ defmodule ElasticSearch do
 end
 ```
 
-Each piece of stream will be encoded as JSON and sent as a new line (conforming to JSON stream format).
+Each piece of stream will be encoded as JSON and sent as a new line (conforming
+to JSON stream format).
+
+### Streaming Response Body
+
+If adapter supports it, you can pass a `response: :stream` option to return
+response body as a
+[Stream](https://elixir-lang.org/docs/stable/elixir/Stream.html)
+
+```elixir
+defmodule OpenAI do
+  use Tesla
+
+  plug Tesla.Middleware.BaseUrl, "https://api.openai.com/v1"
+  plug Tesla.Middleware.JSON
+
+  def completion(messages) do
+    post("/chat/completions", %{model: "gpt-3.5-turbo", messages: messages, stream: true})
+  end
+end
+
+{:ok, env} = OpenAI.completion("What is the meaning of life?")
+env.body
+|> Stream.each(fn chunk -> IO.puts(chunk) end)
+```
 
 ## Multipart
 

Diff for: lib/tesla/adapter/finch.ex

@@ -131,7 +131,9 @@ if Code.ensure_loaded?(Finch) do
                   Task.await(task)
                   nil
               after
-                opts[:receive_timeout] -> nil
+                opts[:receive_timeout] ->
+                  Task.shutdown(task, :brutal_kill)
+                  nil
               end
             end)
 

Diff for: lib/tesla/middleware/json.ex

@@ -108,12 +108,18 @@ defmodule Tesla.Middleware.JSON do
     end
   end
 
+  defp decode_body(body, opts) when is_struct(body, Stream) or is_function(body),
+    do: {:ok, decode_stream(body, opts)}
+
   defp decode_body(body, opts), do: process(body, :decode, opts)
 
   defp decodable?(env, opts), do: decodable_body?(env) && decodable_content_type?(env, opts)
 
   defp decodable_body?(env) do
-    (is_binary(env.body) && env.body != "") || (is_list(env.body) && env.body != [])
+    (is_binary(env.body) && env.body != "") ||
+      (is_list(env.body) && env.body != []) ||
+      is_function(env.body) ||
+      is_struct(env.body, Stream)
   end
 
   defp decodable_content_type?(env, opts) do
@@ -123,6 +129,15 @@ defmodule Tesla.Middleware.JSON do
     end
   end
 
+  defp decode_stream(body, opts) do
+    Stream.map(body, fn chunk ->
+      case decode_body(chunk, opts) do
+        {:ok, item} -> item
+        _ -> chunk
+      end
+    end)
+  end
+
   defp content_types(opts),
     do: @default_content_types ++ Keyword.get(opts, :decode_content_types, [])
 

Diff for: lib/tesla/middleware/sse.ex

@@ -0,0 +1,102 @@
+defmodule Tesla.Middleware.SSE do
+  @moduledoc """
+  Decode Server Sent Events.
+
+  This middleware is mostly useful when streaming response body.
+
+  ## Examples
+
+  ```
+  plug Tesla.Middleware.SSE, only: :data
+
+  ```
+
+  ## Options
+
+  - `:only` - keep only specified keys in event (necessary for using with `JSON` middleware)
+  - `:decode_content_types` - list of additional decodable content-types
+  """
+
+  @behaviour Tesla.Middleware
+
+  @default_content_types ["text/event-stream"]
+
+  @impl Tesla.Middleware
+  def call(env, next, opts) do
+    opts = opts || []
+
+    with {:ok, env} <- Tesla.run(env, next) do
+      decode(env, opts)
+    end
+  end
+
+  def decode(env, opts) do
+    if decodable_content_type?(env, opts) do
+      {:ok, %{env | body: decode_body(env.body, opts)}}
+    else
+      {:ok, env}
+    end
+  end
+
+  defp decode_body(body, opts) when is_struct(body, Stream) or is_function(body) do
+    body
+    |> Stream.chunk_while(
+      "",
+      fn elem, acc ->
+        {lines, [rest]} = (acc <> elem) |> String.split("\n\n") |> Enum.split(-1)
+        {:cont, lines, rest}
+      end,
+      fn
+        "" -> {:cont, ""}
+        acc -> {:cont, acc, ""}
+      end
+    )
+    |> Stream.flat_map(& &1)
+    |> Stream.map(&decode_message/1)
+    |> Stream.flat_map(&only(&1, opts[:only]))
+  end
+
+  defp decode_body(binary, opts) when is_binary(binary) do
+    binary
+    |> String.split("\n\n")
+    |> Enum.map(&decode_message/1)
+    |> Enum.flat_map(&only(&1, opts[:only]))
+  end
+
+  defp decode_message(message) do
+    message
+    |> String.split("\n")
+    |> Enum.map(&decode_body/1)
+    |> Enum.reduce(%{}, fn
+      :empty, acc -> acc
+      {:data, data}, acc -> Map.update(acc, :data, data, &(&1 <> "\n" <> data))
+      {key, value}, acc -> Map.put_new(acc, key, value)
+    end)
+  end
+
+  defp decode_body(": " <> comment), do: {:comment, comment}
+  defp decode_body("data: " <> data), do: {:data, data}
+  defp decode_body("event: " <> event), do: {:event, event}
+  defp decode_body("id: " <> id), do: {:id, id}
+  defp decode_body("retry: " <> retry), do: {:retry, retry}
+  defp decode_body(""), do: :empty
+
+  defp decodable_content_type?(env, opts) do
+    case Tesla.get_header(env, "content-type") do
+      nil -> false
+      content_type -> Enum.any?(content_types(opts), &String.starts_with?(content_type, &1))
+    end
+  end
+
+  defp content_types(opts),
+    do: @default_content_types ++ Keyword.get(opts, :decode_content_types, [])
+
+  defp only(message, nil), do: [message]
+
+  defp only(message, key) do
+    case Map.get(message, key) do
+      nil -> []
+      val -> [val]
+    end
+  end
+end

Diff for: mix.exs

@@ -79,7 +79,7 @@ defmodule Tesla.Mixfile do
       {:excoveralls, "~> 0.8", only: :test},
       {:httparrot, "~> 1.3", only: :test},
       {:ex_doc, "~> 0.21", only: :dev, runtime: false},
-      {:mix_test_watch, "~> 1.0", only: :test, runtime: false},
+      {:mix_test_watch, "~> 1.0", only: [:dev, :test], runtime: false},
       {:dialyxir, "~> 1.0", only: [:dev, :test], runtime: false},
       {:inch_ex, "~> 2.0", only: :docs},
 

Diff for: test/tesla/middleware/json_test.exs

@@ -167,6 +167,32 @@ defmodule Tesla.Middleware.JsonTest do
     end
   end
 
+  describe "Streams" do
+    test "encode stream" do
+      adapter = fn env ->
+        assert IO.iodata_to_binary(Enum.to_list(env.body)) == ~s|{"id":1}\n{"id":2}\n{"id":3}\n|
+      end
+
+      stream = Stream.map(1..3, fn i -> %{id: i} end)
+      Tesla.Middleware.JSON.call(%Tesla.Env{body: stream}, [{:fn, adapter}], [])
+    end
+
+    test "decode stream" do
+      adapter = fn _env ->
+        stream = Stream.map(1..3, fn i -> ~s|{"id": #{i}}\n| end)
+
+        {:ok,
+         %Tesla.Env{
+           headers: [{"content-type", "application/json"}],
+           body: stream
+         }}
+      end
+
+      assert {:ok, env} = Tesla.Middleware.JSON.call(%Tesla.Env{}, [{:fn, adapter}], [])
+      assert Enum.to_list(env.body) == [%{"id" => 1}, %{"id" => 2}, %{"id" => 3}]
+    end
+  end
+
   describe "Multipart" do
     defmodule MultipartClient do
       use Tesla

Diff for: test/tesla/middleware/sse_test.exs

@@ -0,0 +1,124 @@
+defmodule Tesla.Middleware.SSETest do
+  use ExUnit.Case
+
+  @env %Tesla.Env{
+    status: 200,
+    headers: [{"content-type", "text/event-stream"}]
+  }
+
+  describe "Basics" do
+    test "ignore not matching content-type" do
+      adapter = fn _env ->
+        {:ok, %Tesla.Env{headers: [{"content-type", "text/plain"}], body: "test"}}
+      end
+
+      assert {:ok, env} = Tesla.Middleware.SSE.call(%Tesla.Env{}, [{:fn, adapter}], [])
+      assert env.body == "test"
+    end
+
+    test "decode comment" do
+      adapter = fn _env ->
+        {:ok, %{@env | body: ": comment"}}
+      end
+
+      assert {:ok, env} = Tesla.Middleware.SSE.call(%Tesla.Env{}, [{:fn, adapter}], [])
+      assert env.body == [%{comment: "comment"}]
+    end
+
+    test "decode multiple messages" do
+      body = """
+      : this is a test stream
+
+      data: some text
+
+      data: another message
+      data: with two lines
+      """
+
+      adapter = fn _env ->
+        {:ok, %{@env | body: body}}
+      end
+
+      assert {:ok, env} = Tesla.Middleware.SSE.call(%Tesla.Env{}, [{:fn, adapter}], [])
+
+      assert env.body == [
+               %{comment: "this is a test stream"},
+               %{data: "some text"},
+               %{data: "another message\nwith two lines"}
+             ]
+    end
+
+    test "decode named events" do
+      body = """
+      event: userconnect
+      data: {"username": "bobby", "time": "02:33:48"}
+
+      data: Here's a system message of some kind that will get used
+      data: to accomplish some task.
+
+      event: usermessage
+      data: {"username": "bobby", "time": "02:34:11", "text": "Hi everyone."}
+      """
+
+      adapter = fn _env ->
+        {:ok, %{@env | body: body}}
+      end
+
+      assert {:ok, env} = Tesla.Middleware.SSE.call(%Tesla.Env{}, [{:fn, adapter}], [])
+
+      assert env.body == [
+               %{event: "userconnect", data: ~s|{"username": "bobby", "time": "02:33:48"}|},
+               %{
+                 data:
+                   "Here's a system message of some kind that will get used\nto accomplish some task."
+               },
+               %{
+                 event: "usermessage",
+                 data: ~s|{"username": "bobby", "time": "02:34:11", "text": "Hi everyone."}|
+               }
+             ]
+    end
+
+    test "output only data" do
+      body = """
+      : comment1
+
+      event: userconnect
+      data: data1
+
+      data: data2
+      data: data3
+
+      event: usermessage
+      data: data4
+      """
+
+      adapter = fn _env ->
+        {:ok, %{@env | body: body}}
+      end
+
+      assert {:ok, env} = Tesla.Middleware.SSE.call(%Tesla.Env{}, [{:fn, adapter}], only: :data)
+
+      assert env.body == ["data1", "data2\ndata3", "data4"]
+    end
+
+    test "handle stream data" do
+      adapter = fn _env ->
+        chunks = [
+          ~s|dat|,
+          ~s|a: dat|,
+          ~s|a1\n\ndata: data2\n\ndata: d|,
+          ~s|ata3\n\n|
+        ]
+
+        stream = Stream.map(chunks, & &1)
+
+        {:ok, %{@env | body: stream}}
+      end
+
+      assert {:ok, env} = Tesla.Middleware.SSE.call(%Tesla.Env{}, [{:fn, adapter}], [])
+
+      assert Enum.to_list(env.body) == [%{data: "data1"}, %{data: "data2"}, %{data: "data3"}]
+    end
+  end
+end