netget/client/websocket/
actions.rs

1//! WebSocket (RFC 6455) client actions, events and metadata.
2//!
3//! The client mirrors the server's vocabulary deliberately: `send_websocket_text`,
4//! `send_websocket_binary` (with the same `data` + `encoding` pair), `send_websocket_ping` and
5//! `close_websocket` mean the same thing in both directions, so a prompt written for one side
6//! reads correctly on the other. Masking is the only asymmetry and it is handled by the
7//! framing layer — RFC 6455 §5.3 requires every client-to-server frame to be masked with a
8//! fresh 32-bit key, and `tungstenite` does that for `Role::Client` without being asked.
9
10use crate::llm::actions::{
11    client_trait::{Client, ClientActionResult},
12    protocol_trait::Protocol,
13    ActionDefinition, Parameter, ParameterDefinition,
14};
15use crate::protocol::log_template::LogTemplate;
16use crate::protocol::EventType;
17use crate::server::websocket::actions::{decode_outbound_payload, WsOut};
18use crate::state::app_state::AppState;
19use anyhow::{Context, Result};
20use serde_json::json;
21use std::sync::LazyLock;
22use tokio::sync::mpsc;
23use tracing::debug;
24
25/// WebSocket client action handler.
26///
27/// Two shapes, as on the server side: the registry-wide instance used for documentation and
28/// `connect()`, and a per-connection instance that owns the outbound channel and can therefore
29/// actually put a frame on the wire.
30pub struct WebSocketClientProtocol {
31    out_tx: Option<mpsc::UnboundedSender<WsOut>>,
32}
33
34impl Default for WebSocketClientProtocol {
35    fn default() -> Self {
36        Self::new()
37    }
38}
39
40impl WebSocketClientProtocol {
41    pub fn new() -> Self {
42        Self { out_tx: None }
43    }
44
45    pub fn for_connection(out_tx: mpsc::UnboundedSender<WsOut>) -> Self {
46        Self {
47            out_tx: Some(out_tx),
48        }
49    }
50
51    fn send(&self, msg: WsOut) -> Result<()> {
52        let tx = self
53            .out_tx
54            .as_ref()
55            .context("This WebSocket client action needs an open connection")?;
56        tx.send(msg)
57            .map_err(|_| anyhow::anyhow!("WebSocket connection is already closed"))
58    }
59}
60
61impl Protocol for WebSocketClientProtocol {
62    fn get_startup_parameters(&self) -> Vec<ParameterDefinition> {
63        vec![
64            ParameterDefinition {
65                name: "path".to_string(),
66                type_hint: "string".to_string(),
67                description:
68                    "Request path (and optional query string) to open the WebSocket on, e.g. \
69                     \"/ws\" or \"/stream?token=abc\". Defaults to \"/\"."
70                        .to_string(),
71                required: false,
72                example: json!("/ws"),
73                default: None,
74            },
75            ParameterDefinition {
76                name: "subprotocols".to_string(),
77                type_hint: "array".to_string(),
78                description:
79                    "Subprotocols to offer in Sec-WebSocket-Protocol, most preferred first. The \
80                     server picks at most one and it is reported in the \
81                     websocket_client_connected event."
82                        .to_string(),
83                required: false,
84                example: json!(["chat", "superchat"]),
85                default: None,
86            },
87            ParameterDefinition {
88                name: "headers".to_string(),
89                type_hint: "object".to_string(),
90                description: "Extra request headers to send with the upgrade, e.g. \
91                     {\"Authorization\": \"Bearer …\", \"Origin\": \"https://example.com\"}. \
92                     The RFC 6455 handshake headers are set automatically and cannot be \
93                     overridden."
94                    .to_string(),
95                required: false,
96                example: json!({"Origin": "https://example.com"}),
97                default: None,
98            },
99            ParameterDefinition {
100                name: "permessage_deflate".to_string(),
101                type_hint: "boolean".to_string(),
102                description: "Offer RFC 7692 permessage-deflate in the upgrade, so the server may \
103                     compress messages in both directions (default true). Messages reach the \
104                     handler as plain text either way; whether the server accepted is reported \
105                     in websocket_client_connected."
106                    .to_string(),
107                required: false,
108                example: json!(true),
109                default: Some(json!(super::DEFAULT_PERMESSAGE_DEFLATE)),
110            },
111        ]
112    }
113
114    fn get_async_actions(&self, _state: &AppState) -> Vec<ActionDefinition> {
115        client_actions()
116    }
117
118    fn get_sync_actions(&self) -> Vec<ActionDefinition> {
119        client_actions()
120    }
121
122    fn protocol_name(&self) -> &'static str {
123        // Same canonical name and stack as the server side, which is the repo's convention
124        // (TcpClientProtocol is "TCP" too). The two registries are separate, so there is no
125        // collision, and `base_stack: "websocket"` resolves on both.
126        "WebSocket"
127    }
128
129    fn get_event_types(&self) -> Vec<EventType> {
130        get_websocket_client_event_types()
131    }
132
133    fn stack_name(&self) -> &'static str {
134        "ETH>IP>TCP>HTTP>WebSocket"
135    }
136
137    fn keywords(&self) -> Vec<&'static str> {
138        // The client registry's `parse_from_str` has no protocol-name step and matches
139        // keywords by substring, so the bare "websocket" is what makes `base_stack:
140        // "websocket"` (and "websocket_client") resolve at all.
141        vec![
142            "websocket",
143            "web socket",
144            "websocket client",
145            "connect to websocket",
146            "rfc 6455",
147        ]
148    }
149
150    fn metadata(&self) -> crate::protocol::metadata::ProtocolMetadataV2 {
151        use crate::protocol::metadata::{
152            DevelopmentState, PrivilegeRequirement, ProtocolMetadataV2,
153        };
154
155        ProtocolMetadataV2::builder()
156            .state(DevelopmentState::Experimental)
157            .privilege_requirement(PrivilegeRequirement::None)
158            .implementation(
159                "tokio-tungstenite 0.21 connect_async; client-to-server masking, fragmentation \
160                 and the close handshake are handled by the framing layer",
161            )
162            .llm_control(
163                "Every text and binary frame sent, ping payloads, the close code and reason, \
164                 and which subprotocols to offer at connect time",
165            )
166            .e2e_testing(
167                "The NetGet WebSocket server and websocat 1.14.1 in server mode (-s); RFC 7692 \
168                 permessage-deflate against Python websockets 15 (12-bit windows, its \
169                 default) and Node ws 8 (no context takeover), with a relay reading RSV1 off \
170                 the wire (tests/client/websocket/deflate_real_server_test.rs)",
171            )
172            .notes(
173                "ws:// only — wss:// is not supported because the tokio-tungstenite dependency \
174                 is built without a TLS backend. permessage-deflate is offered by default \
175                 (permessage_deflate) and applied beneath tungstenite. Pong frames are logged \
176                 rather than surfaced as an event. Validated against websocat 1.14.1 running as \
177                 the server.",
178            )
179            .build()
180    }
181
182    fn description(&self) -> &'static str {
183        "WebSocket (RFC 6455) client for connecting to ws:// endpoints"
184    }
185
186    fn example_prompt(&self) -> &'static str {
187        "Connect to the WebSocket at 127.0.0.1:9001/ws, send a subscribe message, and summarise every update"
188    }
189
190    fn group_name(&self) -> &'static str {
191        "Application"
192    }
193
194    fn get_startup_examples(&self) -> crate::llm::actions::StartupExamples {
195        use crate::llm::actions::StartupExamples;
196
197        StartupExamples::new(
198            json!({
199                "type": "open_client",
200                "remote_addr": "127.0.0.1:9001",
201                "base_stack": "websocket",
202                "startup_params": {"path": "/ws", "subprotocols": ["chat"]},
203                "instruction": "Connect, send {\"op\":\"subscribe\",\"channel\":\"ticker\"} once connected, \
204                                and report each update you receive."
205            }),
206            json!({
207                "type": "open_client",
208                "remote_addr": "127.0.0.1:9001",
209                "base_stack": "websocket",
210                "event_handlers": [{
211                    "event_pattern": "websocket_client_text_message",
212                    "handler": {
213                        "type": "script",
214                        "language": "python",
215                        "code": "respond([{'type': 'send_websocket_text', 'text': event['text'].upper()}])"
216                    }
217                }]
218            }),
219            json!({
220                "type": "open_client",
221                "remote_addr": "127.0.0.1:9001",
222                "base_stack": "websocket",
223                "event_handlers": [
224                    {
225                        "event_pattern": "websocket_client_connected",
226                        "handler": {
227                            "type": "static",
228                            "actions": [{"type": "send_websocket_text", "text": "hello"}]
229                        }
230                    },
231                    {
232                        "event_pattern": "websocket_client_text_message",
233                        "handler": {
234                            "type": "static",
235                            "actions": [{"type": "close_websocket", "code": 1000, "reason": "done"}]
236                        }
237                    }
238                ]
239            }),
240        )
241    }
242}
243
244impl Client for WebSocketClientProtocol {
245    fn connect(
246        &self,
247        ctx: crate::protocol::ConnectContext,
248    ) -> std::pin::Pin<
249        Box<dyn std::future::Future<Output = anyhow::Result<std::net::SocketAddr>> + Send>,
250    > {
251        Box::pin(async move {
252            use crate::client::websocket::WebSocketClient;
253            WebSocketClient::connect_with_llm_actions(ctx).await
254        })
255    }
256
257    fn execute_action(&self, action: serde_json::Value) -> Result<ClientActionResult> {
258        let action_type = action
259            .get("type")
260            .and_then(|v| v.as_str())
261            .context("Missing 'type' field in action")?;
262
263        match action_type {
264            "send_websocket_text" => {
265                let text = action
266                    .get("text")
267                    .and_then(|v| v.as_str())
268                    .context("Missing 'text' for send_websocket_text")?
269                    .to_string();
270                debug!("WebSocket client -> text {} bytes", text.len());
271                let bytes = text.len();
272                self.send(WsOut::Text(text))?;
273                Ok(ClientActionResult::Custom {
274                    name: "send_websocket_text".to_string(),
275                    data: json!({ "bytes": bytes }),
276                })
277            }
278            "send_websocket_binary" => {
279                let data = action
280                    .get("data")
281                    .and_then(|v| v.as_str())
282                    .context("Missing 'data' for send_websocket_binary")?;
283                let bytes =
284                    decode_outbound_payload(data, action.get("encoding").and_then(|v| v.as_str()))?;
285                debug!("WebSocket client -> binary {} bytes", bytes.len());
286                let len = bytes.len();
287                self.send(WsOut::Binary(bytes))?;
288                Ok(ClientActionResult::Custom {
289                    name: "send_websocket_binary".to_string(),
290                    data: json!({ "bytes": len }),
291                })
292            }
293            "send_websocket_ping" => {
294                let payload = action
295                    .get("payload")
296                    .and_then(|v| v.as_str())
297                    .unwrap_or("")
298                    .as_bytes()
299                    .to_vec();
300                if payload.len() > 125 {
301                    return Err(anyhow::anyhow!(
302                        "Ping payload is {} bytes; RFC 6455 §5.5 limits control frames to 125",
303                        payload.len()
304                    ));
305                }
306                self.send(WsOut::Ping(payload))?;
307                Ok(ClientActionResult::Custom {
308                    name: "send_websocket_ping".to_string(),
309                    data: json!({}),
310                })
311            }
312            "close_websocket" => {
313                let code = action.get("code").and_then(|v| v.as_u64()).unwrap_or(1000);
314                let valid = (1000..=1003).contains(&code)
315                    || (1007..=1011).contains(&code)
316                    || (3000..=4999).contains(&code);
317                if !valid {
318                    return Err(anyhow::anyhow!(
319                        "close code {} is not sendable; use 1000-1003, 1007-1011 or 3000-4999 \
320                         (1005, 1006 and 1015 are reserved)",
321                        code
322                    ));
323                }
324                let reason = action
325                    .get("reason")
326                    .and_then(|v| v.as_str())
327                    .unwrap_or("")
328                    .to_string();
329                if reason.len() > 123 {
330                    return Err(anyhow::anyhow!(
331                        "close reason is {} bytes; the limit is 123",
332                        reason.len()
333                    ));
334                }
335                self.send(WsOut::Close {
336                    code: code as u16,
337                    reason,
338                })?;
339                Ok(ClientActionResult::Disconnect)
340            }
341            "disconnect" => {
342                // Drop the TCP connection without a closing handshake. `close_websocket` is
343                // the polite form and should be preferred.
344                Ok(ClientActionResult::Disconnect)
345            }
346            "wait_for_websocket_data" => Ok(ClientActionResult::WaitForMore),
347            _ => Err(anyhow::anyhow!(
348                "Unknown WebSocket client action: {}",
349                action_type
350            )),
351        }
352    }
353}
354
355// ============================================================================
356// Action definitions
357// ============================================================================
358
359fn client_actions() -> Vec<ActionDefinition> {
360    vec![
361        send_text_action(),
362        send_binary_action(),
363        send_ping_action(),
364        wait_for_data_action(),
365        close_action(),
366        disconnect_action(),
367    ]
368}
369
370fn send_text_action() -> ActionDefinition {
371    ActionDefinition {
372        name: "send_websocket_text".to_string(),
373        description: "Send a WebSocket text frame to the server.".to_string(),
374        parameters: vec![Parameter {
375            name: "text".to_string(),
376            type_hint: "string".to_string(),
377            description: "The message body, sent as UTF-8. Put serialised JSON here as a string."
378                .to_string(),
379            required: true,
380        }],
381        example: json!({"type": "send_websocket_text", "text": "{\"op\":\"subscribe\"}"}),
382        log_template: Some(
383            LogTemplate::new()
384                .with_info("-> WS text")
385                .with_trace("WebSocket client text: {preview(text,200)}"),
386        ),
387    }
388}
389
390fn send_binary_action() -> ActionDefinition {
391    ActionDefinition {
392        name: "send_websocket_binary".to_string(),
393        description:
394            "Send a WebSocket binary frame to the server. 'data' plus 'encoding' is the same \
395             pair a websocket_client_binary_message event carries, so echoing is passing both \
396             fields straight through."
397                .to_string(),
398        parameters: vec![
399            Parameter {
400                name: "data".to_string(),
401                type_hint: "string".to_string(),
402                description: "The payload, interpreted according to 'encoding'.".to_string(),
403                required: true,
404            },
405            Parameter {
406                name: "encoding".to_string(),
407                type_hint: "string".to_string(),
408                description:
409                    "\"utf8\" (default) sends the characters of 'data' unchanged; \"base64\" \
410                     and \"hex\" decode 'data' into bytes. There is no auto-detection."
411                        .to_string(),
412                required: false,
413            },
414        ],
415        example: json!({"type": "send_websocket_binary", "data": "AP/+AQ==", "encoding": "base64"}),
416        log_template: Some(LogTemplate::new().with_info("-> WS binary")),
417    }
418}
419
420fn send_ping_action() -> ActionDefinition {
421    ActionDefinition {
422        name: "send_websocket_ping".to_string(),
423        description:
424            "Send a ping control frame. The server's pong is logged and does not raise an event."
425                .to_string(),
426        parameters: vec![Parameter {
427            name: "payload".to_string(),
428            type_hint: "string".to_string(),
429            description: "Optional payload, at most 125 bytes.".to_string(),
430            required: false,
431        }],
432        example: json!({"type": "send_websocket_ping", "payload": "keepalive"}),
433        log_template: Some(LogTemplate::new().with_debug("WebSocket client ping")),
434    }
435}
436
437fn wait_for_data_action() -> ActionDefinition {
438    ActionDefinition {
439        name: "wait_for_websocket_data".to_string(),
440        description:
441            "Do not act on this message yet; hold it and join the next message of the same kind \
442             onto it before deciding."
443                .to_string(),
444        parameters: vec![],
445        example: json!({"type": "wait_for_websocket_data"}),
446        log_template: Some(LogTemplate::new().with_debug("WebSocket client waiting for more")),
447    }
448}
449
450fn close_action() -> ActionDefinition {
451    ActionDefinition {
452        name: "close_websocket".to_string(),
453        description:
454            "Close politely: send a close frame with a status code and let the server answer."
455                .to_string(),
456        parameters: vec![
457            Parameter {
458                name: "code".to_string(),
459                type_hint: "integer".to_string(),
460                description:
461                    "Close status code (default 1000, normal). 1001 going away, 1008 policy \
462                     violation, 1011 internal error, or 3000-4999 for your own meanings."
463                        .to_string(),
464                required: false,
465            },
466            Parameter {
467                name: "reason".to_string(),
468                type_hint: "string".to_string(),
469                description: "Optional explanation, at most 123 bytes.".to_string(),
470                required: false,
471            },
472        ],
473        example: json!({"type": "close_websocket", "code": 1000, "reason": "done"}),
474        log_template: Some(LogTemplate::new().with_info("WS client close {code}")),
475    }
476}
477
478fn disconnect_action() -> ActionDefinition {
479    ActionDefinition {
480        name: "disconnect".to_string(),
481        description:
482            "Drop the connection immediately without a closing handshake. Prefer close_websocket."
483                .to_string(),
484        parameters: vec![],
485        example: json!({"type": "disconnect"}),
486        log_template: Some(LogTemplate::new().with_info("WS client disconnect")),
487    }
488}
489
490// ============================================================================
491// Event types — every one of these is emitted by src/client/websocket/mod.rs
492// ============================================================================
493
494fn connected_actions() -> Vec<ActionDefinition> {
495    vec![
496        send_text_action(),
497        send_binary_action(),
498        send_ping_action(),
499        close_action(),
500        disconnect_action(),
501    ]
502}
503
504pub static WEBSOCKET_CLIENT_CONNECTED_EVENT: LazyLock<EventType> = LazyLock::new(|| {
505    EventType::new(
506        "websocket_client_connected",
507        "The WebSocket handshake completed. Send the first message now if the protocol expects \
508         the client to speak first.",
509        json!({"type": "send_websocket_text", "text": "{\"op\":\"subscribe\"}"}),
510    )
511    .with_parameters(vec![
512        Parameter {
513            name: "remote_addr".to_string(),
514            type_hint: "string".to_string(),
515            description: "Address of the server that was connected to".to_string(),
516            required: true,
517        },
518        Parameter {
519            name: "path".to_string(),
520            type_hint: "string".to_string(),
521            description: "Request path the connection was opened on".to_string(),
522            required: true,
523        },
524        Parameter {
525            name: "subprotocol".to_string(),
526            type_hint: "string".to_string(),
527            description:
528                "The subprotocol the server agreed on, or empty if it chose none. Never a value \
529                 that was not offered."
530                    .to_string(),
531            required: false,
532        },
533        Parameter {
534            name: "permessage_deflate".to_string(),
535            type_hint: "object".to_string(),
536            description: "The RFC 7692 compression parameters the server accepted, or null when \
537                 the connection is uncompressed"
538                .to_string(),
539            required: false,
540        },
541    ])
542    .with_actions(connected_actions())
543    .with_log_template(
544        LogTemplate::new()
545            .with_info("WS client connected to {remote_addr}{path}")
546            .with_debug("WebSocket client connected subprotocol={subprotocol}"),
547    )
548});
549
550pub static WEBSOCKET_CLIENT_TEXT_MESSAGE_EVENT: LazyLock<EventType> = LazyLock::new(|| {
551    EventType::new(
552        "websocket_client_text_message",
553        "A text message arrived from the server.",
554        json!({"type": "send_websocket_text", "text": "ack"}),
555    )
556    .with_parameters(vec![
557        Parameter {
558            name: "text".to_string(),
559            type_hint: "string".to_string(),
560            description: "The complete message, reassembled from any fragments".to_string(),
561            required: true,
562        },
563        Parameter {
564            name: "message_bytes".to_string(),
565            type_hint: "integer".to_string(),
566            description: "Length of the message in bytes".to_string(),
567            required: true,
568        },
569    ])
570    .with_actions({
571        let mut a = connected_actions();
572        a.push(wait_for_data_action());
573        a
574    })
575    .with_alternative_example(json!({"type": "close_websocket", "code": 1000, "reason": "done"}))
576    .with_log_template(
577        LogTemplate::new()
578            .with_info("WS client <- text {message_bytes}B")
579            .with_trace("WebSocket client text: {preview(text,200)}"),
580    )
581});
582
583pub static WEBSOCKET_CLIENT_BINARY_MESSAGE_EVENT: LazyLock<EventType> = LazyLock::new(|| {
584    EventType::new(
585        "websocket_client_binary_message",
586        "A binary message arrived from the server. To echo it back unchanged, pass this event's \
587         'data' AND its 'encoding' straight into send_websocket_binary.",
588        json!({"type": "send_websocket_binary", "data": "AP/+AQ==", "encoding": "base64"}),
589    )
590    .with_parameters(vec![
591        Parameter {
592            name: "data".to_string(),
593            type_hint: "string".to_string(),
594            description: "The payload, read according to the 'encoding' field".to_string(),
595            required: true,
596        },
597        Parameter {
598            name: "encoding".to_string(),
599            type_hint: "string".to_string(),
600            description: "\"utf8\" when every byte is printable ASCII, otherwise \"base64\". \
601                 send_websocket_binary accepts the same values, so the pair round-trips."
602                .to_string(),
603            required: true,
604        },
605        Parameter {
606            name: "message_bytes".to_string(),
607            type_hint: "integer".to_string(),
608            description: "Length of the decoded message in bytes".to_string(),
609            required: true,
610        },
611    ])
612    .with_actions({
613        let mut a = connected_actions();
614        a.push(wait_for_data_action());
615        a
616    })
617    .with_log_template(LogTemplate::new().with_info("WS client <- binary {message_bytes}B"))
618});
619
620pub static WEBSOCKET_CLIENT_CLOSED_EVENT: LazyLock<EventType> = LazyLock::new(|| {
621    EventType::new(
622        "websocket_client_closed",
623        "The server closed the connection. Record what happened; nothing more can be sent.",
624        json!({"type": "show_message", "message": "server closed the WebSocket"}),
625    )
626    .with_parameters(vec![
627        Parameter {
628            name: "code".to_string(),
629            type_hint: "integer".to_string(),
630            description: "Close status code, or 1005 when the server sent none".to_string(),
631            required: true,
632        },
633        Parameter {
634            name: "reason".to_string(),
635            type_hint: "string".to_string(),
636            description: "Close reason text, or empty".to_string(),
637            required: false,
638        },
639    ])
640    // Nothing can go on the wire after the closing handshake, so the common actions
641    // (show_message, set_memory, append_to_log) are the only honest vocabulary here.
642    .with_no_actions()
643    .with_log_template(LogTemplate::new().with_info("WS client closed code={code}"))
644});
645
646pub fn get_websocket_client_event_types() -> Vec<EventType> {
647    vec![
648        WEBSOCKET_CLIENT_CONNECTED_EVENT.clone(),
649        WEBSOCKET_CLIENT_TEXT_MESSAGE_EVENT.clone(),
650        WEBSOCKET_CLIENT_BINARY_MESSAGE_EVENT.clone(),
651        WEBSOCKET_CLIENT_CLOSED_EVENT.clone(),
652    ]
653}
netget/client/websocket/actions.rs

netget/client/websocket/
actions.rs
