Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 22 additions & 12 deletions kinode/src/http/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -265,7 +265,7 @@ async fn listen_to_stream(
match message {
Ok(msg) => {
// Handle different types of incoming WebSocket messages
let (body, blob) = match msg {
let (body, blob, should_exit) = match msg {
TungsteniteMessage::Text(text) => (
HttpClientRequest::WebSocketPush {
channel_id,
Expand All @@ -275,6 +275,7 @@ async fn listen_to_stream(
mime: Some("text/plain".into()),
bytes: text.into_bytes(),
}),
false,
),
TungsteniteMessage::Binary(bytes) => (
HttpClientRequest::WebSocketPush {
Expand All @@ -285,42 +286,51 @@ async fn listen_to_stream(
mime: Some("application/octet-stream".into()),
bytes,
}),
false,
),
TungsteniteMessage::Close(_) => {
// remove the websocket from the map
ws_streams.remove(&(target.process.clone(), channel_id));

(HttpClientRequest::WebSocketClose { channel_id }, None)
(HttpClientRequest::WebSocketClose { channel_id }, None, true)
}
TungsteniteMessage::Ping(_) => (
HttpClientRequest::WebSocketPush {
channel_id,
message_type: WsMessageType::Ping,
},
None,
false,
),
TungsteniteMessage::Pong(_) => (
HttpClientRequest::WebSocketPush {
channel_id,
message_type: WsMessageType::Pong,
},
None,
false,
),
_ => {
// should never get a TungsteniteMessage::Frame, ignore if we do
continue;
}
};

handle_ws_message(
our.clone(),
id,
target.clone(),
body,
blob,
send_to_loop.clone(),
)
.await;
if ws_streams.contains_key(&(target.process.clone(), channel_id)) || should_exit {
handle_ws_message(
our.clone(),
id,
target.clone(),
body,
blob,
send_to_loop.clone(),
)
.await;
}

if should_exit {
break;
}
}
Err(e) => {
println!("WebSocket Client Error ({}): {:?}", channel_id, e);
Expand Down Expand Up @@ -633,7 +643,7 @@ async fn close_ws_connection(
ws_streams: WebSocketStreams,
_print_tx: PrintSender,
) -> Result<HttpClientResponse, HttpClientError> {
let Some(mut ws_sink) = ws_streams.get_mut(&(target.process.clone(), channel_id)) else {
let Some((_, mut ws_sink)) = ws_streams.remove(&(target.process.clone(), channel_id)) else {
return Err(HttpClientError::WsCloseFailed { channel_id });
};

Expand Down