From 201dcb019280f7a6c95637bbdf3f322eab859f55 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 17 Oct 2025 00:50:29 -0400 Subject: [PATCH] wreq websocket wip --- Cargo.lock | 30 ++++ Cargo.toml | 2 +- crates/api_clients/grok_client/src/lib.rs | 6 +- .../create_listen_websocket.rs | 152 ++++++++++++++++++ .../src/requests/listen_websocket/mod.rs | 1 + .../src/test_utils/get_test_cookies.rs | 9 ++ .../grok_client/src/test_utils/mod.rs | 2 + 7 files changed, 200 insertions(+), 2 deletions(-) create mode 100644 crates/api_clients/grok_client/src/requests/listen_websocket/create_listen_websocket.rs create mode 100644 crates/api_clients/grok_client/src/test_utils/get_test_cookies.rs create mode 100644 crates/api_clients/grok_client/src/test_utils/mod.rs diff --git a/Cargo.lock b/Cargo.lock index 1e1d2607ca..5ec746e1ad 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11871,6 +11871,18 @@ dependencies = [ "tokio-reactor", ] +[[package]] +name = "tokio-tungstenite" +version = "0.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "489a59b6730eda1b0171fcfda8b121f4bee2b35cba8645ca35c5f7ba3eb736c1" +dependencies = [ + "futures-util", + "log", + "tokio", + "tungstenite", +] + [[package]] name = "tokio-uds" version = "0.2.7" @@ -12301,6 +12313,23 @@ dependencies = [ "workspace-build-acceleration", ] +[[package]] +name = "tungstenite" +version = "0.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eadc29d668c91fcc564941132e17b28a7ceb2f3ebf0b9dae3e03fd7a6748eb0d" +dependencies = [ + "bytes 1.10.1", + "data-encoding", + "http 1.3.1", + "httparse", + "log", + "rand 0.9.1", + "sha1 0.10.6", + "thiserror 2.0.12", + "utf-8", +] + [[package]] name = "typed-builder" version = "0.22.0" @@ -13869,6 +13898,7 @@ dependencies = [ "sync_wrapper 1.0.2", "tokio", "tokio-boring2", + "tokio-tungstenite", "tower 0.5.2", "tower-http", "want", diff --git a/Cargo.toml b/Cargo.toml index 91a1dab38c..f855679886 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -269,5 +269,5 @@ url = { version = "2.5.4", features = ["serde"] } # choco install pkgconfiglite -y # choco install llvm -y # choco install nasm -y -wreq = { version = "6.0.0-rc.20", features = ["json", "multipart", "cookies"] } +wreq = { version = "6.0.0-rc.20", features = ["json", "multipart", "cookies", "ws"] } wreq-util = "3.0.0-rc.5" diff --git a/crates/api_clients/grok_client/src/lib.rs b/crates/api_clients/grok_client/src/lib.rs index 8de4dc7fc4..6f3ea9c770 100644 --- a/crates/api_clients/grok_client/src/lib.rs +++ b/crates/api_clients/grok_client/src/lib.rs @@ -1,5 +1,9 @@ //! A Grok API client. +#[cfg(test)] +pub (crate) mod test_utils; + +// User lib pub mod credentials; pub mod error; -pub mod requests; \ No newline at end of file +pub mod requests; diff --git a/crates/api_clients/grok_client/src/requests/listen_websocket/create_listen_websocket.rs b/crates/api_clients/grok_client/src/requests/listen_websocket/create_listen_websocket.rs new file mode 100644 index 0000000000..5bb469d059 --- /dev/null +++ b/crates/api_clients/grok_client/src/requests/listen_websocket/create_listen_websocket.rs @@ -0,0 +1,152 @@ +use log::info; +use wreq::Client; +use wreq::header::{ACCEPT, ACCEPT_ENCODING, ACCEPT_LANGUAGE, CACHE_CONTROL, CONNECTION, COOKIE, ORIGIN, PRAGMA, UPGRADE, USER_AGENT}; +use wreq::ws::message::Message; +use wreq_util::Emulation; +use crate::error::grok_client_error::GrokClientError; +use crate::error::grok_error::GrokError; +use crate::error::grok_generic_api_error::GrokGenericApiError; + +const WEBSOCKET_URL: &str = "wss://grok.com/ws/imagine/listen"; + +pub struct CreateListenWebsocketArgs<'a> { + pub cookies: &'a str, +} + +pub async fn create_listen_websocket(args: CreateListenWebsocketArgs<'_>) -> Result<(), GrokError> { + println!("Building client..."); + info!("Building client..."); + + let client = Client::builder() + .emulation(Emulation::Firefox139) + .build() + .map_err(|err| GrokClientError::WreqClientError(err))?; + + println!("Configuring client..."); + info!("Configuring client..."); + + let builder = client.get(WEBSOCKET_URL) + .header(ACCEPT, "*/*") + .header(USER_AGENT, "User-Agent: Mozilla/5.0 (Macintosh; Intel Mac OS X 10.15; rv:144.0) Gecko/20100101 Firefox/144.0") + .header(ACCEPT_LANGUAGE, "en-US,en;q=0.5") + .header(ACCEPT_ENCODING, "gzip, deflate, br, zstd") + .header("Sec-WebSocket-Version", "13") + .header(ORIGIN, "https://grok.com") + .header("Sec-WebSocket-Extensions", "permessage-deflate") + // Sec-WebSocket-Key: BhBXbFSG6/1xcZVq4ySxcg== + .header("Sec-GPC", "1") + .header(CONNECTION, "keep-alive, Upgrade") + .header(COOKIE, args.cookies.to_string()) + .header("Sec-Fetch-Dest", "empty") + .header("Sec-Fetch-Mode", "websocket") + .header("Sec-Fetch-Site", "same-origin") + .header(PRAGMA, "no-cache") + .header(CACHE_CONTROL, "no-cache") + //.header(UPGRADE, "websocket") + //.map_err(|err| GrokClientError::WreqClientError(err))?; + + println!("Sending..."); + info!("Sending..."); + + let response = builder.send() + .await + .map_err(|err| GrokClientError::WreqClientError(err))?; + + let status = response.status(); + + for (k, v) in response.headers().iter() { + println!("Header: {}: {:?}", k, v); + } + + println!("Upgrading..."); + info!("Upgrading..."); + + let upgraded = response + .upgrade() + .await + .map_err(|err| GrokClientError::WreqClientError(err))?; + + println!("Status: {}", status); + + println!("Into websocket..."); + info!("Into websocket..."); + + let mut websocket = upgraded.into_websocket() + .await + .map_err(|err| GrokGenericApiError::WreqError(err))?; + + let message = r#" + {"type":"conversation.item.create","timestamp":1760675883072,"item":{"type":"message","content":[{"requestId":"3cedf20e-f51c-db5d-a124-bcec05feedf1","text":"A pirannah","type":"input_text","properties":{"section_count":0,"is_kids_mode":false,"enable_nsfw":true,"skip_upsampler":false,"is_initial":false}}]}} + "#.trim().to_string(); + + let wire_message = Message::text(message); + + websocket.send(wire_message) + .await + .map_err(|err| GrokGenericApiError::WreqError(err))?; + + loop { + let maybe_message = websocket.recv().await; + + if let Some(message) = maybe_message { + let message = message.map_err(|err| GrokGenericApiError::WreqError(err))?; + + match message { + Message::Text(text) => { + println!("Received text message: {}", text); + } + Message::Binary(bin) => { + println!("Received binary message: {:?}", bin); + } + Message::Ping(ping) => { + println!("Received ping: {:?}", ping); + } + Message::Pong(pong) => { + println!("Received pong: {:?}", pong); + } + Message::Close(close_frame) => { + println!("Received close: {:?}", close_frame); + break; + } + _ => { + println!("Received other message: {:?}", message); + } + } + + tokio::time::sleep(std::time::Duration::from_millis(15_000)).await; + } + } + + + Ok(()) +} + +#[cfg(test)] +mod tests { + use errors::AnyhowResult; + use crate::test_utils::get_test_cookies::get_test_cookies; + use super::*; + + #[tokio::test] + #[ignore] // manually test + async fn create() -> AnyhowResult<()> { + let cookies = get_test_cookies()?; + let args = CreateListenWebsocketArgs { + cookies: &cookies, + }; + let result = create_listen_websocket(args).await; + + match result { + Ok(ok) => { + println!("Okay"); + } + Err(err) => { + println!("{:?}", err); + } + } + + assert_eq!(1, 2); + + Ok(()) + } +} diff --git a/crates/api_clients/grok_client/src/requests/listen_websocket/mod.rs b/crates/api_clients/grok_client/src/requests/listen_websocket/mod.rs index e69de29bb2..6b7c343a03 100644 --- a/crates/api_clients/grok_client/src/requests/listen_websocket/mod.rs +++ b/crates/api_clients/grok_client/src/requests/listen_websocket/mod.rs @@ -0,0 +1 @@ +mod create_listen_websocket; \ No newline at end of file diff --git a/crates/api_clients/grok_client/src/test_utils/get_test_cookies.rs b/crates/api_clients/grok_client/src/test_utils/get_test_cookies.rs new file mode 100644 index 0000000000..d7a2336c03 --- /dev/null +++ b/crates/api_clients/grok_client/src/test_utils/get_test_cookies.rs @@ -0,0 +1,9 @@ +use errors::AnyhowResult; +use std::fs::read_to_string; + +#[cfg(test)] +pub fn get_test_cookies() -> AnyhowResult { + let cookies = read_to_string("/Users/bt/Artcraft/credentials/grok_cookies.txt")?; + let cookies = cookies.trim().to_string(); + Ok(cookies) +} diff --git a/crates/api_clients/grok_client/src/test_utils/mod.rs b/crates/api_clients/grok_client/src/test_utils/mod.rs new file mode 100644 index 0000000000..6dce54acd7 --- /dev/null +++ b/crates/api_clients/grok_client/src/test_utils/mod.rs @@ -0,0 +1,2 @@ +#[cfg(test)] +pub (crate) mod get_test_cookies;