Skip to main content

netget/client/sse/
actions.rs

1//! What the model does as a server-sent-events client: hear each event of a stream, and decide
2//! whether to reconnect or stop. Rust owns the HTTP request, the event-stream parsing, and the
3//! reconnects with Last-Event-ID that EventSource makes.
4use crate::llm::actions::{
5    client_trait::{Client, ClientActionResult},
6    protocol_trait::Protocol,
7    ActionDefinition, ParameterDefinition, StartupExamples,
8};
9use crate::protocol::{ConnectContext, EventType};
10use crate::server::sse::actions::{action, p};
11use crate::state::app_state::AppState;
12use anyhow::{bail, ensure, Result};
13use serde_json::{json, Value};
14use std::sync::LazyLock;
15
16pub const RECONNECT: &str = "sse_reconnect";
17
18#[derive(Default)]
19pub struct SseClientProtocol;
20impl SseClientProtocol {
21    pub fn new() -> Self {
22        Self
23    }
24}
25
26pub fn actions() -> Vec<ActionDefinition> {
27    vec![
28        action(
29            RECONNECT,
30            "Drop the stream and open it again now, sending Last-Event-ID so the server can resume",
31            vec![p(
32                "last_event_id",
33                "string",
34                "The Last-Event-ID to send (default the last event ID seen)",
35                false,
36            )],
37            json!({"type": RECONNECT}),
38        ),
39        action(
40            "disconnect",
41            "Close the stream and stop; no further reconnects",
42            vec![],
43            json!({"type": "disconnect"}),
44        ),
45    ]
46}
47
48fn event(id: &str, description: &str, params: Vec<crate::llm::actions::Parameter>) -> EventType {
49    EventType::new(id, description, json!({"type": "disconnect"}))
50        .with_parameters(params)
51        .with_actions(actions())
52}
53
54pub static CONNECTED_EVENT: LazyLock<EventType> = LazyLock::new(|| {
55    event(
56        "sse_connected",
57        "The stream opened (200 text/event-stream).",
58        vec![
59            p(
60                "url",
61                "string",
62                "The stream URL this connection requested",
63                true,
64            ),
65            p(
66                "last_event_id",
67                "string",
68                "The Last-Event-ID this connection sent (empty on the first)",
69                true,
70            ),
71            p(
72                "attempt",
73                "number",
74                "Which connection this is: 1 for the first, more after reconnects",
75                true,
76            ),
77        ],
78    )
79});
80
81pub static MESSAGE_EVENT: LazyLock<EventType> = LazyLock::new(|| {
82    event(
83        "sse_event_received",
84        "An event arrived on the stream.",
85        vec![
86            p(
87                "event",
88                "string",
89                "The event type (message when the server named none)",
90                true,
91            ),
92            p(
93                "data",
94                "string",
95                "The event's data, lines joined with newlines",
96                true,
97            ),
98            p(
99                "id",
100                "string",
101                "The last event ID after this event (sent back as Last-Event-ID on a reconnect)",
102                true,
103            ),
104        ],
105    )
106});
107
108pub static ENDED_EVENT: LazyLock<EventType> = LazyLock::new(|| {
109    event(
110        "sse_stream_ended",
111        "The stream stopped: it closed, the server refused it, or reconnects ran out.",
112        vec![
113            p(
114                "reason",
115                "string",
116                "closed, refused, error or no_more_reconnects",
117                true,
118            ),
119            p(
120                "status",
121                "number",
122                "The HTTP status when the server refused (204 means stop for good)",
123                false,
124            ),
125            p("body", "string", "The start of a refusal's body", false),
126            p("error", "string", "What failed, for an error", false),
127            p(
128                "reconnecting_in_ms",
129                "number",
130                "Set when the client will reconnect by itself after this long",
131                false,
132            ),
133        ],
134    )
135});
136
137pub fn check(v: &Value) -> Result<()> {
138    match v["type"].as_str().unwrap_or_default() {
139        RECONNECT => {
140            if let Some(id) = v.get("last_event_id").filter(|x| !x.is_null()) {
141                let id = id.as_str().unwrap_or_default();
142                ensure!(
143                    id.len() <= 1024 && !id.contains(['\r', '\n', '\0']),
144                    "last_event_id is one line"
145                );
146            }
147        }
148        other => bail!("Unknown SSE client action {other:?}"),
149    }
150    Ok(())
151}
152
153impl Protocol for SseClientProtocol {
154    fn protocol_name(&self) -> &'static str {
155        "SSE"
156    }
157    fn stack_name(&self) -> &'static str {
158        "ETH>IP>TCP>HTTP>SSE"
159    }
160    fn keywords(&self) -> Vec<&'static str> {
161        vec![
162            "sse client",
163            "eventsource",
164            "server-sent events",
165            "event stream",
166        ]
167    }
168    fn description(&self) -> &'static str {
169        "Server-sent events client: reads an event stream like EventSource, reconnecting with Last-Event-ID"
170    }
171    fn get_async_actions(&self, _: &AppState) -> Vec<ActionDefinition> {
172        actions()
173    }
174    fn get_sync_actions(&self) -> Vec<ActionDefinition> {
175        vec![]
176    }
177    fn get_event_types(&self) -> Vec<EventType> {
178        vec![
179            CONNECTED_EVENT.clone(),
180            MESSAGE_EVENT.clone(),
181            ENDED_EVENT.clone(),
182        ]
183    }
184    fn get_startup_parameters(&self) -> Vec<ParameterDefinition> {
185        vec![
186            ParameterDefinition {
187                name: "last_event_id".into(),
188                type_hint: "string".into(),
189                description: "Last-Event-ID to send on the first connection, to resume a stream"
190                    .into(),
191                required: false,
192                example: json!("42"),
193                default: None,
194            },
195            ParameterDefinition {
196                name: "max_reconnects".into(),
197                type_hint: "number".into(),
198                description: "How many times (0..=100) a closed stream is reopened by itself"
199                    .into(),
200                required: false,
201                example: json!(5),
202                default: Some(json!(super::MAX_RECONNECTS)),
203            },
204        ]
205    }
206    fn metadata(&self) -> crate::protocol::metadata::ProtocolMetadataV2 {
207        use crate::protocol::metadata::*;
208        ProtocolMetadataV2::builder()
209            .state(DevelopmentState::Experimental)
210            .privilege_requirement(PrivilegeRequirement::None)
211            .implementation("Hand-rolled HTTP/1.1 GET over Tokio TCP (Content-Length, chunked and close-delimited bodies), the WHATWG event-stream parser shared with the SSE server, and EventSource's reconnects: the server's retry hint, Last-Event-ID, no reconnect after a 204 or any other refusal")
212            .llm_control("What to make of each event, and whether to reconnect or stop")
213            .e2e_testing("tests/client/sse: sse-starlette (Starlette under uvicorn) and aiohttp-sse streams, each closing mid-stream so the reconnect's Last-Event-ID is checked on the server")
214            .notes("http:// only (no TLS), no redirects, no request headers beyond Accept, Cache-Control and Last-Event-ID. Lines are capped at 64 KiB and an event's data at 1 MiB. A handler chain stops after 8 follow-ups.")
215            .max_inbound_bytes(crate::server::sse::wire::MAX_EVENT_BYTES)
216            .build()
217    }
218    fn example_prompt(&self) -> &'static str {
219        "Subscribe to http://127.0.0.1:8080/news and tell me each headline as it arrives"
220    }
221    fn get_startup_examples(&self) -> StartupExamples {
222        let llm = json!({"type":"open_client","protocol":"sse","remote_addr":"http://127.0.0.1:8080/news",
223            "instruction":"Report each event; disconnect after an event of type bye"});
224        let mut static_example = llm.clone();
225        static_example["event_handlers"] =
226            json!([{"event_pattern":"*","handler":{"type":"static","actions":[]}}]);
227        let mut scripted = llm.clone();
228        scripted["event_handlers"] = json!([{"event_pattern":"*","handler":{"type":"script","language":"python",
229            "code":"import json,sys\ni=json.load(sys.stdin)\nbye=i['event_type_id']=='sse_event_received' and i['event']['event']=='bye'\nprint(json.dumps({'actions':[{'type':'disconnect'}] if bye else []}))"}}]);
230        StartupExamples::new(llm, scripted, static_example)
231    }
232    fn group_name(&self) -> &'static str {
233        "Web & File"
234    }
235}
236
237impl Client for SseClientProtocol {
238    fn connect(
239        &self,
240        ctx: ConnectContext,
241    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<std::net::SocketAddr>> + Send>>
242    {
243        Box::pin(super::connect(ctx))
244    }
245    fn execute_action(&self, v: Value) -> Result<ClientActionResult> {
246        let name = v["type"].as_str().unwrap_or_default().to_string();
247        if name == "disconnect" {
248            return Ok(ClientActionResult::Disconnect);
249        }
250        check(&v)?;
251        Ok(ClientActionResult::Custom { name, data: v })
252    }
253}