From f664a53f76b35ce06e356343ed935c77a5b998f2 Mon Sep 17 00:00:00 2001
From: Steffen Deusch <steffen@deusch.me>
Date: Tue, 30 Jun 2026 12:42:12 +0200
Subject: [PATCH 1/3] Fix register_before_send with upgrade_adapter

Adds :set_upgrade state and runs before_send before calling upgrade.
---
 lib/plug/conn.ex | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/lib/plug/conn.ex b/lib/plug/conn.ex
index 5780ca71..86a31a29 100644
--- a/lib/plug/conn.ex
+++ b/lib/plug/conn.ex
@@ -294,7 +294,7 @@ defmodule Plug.Conn do
   alias Plug.Conn
   @epoch {{1970, 1, 1}, {0, 0, 0}}
   @already_sent {:plug_conn, :sent}
-  @unsent [:unset, :set, :set_chunked, :set_file]
+  @unsent [:unset, :set, :set_upgrade, :set_chunked, :set_file]
 
   @doc """
   Assigns a value to a key in the connection.
@@ -1471,10 +1471,11 @@ defmodule Plug.Conn do
   @spec upgrade_adapter(t, atom, term) :: t
   def upgrade_adapter(%Conn{adapter: {adapter, payload}, state: state} = conn, protocol, args)
       when state in @unsent do
+    conn = run_before_send(conn, :set_upgrade)
+
     case adapter.upgrade(payload, protocol, args) do
       {:ok, payload} ->
-        conn = run_before_send(conn, :upgraded)
-        %{conn | adapter: {adapter, payload}}
+        %{conn | state: :upgraded, adapter: {adapter, payload}}
 
       {:error, :not_supported} ->
         raise ArgumentError, "upgrade to #{protocol} not supported by #{inspect(adapter)}"

From d63611f5284b4c9fce41613205fc0df3a81dd310 Mon Sep 17 00:00:00 2001
From: Steffen Deusch <steffen@deusch.me>
Date: Tue, 30 Jun 2026 12:45:45 +0200
Subject: [PATCH 2/3] test

---
 test/plug/conn_test.exs | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/test/plug/conn_test.exs b/test/plug/conn_test.exs
index d5080944..30a40973 100644
--- a/test/plug/conn_test.exs
+++ b/test/plug/conn_test.exs
@@ -510,9 +510,18 @@ defmodule Plug.ConnTest do
   test "upgrade_adapter/3 runs before_send callbacks" do
     conn =
       conn(:get, "/foo")
-      |> register_before_send(&assign(&1, :ran_before_send, true))
+      |> register_before_send(fn conn ->
+        send(self(), {:state, conn.state})
+
+        conn
+        |> put_resp_header("x-test", "UPGRADE")
+        |> assign(:ran_before_send, true)
+      end)
       |> upgrade_adapter(:supported, opt: :supported)
 
+    assert_receive {:state, :set_upgrade}
+
+    assert {"x-test", "UPGRADE"} in conn.resp_headers
     assert conn.assigns[:ran_before_send] == true
   end
 

From 8e571d05b840ac89634dad518ff09ac3c57cb56d Mon Sep 17 00:00:00 2001
From: Steffen Deusch <steffen@deusch.me>
Date: Tue, 30 Jun 2026 13:32:40 +0200
Subject: [PATCH 3/3] set status to 101 (switching protocols)

---
 lib/plug/conn.ex        | 2 +-
 test/plug/conn_test.exs | 1 +
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/lib/plug/conn.ex b/lib/plug/conn.ex
index 86a31a29..d7177966 100644
--- a/lib/plug/conn.ex
+++ b/lib/plug/conn.ex
@@ -1471,7 +1471,7 @@ defmodule Plug.Conn do
   @spec upgrade_adapter(t, atom, term) :: t
   def upgrade_adapter(%Conn{adapter: {adapter, payload}, state: state} = conn, protocol, args)
       when state in @unsent do
-    conn = run_before_send(conn, :set_upgrade)
+    conn = run_before_send(%{conn | status: 101}, :set_upgrade)
 
     case adapter.upgrade(payload, protocol, args) do
       {:ok, payload} ->
diff --git a/test/plug/conn_test.exs b/test/plug/conn_test.exs
index 30a40973..95af3c85 100644
--- a/test/plug/conn_test.exs
+++ b/test/plug/conn_test.exs
@@ -512,6 +512,7 @@ defmodule Plug.ConnTest do
       conn(:get, "/foo")
       |> register_before_send(fn conn ->
         send(self(), {:state, conn.state})
+        assert conn.status == 101
 
         conn
         |> put_resp_header("x-test", "UPGRADE")