netget/client/socketio/
mod.rs

1//! Socket.IO v5 client over Engine.IO v4 (direct WebSocket or HTTP long-polling).
2pub mod actions;
3use crate::client::http_fetch::FetchClient;
4use crate::client::llm_budget::call_llm_for_client;
5use crate::llm::actions::client_trait::{Client, ClientActionResult};
6use crate::logging::emit::Log;
7use crate::protocol::{ConnectContext, Event};
8use crate::server::socketio::packet::{self, Eio, Kind, Sio};
9use crate::state::client_handles::{ClientCommand, ClientSendOutcome};
10use crate::state::{AccessLogOwner, ClientStatus};
11pub use actions::SocketIoClientProtocol;
12use anyhow::{bail, ensure, Context, Result};
13use futures::{SinkExt, StreamExt};
14use serde_json::{json, Value};
15use std::collections::{HashMap, HashSet};
16use std::net::SocketAddr;
17use std::sync::Arc;
18use std::time::Duration;
19use tokio::sync::{mpsc, Mutex};
20use tokio_tungstenite::tungstenite::Message;
21
22pub const DEFAULT_TRANSPORT: &str = "websocket";
23const TIMEOUT: Duration = Duration::from_secs(30);
24
25type WsStream =
26    tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>;
27
28enum Transport {
29    Ws(Mutex<futures::stream::SplitSink<WsStream, Message>>),
30    Polling { fetch: FetchClient, url: String },
31}
32
33impl Transport {
34    fn name(&self) -> &'static str {
35        match self {
36            Transport::Ws(_) => "websocket",
37            Transport::Polling { .. } => "polling",
38        }
39    }
40
41    async fn send(&self, packets: Vec<String>) -> Result<()> {
42        match self {
43            Transport::Ws(sink) => {
44                let mut s = sink.lock().await;
45                for p in packets {
46                    s.send(Message::Text(p)).await?;
47                }
48            }
49            Transport::Polling { fetch, url } => {
50                let r = fetch
51                    .post(url)
52                    .header("Content-Type", "text/plain;charset=UTF-8")
53                    .body(packet::join_payload(&packets))
54                    .send()
55                    .await?;
56                ensure!(
57                    r.status().as_u16() == 200,
58                    "polling POST answered HTTP {}",
59                    r.status()
60                );
61            }
62        }
63        Ok(())
64    }
65}
66
67fn open_info(open: &str) -> Result<String> {
68    let v: Value = serde_json::from_str(open).context("open packet is not JSON")?;
69    let sid = v["sid"]
70        .as_str()
71        .filter(|s| !s.is_empty() && s.len() <= 128)
72        .context("open packet has no sid")?;
73    Ok(sid.to_owned())
74}
75
76pub async fn connect(ctx: ConnectContext) -> Result<SocketAddr> {
77    let p = ctx.startup_params.as_ref();
78    let path = p
79        .map(|p| p.get_optional_string("path"))
80        .transpose()?
81        .flatten()
82        .unwrap_or_else(|| crate::server::socketio::DEFAULT_PATH.to_owned());
83    ensure!(
84        path.starts_with('/') && path.len() <= 128 && !path.contains(['?', '#', ' ']),
85        "path must be an absolute path"
86    );
87    let path = format!("{}/", path.trim_end_matches('/'));
88    let transport = p
89        .map(|p| p.get_optional_string("transport"))
90        .transpose()?
91        .flatten()
92        .unwrap_or_else(|| DEFAULT_TRANSPORT.to_owned());
93    ensure!(
94        transport == "websocket" || transport == "polling",
95        "transport must be websocket or polling"
96    );
97    let namespaces: Vec<String> = match p
98        .map(|p| p.get_optional_array("namespaces"))
99        .transpose()?
100        .flatten()
101    {
102        None => vec!["/".into()],
103        Some(list) => list
104            .iter()
105            .map(|n| {
106                n.as_str()
107                    .filter(|n| packet::namespace_ok(n))
108                    .map(str::to_owned)
109                    .ok_or_else(|| anyhow::anyhow!("namespaces look like /name"))
110            })
111            .collect::<Result<_>>()?,
112    };
113    let auth = p
114        .map(|p| p.get_optional_object("auth"))
115        .transpose()?
116        .flatten()
117        .map(|m| Value::Object(m.clone()));
118    let (incoming_tx, incoming_rx) = mpsc::channel::<Result<Eio>>(256);
119    let local: SocketAddr = "0.0.0.0:0".parse()?;
120    let transport: Arc<Transport> = if transport == "websocket" {
121        let url = format!("ws://{}{path}?EIO=4&transport=websocket", ctx.remote_addr);
122        let mut cfg = tokio_tungstenite::tungstenite::protocol::WebSocketConfig::default();
123        cfg.max_message_size = Some(packet::MAX_PAYLOAD);
124        let (ws, _) = tokio::time::timeout(
125            TIMEOUT,
126            tokio_tungstenite::connect_async_with_config(url.as_str(), Some(cfg), false),
127        )
128        .await
129        .context("WebSocket connect timed out")??;
130        let (sink, mut stream) = ws.split();
131        let first = tokio::time::timeout(TIMEOUT, stream.next())
132            .await
133            .context("no Engine.IO open packet")?;
134        let Some(Ok(Message::Text(t))) = first else {
135            bail!("the server did not open an Engine.IO session")
136        };
137        let Eio::Open(info) = Eio::decode(&t)? else {
138            bail!("expected the Engine.IO open packet")
139        };
140        open_info(&info)?;
141        let reader = tokio::spawn(async move {
142            while let Some(m) = stream.next().await {
143                let item = match m {
144                    Ok(Message::Text(t)) => Eio::decode(&t),
145                    Ok(Message::Close(_)) | Err(_) => break,
146                    Ok(_) => continue,
147                };
148                if incoming_tx.send(item).await.is_err() {
149                    break;
150                }
151            }
152        });
153        ctx.state.register_client_task(ctx.client_id, reader).await;
154        Arc::new(Transport::Ws(Mutex::new(sink)))
155    } else {
156        let base = format!("http://{}{path}?EIO=4&transport=polling", ctx.remote_addr);
157        crate::client::http_fetch::check_url(&base)?;
158        #[cfg(not(target_arch = "wasm32"))]
159        let fetch = FetchClient::from_reqwest(
160            crate::llm::ollama_client::configured_for_endpoint(
161                reqwest::Client::builder().timeout(Duration::from_secs(120)),
162                &base,
163            )
164            .build()?,
165        );
166        #[cfg(target_arch = "wasm32")]
167        let fetch = FetchClient::transport(Duration::from_secs(120));
168        let fetch = fetch
169            .with_max_body(packet::MAX_PAYLOAD)
170            .with_user_agent("netget-socketio");
171        let r = fetch
172            .get(&base)
173            .send()
174            .await
175            .context("Engine.IO handshake")?;
176        ensure!(
177            r.status().as_u16() == 200,
178            "Engine.IO handshake answered HTTP {}",
179            r.status()
180        );
181        let body = String::from_utf8(r.bytes().await?.to_vec()).context("handshake is not text")?;
182        let packets = packet::split_payload(&body)?;
183        let Some(Eio::Open(info)) = packets.first().cloned() else {
184            bail!("expected the Engine.IO open packet")
185        };
186        let sid = open_info(&info)?;
187        let url = format!("{base}&sid={sid}");
188        for p in packets.into_iter().skip(1) {
189            incoming_tx.try_send(Ok(p))?;
190        }
191        let poll_fetch = fetch.clone();
192        let poll_url = url.clone();
193        let reader = tokio::spawn(async move {
194            loop {
195                let r = match poll_fetch.get(&poll_url).send().await {
196                    Ok(r) if r.status().as_u16() == 200 => r,
197                    _ => break,
198                };
199                let Ok(body) = r.bytes().await else { break };
200                let Ok(text) = String::from_utf8(body.to_vec()) else {
201                    break;
202                };
203                let items: Vec<Result<Eio>> = match packet::split_payload(&text) {
204                    Ok(list) => list.into_iter().map(Ok).collect(),
205                    Err(e) => vec![Err(e)],
206                };
207                for item in items {
208                    let closing = matches!(item, Ok(Eio::Close));
209                    if incoming_tx.send(item).await.is_err() || closing {
210                        return;
211                    }
212                }
213            }
214        });
215        ctx.state.register_client_task(ctx.client_id, reader).await;
216        Arc::new(Transport::Polling { fetch, url })
217    };
218    let connects: Vec<String> = namespaces
219        .iter()
220        .map(|n| Eio::Message(Sio::new(Kind::Connect, n, None, auth.clone()).encode()).encode())
221        .collect();
222    transport
223        .send(connects)
224        .await
225        .context("sending namespace CONNECT")?;
226    ctx.state
227        .update_client_status(ctx.client_id, ClientStatus::Connected)
228        .await;
229    let external =
230        crate::client::command_support::register_command_channel(&ctx.state, ctx.client_id).await;
231    let (internal_tx, internal_rx) = mpsc::channel::<Value>(16);
232    let (event_tx, mut event_rx) = mpsc::channel::<Event>(64);
233    let events_ctx = ctx.clone();
234    let dispatcher = tokio::spawn(async move {
235        let protocol = SocketIoClientProtocol;
236        while let Some(event) = event_rx.recv().await {
237            let instruction = events_ctx
238                .state
239                .get_instruction_for_client(events_ctx.client_id)
240                .await
241                .unwrap_or_default();
242            let memory = events_ctx
243                .state
244                .get_memory_for_client(events_ctx.client_id)
245                .await
246                .unwrap_or_default();
247            match call_llm_for_client(
248                &events_ctx.llm_client,
249                &events_ctx.state,
250                events_ctx.client_id.to_string(),
251                &instruction,
252                &memory,
253                Some(&event),
254                &protocol,
255                &events_ctx.status_tx,
256            )
257            .await
258            {
259                Ok(result) => {
260                    if let Some(memory) = result.memory_updates {
261                        events_ctx
262                            .state
263                            .set_memory_for_client(events_ctx.client_id, memory)
264                            .await;
265                    }
266                    for action in result.actions {
267                        if internal_tx.send(action).await.is_err() {
268                            return;
269                        }
270                    }
271                }
272                Err(e) => Log::new(Some(&events_ctx.status_tx))
273                    .warn(format!("Socket.IO client handler: {e}")),
274            }
275        }
276    });
277    ctx.state
278        .register_client_task(ctx.client_id, dispatcher)
279        .await;
280    let session_ctx = ctx.clone();
281    let task = tokio::spawn(async move {
282        let result = session(
283            &session_ctx,
284            &transport,
285            auth,
286            external,
287            internal_rx,
288            incoming_rx,
289            event_tx,
290        )
291        .await;
292        let _ = transport.send(vec![Eio::Close.encode()]).await;
293        let status = match result {
294            Ok(()) => ClientStatus::Disconnected,
295            Err(e) => {
296                Log::new(Some(&session_ctx.status_tx)).warn(format!("Socket.IO client ended: {e}"));
297                ClientStatus::Error(e.to_string())
298            }
299        };
300        session_ctx
301            .state
302            .update_client_status(session_ctx.client_id, status)
303            .await;
304        session_ctx
305            .state
306            .remove_client_handle(session_ctx.client_id)
307            .await;
308        let _ = session_ctx.status_tx.send("__UPDATE_UI__".into());
309    });
310    ctx.state.register_client_task(ctx.client_id, task).await;
311    Ok(local)
312}
313
314fn reply(command: Option<ClientCommand>, outcome: Result<ClientSendOutcome>) {
315    if let Some(c) = command {
316        crate::client::command_support::reply(c, outcome);
317    }
318}
319
320#[allow(clippy::too_many_arguments)]
321async fn session(
322    ctx: &ConnectContext,
323    transport: &Transport,
324    auth: Option<Value>,
325    mut external: mpsc::Receiver<ClientCommand>,
326    mut internal: mpsc::Receiver<Value>,
327    mut incoming: mpsc::Receiver<Result<Eio>>,
328    events: mpsc::Sender<Event>,
329) -> Result<()> {
330    let mut connected: HashSet<String> = HashSet::new();
331    let mut pending: HashMap<u64, (String, String)> = HashMap::new();
332    let mut next_ack: u64 = 1;
333    loop {
334        let (action, command) = tokio::select! {
335            c = external.recv() => match c { Some(c) => (c.action.clone(), Some(c)), None => return Ok(()) },
336            a = internal.recv() => match a { Some(a) => (a, None), None => return Ok(()) },
337            m = incoming.recv() => {
338                let p = match m {
339                    Some(Ok(p)) => p,
340                    Some(Err(e)) => bail!("protocol error from the server: {e}"),
341                    None => {
342                        events.send(Event::new(&actions::DISCONNECTED_EVENT, json!({"namespace": "*", "reason": "transport close"}))).await.ok();
343                        return Ok(());
344                    }
345                };
346                match p {
347                    Eio::Ping(d) => transport.send(vec![Eio::Pong(d).encode()]).await?,
348                    Eio::Close => {
349                        events.send(Event::new(&actions::DISCONNECTED_EVENT, json!({"namespace": "*", "reason": "io server disconnect"}))).await.ok();
350                        return Ok(());
351                    }
352                    Eio::Message(m) => {
353                        let sio = match Sio::decode(&m) {
354                            Ok(s) => s,
355                            Err(e) => {
356                                Log::new(Some(&ctx.status_tx)).warn(format!("Socket.IO client: {e}"));
357                                continue;
358                            }
359                        };
360                        let event = match sio.kind {
361                            Kind::Connect => {
362                                connected.insert(sio.nsp.clone());
363                                let sid = sio.data.as_ref().and_then(|d| d["sid"].as_str()).unwrap_or_default().to_owned();
364                                Event::new(&actions::CONNECTED_EVENT, json!({"namespace": sio.nsp, "socket_id": sid, "transport": transport.name()}))
365                            }
366                            Kind::ConnectError => {
367                                let d = sio.data.clone().unwrap_or(Value::Null);
368                                let message = d["message"].as_str().or(d.as_str()).unwrap_or("refused").to_owned();
369                                Event::new(&actions::CONNECT_ERROR_EVENT, json!({"namespace": sio.nsp, "message": message, "data": d.get("data")}))
370                            }
371                            Kind::Disconnect => {
372                                connected.remove(&sio.nsp);
373                                Event::new(&actions::DISCONNECTED_EVENT, json!({"namespace": sio.nsp, "reason": "io server disconnect"}))
374                            }
375                            Kind::Event => match sio.event() {
376                                Ok((name, args)) => Event::new(&actions::EVENT_EVENT, json!({"namespace": sio.nsp, "event": name, "args": args, "ack_id": sio.id})),
377                                Err(e) => {
378                                    Log::new(Some(&ctx.status_tx)).warn(format!("Socket.IO client: {e}"));
379                                    continue;
380                                }
381                            },
382                            Kind::Ack => {
383                                let Some((nsp, event)) = sio.id.and_then(|id| pending.remove(&id)) else { continue };
384                                Event::new(&actions::ACK_EVENT, json!({"namespace": nsp, "event": event, "args": sio.data.and_then(|d| d.as_array().cloned()).unwrap_or_default()}))
385                            }
386                            Kind::BinaryEvent | Kind::BinaryAck => continue,
387                        };
388                        events.send(event).await.context("Socket.IO event consumer stopped")?;
389                    }
390                    Eio::Open(_) | Eio::Pong(_) | Eio::Upgrade | Eio::Noop => {}
391                }
392                continue;
393            }
394        };
395        match SocketIoClientProtocol.execute_action(action.clone()) {
396            Ok(ClientActionResult::Disconnect) => {
397                reply(command, Ok(ClientSendOutcome::Disconnected));
398                return Ok(());
399            }
400            Ok(_) => {}
401            Err(e) => {
402                reply(
403                    command,
404                    Ok(ClientSendOutcome::Rejected {
405                        error: e.to_string(),
406                    }),
407                );
408                continue;
409            }
410        }
411        let nsp = action["namespace"].as_str().unwrap_or("/").to_owned();
412        let built: Result<Sio> = match action["type"].as_str().unwrap_or_default() {
413            "socketio_emit" => {
414                if !connected.contains(&nsp) {
415                    Err(anyhow::anyhow!("not connected to {nsp}"))
416                } else {
417                    let id = action["ack"].as_bool().unwrap_or(false).then(|| {
418                        let id = next_ack;
419                        next_ack += 1;
420                        pending.insert(
421                            id,
422                            (
423                                nsp.clone(),
424                                action["event"].as_str().unwrap_or_default().to_owned(),
425                            ),
426                        );
427                        id
428                    });
429                    let mut payload = vec![action["event"].clone()];
430                    payload.extend(action["args"].as_array().cloned().unwrap_or_default());
431                    Ok(Sio::new(Kind::Event, &nsp, id, Some(Value::Array(payload))))
432                }
433            }
434            "socketio_ack" => Ok(Sio::new(
435                Kind::Ack,
436                &nsp,
437                action["ack_id"].as_u64(),
438                Some(Value::Array(
439                    action["args"].as_array().cloned().unwrap_or_default(),
440                )),
441            )),
442            "socketio_connect_namespace" => Ok(Sio::new(
443                Kind::Connect,
444                &nsp,
445                None,
446                action
447                    .get("auth")
448                    .filter(|a| a.is_object())
449                    .cloned()
450                    .or_else(|| auth.clone()),
451            )),
452            _ => {
453                connected.remove(&nsp);
454                Ok(Sio::new(Kind::Disconnect, &nsp, None, None))
455            }
456        };
457        let packet = match built {
458            Ok(p) => Eio::Message(p.encode()).encode(),
459            Err(e) => {
460                reply(
461                    command,
462                    Ok(ClientSendOutcome::Rejected {
463                        error: e.to_string(),
464                    }),
465                );
466                continue;
467            }
468        };
469        let n = packet.len();
470        let sent = transport.send(vec![packet]).await;
471        if command.is_some() {
472            ctx.state
473                .record_access_log(
474                    AccessLogOwner::Client(ctx.client_id.as_u32()),
475                    "Socket.IO",
476                    None,
477                    "injected_action",
478                    json!({"type": action["type"], "event": action["event"]}),
479                    vec![json!({"ok": sent.is_ok()})],
480                )
481                .await;
482        }
483        match sent {
484            Ok(()) => reply(command, Ok(ClientSendOutcome::Sent { bytes_sent: n })),
485            Err(e) => {
486                reply(command, Err(anyhow::anyhow!("{e}")));
487                return Err(e);
488            }
489        }
490    }
491}
netget/client/socketio/mod.rs

netget/client/socketio/
mod.rs
