From 047e518edf563cfa84f5c964d1de4c341513cbd9 Mon Sep 17 00:00:00 2001
From: Jonathan Kelley <jkelleyrtp@gmail.com>
Date: Fri, 24 Oct 2025 19:35:59 -0700
Subject: [PATCH] feat: allow `.set()` on UseWebsocket to drive the connection
 manually

---
 packages/fullstack/src/payloads/websocket.rs | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/packages/fullstack/src/payloads/websocket.rs b/packages/fullstack/src/payloads/websocket.rs
index b0928b5be..6737adc36 100644
--- a/packages/fullstack/src/payloads/websocket.rs
+++ b/packages/fullstack/src/payloads/websocket.rs
@@ -233,6 +233,20 @@ impl<In, Out, E> UseWebsocket<In, Out, E> {
         result
     }
 
+    /// Set the WebSocket connection.
+    ///
+    /// This method takes a `Result<Websocket<In, Out, E>, Err>`, allowing you to drive the connection
+    /// into an errored state manually.
+    pub fn set<Err: Into<CapturedError>>(&mut self, socket: Result<Websocket<In, Out, E>, Err>) {
+        match socket {
+            Ok(_) => self.status.set(WebsocketState::Open),
+            Err(_) => self.status.set(WebsocketState::FailedToConnect),
+        }
+
+        self.connection.set(Some(socket.map_err(|e| e.into())));
+        self.waker.wake(());
+    }
+
     /// Mark the WebSocket as closed. This is called internally when the connection is closed.
     fn received_shutdown(&self) {
         let mut _self = *self;