Skip to main content

netget/client/graphql/
mod.rs

1//! GraphQL over HTTP client.
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::graphql::{engine, ws, GRAPHQL_RESPONSE};
9use crate::state::client_handles::{ClientCommand, ClientSendOutcome};
10use crate::state::{AccessLogOwner, ClientStatus};
11pub use actions::GraphqlClientProtocol;
12use anyhow::{ensure, Context, Result};
13use apollo_compiler::executable::OperationType;
14use futures::{SinkExt, StreamExt};
15use serde_json::{json, Value};
16use std::collections::HashMap;
17use std::net::SocketAddr;
18use std::time::Duration;
19use tokio::sync::mpsc;
20use tokio_tungstenite::tungstenite::{client::IntoClientRequest, Message};
21
22pub const DEFAULT_INTROSPECT: bool = true;
23const TIMEOUT: Duration = Duration::from_secs(30);
24const ACCEPT: &str = "application/graphql-response+json, application/json;q=0.9";
25
26pub async fn connect(ctx: ConnectContext) -> Result<SocketAddr> {
27    let p = ctx.startup_params.as_ref();
28    let endpoint = p
29        .map(|p| p.get_optional_string("endpoint"))
30        .transpose()?
31        .flatten()
32        .unwrap_or_else(|| crate::server::graphql::DEFAULT_ENDPOINT.to_owned());
33    ensure!(
34        endpoint.starts_with('/') && endpoint.len() <= 256 && !endpoint.contains(['#', ' ']),
35        "endpoint must be an absolute path"
36    );
37    let introspect = p
38        .map(|p| p.get_optional_bool("introspect"))
39        .transpose()?
40        .flatten()
41        .unwrap_or(DEFAULT_INTROSPECT);
42    let url = format!("http://{}{endpoint}", ctx.remote_addr);
43    crate::client::http_fetch::check_url(&url)?;
44    #[cfg(not(target_arch = "wasm32"))]
45    let fetch = FetchClient::from_reqwest(
46        crate::llm::ollama_client::configured_for_endpoint(
47            reqwest::Client::builder()
48                .timeout(TIMEOUT)
49                .redirect(reqwest::redirect::Policy::none()),
50            &url,
51        )
52        .build()?,
53    );
54    #[cfg(target_arch = "wasm32")]
55    let fetch = FetchClient::transport(TIMEOUT);
56    let fetch = fetch
57        .with_max_body(engine::MAX_BODY_BYTES)
58        .with_user_agent("netget-graphql");
59    let mut connected = json!({"url": url});
60    if introspect {
61        match exchange(
62            &fetch,
63            &url,
64            engine::CLIENT_INTROSPECTION,
65            None,
66            &json!({}),
67            false,
68        )
69        .await
70        {
71            Ok(r) if r["data"].is_object() && r.get("errors").is_none() => {
72                connected["root_fields"] = engine::root_signatures(&r["data"]);
73            }
74            Ok(r) => {
75                connected["introspection_error"] = json!(r
76                    .get("error")
77                    .or_else(|| r["errors"].get(0).map(|e| &e["message"]))
78                    .and_then(Value::as_str)
79                    .unwrap_or("introspection gave no schema"))
80            }
81            Err(e) => bail_connect(&e)?,
82        }
83    }
84    let local: SocketAddr = "0.0.0.0:0".parse()?;
85    ctx.state
86        .update_client_status(ctx.client_id, ClientStatus::Connected)
87        .await;
88    let external =
89        crate::client::command_support::register_command_channel(&ctx.state, ctx.client_id).await;
90    let (internal_tx, internal_rx) = mpsc::channel::<Value>(16);
91    let (event_tx, mut event_rx) = mpsc::channel::<Event>(16);
92    event_tx.try_send(Event::new(&actions::CONNECTED_EVENT, connected))?;
93    let events_ctx = ctx.clone();
94    let dispatcher = tokio::spawn(async move {
95        let protocol = GraphqlClientProtocol;
96        while let Some(event) = event_rx.recv().await {
97            let instruction = events_ctx
98                .state
99                .get_instruction_for_client(events_ctx.client_id)
100                .await
101                .unwrap_or_default();
102            let memory = events_ctx
103                .state
104                .get_memory_for_client(events_ctx.client_id)
105                .await
106                .unwrap_or_default();
107            match call_llm_for_client(
108                &events_ctx.llm_client,
109                &events_ctx.state,
110                events_ctx.client_id.to_string(),
111                &instruction,
112                &memory,
113                Some(&event),
114                &protocol,
115                &events_ctx.status_tx,
116            )
117            .await
118            {
119                Ok(result) => {
120                    if let Some(memory) = result.memory_updates {
121                        events_ctx
122                            .state
123                            .set_memory_for_client(events_ctx.client_id, memory)
124                            .await;
125                    }
126                    for action in result.actions {
127                        if internal_tx.send(action).await.is_err() {
128                            return;
129                        }
130                    }
131                }
132                Err(e) => Log::new(Some(&events_ctx.status_tx))
133                    .warn(format!("GraphQL client handler: {e}")),
134            }
135        }
136    });
137    ctx.state
138        .register_client_task(ctx.client_id, dispatcher)
139        .await;
140    let session_ctx = ctx.clone();
141    let task = tokio::spawn(async move {
142        let result = session(&session_ctx, &fetch, &url, external, internal_rx, event_tx).await;
143        let status = match result {
144            Ok(()) => ClientStatus::Disconnected,
145            Err(e) => {
146                Log::new(Some(&session_ctx.status_tx)).warn(format!("GraphQL client ended: {e}"));
147                ClientStatus::Error(e.to_string())
148            }
149        };
150        session_ctx
151            .state
152            .update_client_status(session_ctx.client_id, status)
153            .await;
154        session_ctx
155            .state
156            .remove_client_handle(session_ctx.client_id)
157            .await;
158        let _ = session_ctx.status_tx.send("__UPDATE_UI__".into());
159    });
160    ctx.state.register_client_task(ctx.client_id, task).await;
161    Ok(local)
162}
163
164/// A transport failure on the first request means there is no server to talk to.
165fn bail_connect(e: &anyhow::Error) -> Result<()> {
166    Err(anyhow::anyhow!("GraphQL endpoint unreachable: {e}"))
167}
168
169type WsStream =
170    tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>;
171
172/// The graphql-transport-ws socket, opened on the first subscription.
173struct Socket {
174    sink: futures::stream::SplitSink<WsStream, Message>,
175    stream: futures::stream::SplitStream<WsStream>,
176    /// Active subscriptions: id → operation name.
177    active: HashMap<String, Option<String>>,
178    next_id: u64,
179}
180
181async fn send_ws(sink: &mut futures::stream::SplitSink<WsStream, Message>, v: Value) -> Result<()> {
182    sink.send(Message::Text(v.to_string())).await?;
183    Ok(())
184}
185
186/// Connect with the graphql-transport-ws subprotocol and complete connection_init/ack.
187async fn open_socket(url: &str) -> Result<Socket> {
188    let ws_url = url.replacen("http://", "ws://", 1);
189    let mut request = ws_url.as_str().into_client_request()?;
190    request.headers_mut().insert(
191        "Sec-WebSocket-Protocol",
192        tokio_tungstenite::tungstenite::http::HeaderValue::from_static(ws::SUBPROTOCOL),
193    );
194    let (socket, response) = tokio::time::timeout(
195        TIMEOUT,
196        tokio_tungstenite::connect_async_with_config(request, Some(ws::ws_config()), false),
197    )
198    .await
199    .context("WebSocket connect timed out")??;
200    ensure!(
201        response
202            .headers()
203            .get("sec-websocket-protocol")
204            .and_then(|v| v.to_str().ok())
205            == Some(ws::SUBPROTOCOL),
206        "the server did not agree to graphql-transport-ws"
207    );
208    let (mut sink, mut stream) = socket.split();
209    send_ws(&mut sink, json!({"type": "connection_init", "payload": {}})).await?;
210    let deadline = tokio::time::Instant::now() + ws::CONNECTION_INIT_TIMEOUT;
211    loop {
212        let message = tokio::time::timeout_at(deadline, stream.next())
213            .await
214            .context("no connection_ack in time")?
215            .context("socket closed before connection_ack")??;
216        let Message::Text(text) = message else {
217            continue;
218        };
219        let v: Value = serde_json::from_str(&text).context("server sent a non-JSON message")?;
220        match v["type"].as_str() {
221            Some("connection_ack") => break,
222            Some("ping") => send_ws(&mut sink, json!({"type": "pong"})).await?,
223            Some("pong") => {}
224            other => anyhow::bail!("expected connection_ack, got {other:?}"),
225        }
226    }
227    Ok(Socket {
228        sink,
229        stream,
230        active: HashMap::new(),
231        next_id: 1,
232    })
233}
234
235enum Step {
236    Command(Option<ClientCommand>),
237    Action(Option<Value>),
238    Socket(Option<Result<Message, tokio_tungstenite::tungstenite::Error>>),
239}
240
241fn reply(command: Option<ClientCommand>, outcome: Result<ClientSendOutcome>) {
242    if let Some(c) = command {
243        crate::client::command_support::reply(c, outcome);
244    }
245}
246
247fn rejected(error: impl ToString) -> Result<ClientSendOutcome> {
248    Ok(ClientSendOutcome::Rejected {
249        error: error.to_string(),
250    })
251}
252
253async fn session(
254    ctx: &ConnectContext,
255    fetch: &FetchClient,
256    url: &str,
257    mut external: mpsc::Receiver<ClientCommand>,
258    mut internal: mpsc::Receiver<Value>,
259    events: mpsc::Sender<Event>,
260) -> Result<()> {
261    let mut socket: Option<Socket> = None;
262    loop {
263        let step = {
264            let incoming = async {
265                match socket.as_mut() {
266                    Some(s) => s.stream.next().await,
267                    None => std::future::pending().await,
268                }
269            };
270            tokio::select! {
271                c = external.recv() => Step::Command(c),
272                a = internal.recv() => Step::Action(a),
273                m = incoming => Step::Socket(m),
274            }
275        };
276        let (action, command) = match step {
277            Step::Command(Some(c)) => (c.action.clone(), Some(c)),
278            Step::Action(Some(a)) => (a, None),
279            Step::Command(None) | Step::Action(None) => return Ok(()),
280            Step::Socket(message) => {
281                if !socket_message(ctx, &mut socket, message, &events).await? {
282                    socket = None;
283                }
284                continue;
285            }
286        };
287        match GraphqlClientProtocol.execute_action(action.clone()) {
288            Ok(ClientActionResult::Disconnect) => {
289                if let Some(mut s) = socket.take() {
290                    let _ = s.sink.close().await;
291                }
292                reply(command, Ok(ClientSendOutcome::Disconnected));
293                return Ok(());
294            }
295            Ok(_) => {}
296            Err(e) => {
297                reply(command, rejected(e));
298                continue;
299            }
300        }
301        match action["type"].as_str() {
302            Some("graphql_subscribe") => {
303                let outcome = subscribe(url, &mut socket, &action).await;
304                reply(command, outcome);
305            }
306            Some("graphql_unsubscribe") => {
307                let sub = action["subscription_id"].as_str().unwrap_or_default();
308                let outcome = match socket.as_mut() {
309                    Some(s) if s.active.contains_key(sub) => {
310                        s.active.remove(sub);
311                        send_ws(&mut s.sink, json!({"type": "complete", "id": sub}))
312                            .await
313                            .map(|()| ClientSendOutcome::Sent {
314                                bytes_sent: sub.len(),
315                            })
316                    }
317                    _ => rejected(format!("no active subscription {sub}")),
318                };
319                reply(command, outcome);
320            }
321            _ => query(ctx, fetch, url, &action, command, &events).await?,
322        }
323    }
324}
325
326async fn subscribe(
327    url: &str,
328    socket: &mut Option<Socket>,
329    action: &Value,
330) -> Result<ClientSendOutcome> {
331    if socket.is_none() {
332        match open_socket(url).await {
333            Ok(s) => *socket = Some(s),
334            Err(e) => return rejected(format!("graphql-transport-ws connection failed: {e}")),
335        }
336    }
337    let s = socket.as_mut().expect("opened above");
338    if s.active.len() >= ws::MAX_SUBSCRIPTIONS {
339        return rejected(format!("at most {} subscriptions", ws::MAX_SUBSCRIPTIONS));
340    }
341    let sub = s.next_id.to_string();
342    s.next_id += 1;
343    let mut payload = json!({"query": action["query"]});
344    if let Some(n) = action["operation_name"].as_str() {
345        payload["operationName"] = json!(n);
346    }
347    if action["variables"]
348        .as_object()
349        .is_some_and(|v| !v.is_empty())
350    {
351        payload["variables"] = action["variables"].clone();
352    }
353    let message = json!({"type": "subscribe", "id": sub, "payload": payload});
354    let n = message.to_string().len();
355    send_ws(&mut s.sink, message).await?;
356    let (_, name) = engine::parse_operation(
357        action["query"].as_str().unwrap_or_default(),
358        action["operation_name"].as_str(),
359    )?;
360    s.active.insert(sub, name);
361    Ok(ClientSendOutcome::Sent { bytes_sent: n })
362}
363
364/// Handle one socket message. `Ok(false)` means the socket is gone; every subscription still
365/// active hears about it as an error event.
366async fn socket_message(
367    ctx: &ConnectContext,
368    socket: &mut Option<Socket>,
369    message: Option<Result<Message, tokio_tungstenite::tungstenite::Error>>,
370    events: &mpsc::Sender<Event>,
371) -> Result<bool> {
372    let s = socket.as_mut().expect("only polled while open");
373    let closed = |why: String| async move {
374        Log::new(Some(&ctx.status_tx)).warn(format!("GraphQL WebSocket: {why}"));
375        why
376    };
377    let text = match message {
378        Some(Ok(Message::Text(t))) => t,
379        Some(Ok(Message::Ping(_) | Message::Pong(_) | Message::Frame(_))) => return Ok(true),
380        Some(Ok(Message::Close(frame))) => {
381            let why = closed(match frame {
382                Some(f) => format!("closed by the server ({} {})", u16::from(f.code), f.reason),
383                None => "closed by the server".into(),
384            })
385            .await;
386            return end_all(s, &why, events).await.map(|()| false);
387        }
388        Some(Ok(Message::Binary(_))) => {
389            let why = closed("binary frame from the server".into()).await;
390            return end_all(s, &why, events).await.map(|()| false);
391        }
392        None | Some(Err(_)) => {
393            let why = closed("connection lost".into()).await;
394            return end_all(s, &why, events).await.map(|()| false);
395        }
396    };
397    let v: Value = match serde_json::from_str(&text) {
398        Ok(v @ Value::Object(_)) if engine::budget_ok(&v) => v,
399        _ => {
400            let why = closed("unparsable message from the server".into()).await;
401            let _ = s.sink.close().await;
402            return end_all(s, &why, events).await.map(|()| false);
403        }
404    };
405    let sub = v["id"].as_str().unwrap_or_default().to_owned();
406    let event = match v["type"].as_str() {
407        Some("ping") => {
408            send_ws(&mut s.sink, json!({"type": "pong"})).await?;
409            return Ok(true);
410        }
411        Some("pong") => return Ok(true),
412        Some(kind @ ("next" | "error" | "complete")) => {
413            let Some(name) = s.active.get(&sub).cloned() else {
414                // A message for a subscription we ended: allowed to race, ignored.
415                return Ok(true);
416            };
417            match kind {
418                "next" => {
419                    if let Err(e) = engine::check_response(&v["payload"]) {
420                        Log::new(Some(&ctx.status_tx))
421                            .warn(format!("GraphQL subscription {sub}: bad next payload: {e}"));
422                        return Ok(true);
423                    }
424                    let mut data = json!({"subscription_id": sub, "operation_name": name});
425                    for k in ["data", "errors"] {
426                        if let Some(x) = v["payload"].get(k) {
427                            data[k] = x.clone();
428                        }
429                    }
430                    Event::new(&actions::SUBSCRIPTION_EVENT, data)
431                }
432                "error" => {
433                    s.active.remove(&sub);
434                    let errors = v["payload"]
435                        .as_array()
436                        .filter(|e| !e.is_empty())
437                        .cloned()
438                        .unwrap_or_else(|| vec![json!({"message": "error without details"})]);
439                    Event::new(
440                        &actions::SUBSCRIPTION_ERROR_EVENT,
441                        json!({"subscription_id": sub, "operation_name": name, "errors": errors}),
442                    )
443                }
444                _ => {
445                    s.active.remove(&sub);
446                    Event::new(
447                        &actions::SUBSCRIPTION_COMPLETE_EVENT,
448                        json!({"subscription_id": sub, "operation_name": name}),
449                    )
450                }
451            }
452        }
453        _ => {
454            let why = closed(format!("unexpected message type {:?}", v["type"])).await;
455            let _ = s.sink.close().await;
456            return end_all(s, &why, events).await.map(|()| false);
457        }
458    };
459    events
460        .send(event)
461        .await
462        .context("GraphQL event consumer stopped")?;
463    Ok(true)
464}
465
466async fn end_all(s: &mut Socket, why: &str, events: &mpsc::Sender<Event>) -> Result<()> {
467    for (sub, name) in s.active.drain() {
468        events
469            .send(Event::new(
470                &actions::SUBSCRIPTION_ERROR_EVENT,
471                json!({"subscription_id": sub, "operation_name": name, "errors": [{"message": format!("WebSocket {why}")}]}),
472            ))
473            .await
474            .context("GraphQL event consumer stopped")?;
475    }
476    Ok(())
477}
478
479async fn query(
480    ctx: &ConnectContext,
481    fetch: &FetchClient,
482    url: &str,
483    action: &Value,
484    command: Option<ClientCommand>,
485    events: &mpsc::Sender<Event>,
486) -> Result<()> {
487    let (op_type, op_name) = match engine::parse_operation(
488        action["query"].as_str().unwrap_or_default(),
489        action["operation_name"].as_str(),
490    ) {
491        Ok(v) => v,
492        Err(e) => {
493            reply(command, rejected(e));
494            return Ok(());
495        }
496    };
497    let use_get = action["use_get"].as_bool().unwrap_or(false);
498    if use_get && op_type == OperationType::Mutation {
499        reply(command, rejected("mutations cannot be sent with GET"));
500        return Ok(());
501    }
502    let query = action["query"].as_str().unwrap_or_default();
503    let variables = action.get("variables").cloned().unwrap_or(json!({}));
504    let outcome = exchange(
505        fetch,
506        url,
507        query,
508        action["operation_name"].as_str(),
509        &variables,
510        use_get,
511    )
512    .await;
513    if command.is_some() {
514        ctx.state
515            .record_access_log(
516                AccessLogOwner::Client(ctx.client_id.as_u32()),
517                "GraphQL",
518                None,
519                "injected_action",
520                json!({"operation_type": engine::operation_type_name(op_type)}),
521                vec![json!({"ok": outcome.is_ok()})],
522            )
523            .await;
524    }
525    match outcome {
526        Ok(mut data) => {
527            data["operation_type"] = json!(engine::operation_type_name(op_type));
528            data["operation_name"] = json!(op_name);
529            reply(
530                command,
531                Ok(ClientSendOutcome::Sent {
532                    bytes_sent: query.len(),
533                }),
534            );
535            events
536                .send(Event::new(&actions::RESPONSE_EVENT, data))
537                .await
538                .context("GraphQL event consumer stopped")?;
539        }
540        Err(e) => {
541            Log::new(Some(&ctx.status_tx)).warn(format!("GraphQL request failed: {e}"));
542            reply(command, Err(e));
543        }
544    }
545    Ok(())
546}
547
548/// One GraphQL-over-HTTP request. Transport failures are `Err`; any HTTP answer is `Ok` with
549/// either the checked GraphQL body or an `error` saying why it was not one.
550async fn exchange(
551    fetch: &FetchClient,
552    url: &str,
553    query: &str,
554    operation_name: Option<&str>,
555    variables: &Value,
556    use_get: bool,
557) -> Result<Value> {
558    let request = if use_get {
559        let mut params = vec![("query", query.to_owned())];
560        if let Some(n) = operation_name {
561            params.push(("operationName", n.to_owned()));
562        }
563        if variables.as_object().is_some_and(|v| !v.is_empty()) {
564            params.push(("variables", variables.to_string()));
565        }
566        fetch.get(url).query(&params)
567    } else {
568        let mut body = json!({"query": query});
569        if let Some(n) = operation_name {
570            body["operationName"] = json!(n);
571        }
572        if variables.as_object().is_some_and(|v| !v.is_empty()) {
573            body["variables"] = variables.clone();
574        }
575        fetch
576            .post(url)
577            .header("Content-Type", "application/json")
578            .body(body.to_string())
579    };
580    let response = request.header("Accept", ACCEPT).send().await?;
581    let status = response.status().as_u16();
582    let media = response
583        .headers()
584        .get("content-type")
585        .and_then(|v| v.to_str().ok())
586        .unwrap_or("")
587        .split(';')
588        .next()
589        .unwrap_or("")
590        .trim()
591        .to_ascii_lowercase();
592    let bytes = response.bytes().await?;
593    let mut out = json!({"status": status});
594    if media != GRAPHQL_RESPONSE && media != "application/json" {
595        out["error"] = json!(format!(
596            "HTTP {status} answered with {} rather than a GraphQL response",
597            if media.is_empty() {
598                "no content type"
599            } else {
600                &media
601            }
602        ));
603        return Ok(out);
604    }
605    out["media_type"] = json!(media);
606    let body = match serde_json::from_slice::<Value>(&bytes) {
607        Ok(b) if engine::budget_ok(&b) => b,
608        _ => {
609            out["error"] = json!(format!(
610                "HTTP {status}: body is not a bounded JSON document"
611            ));
612            return Ok(out);
613        }
614    };
615    if let Err(e) = engine::check_response(&body) {
616        out["error"] = json!(format!("HTTP {status}: not a GraphQL response: {e}"));
617        return Ok(out);
618    }
619    // graphql-response+json: a body without data must not come with 2xx.
620    if media == GRAPHQL_RESPONSE && body.get("data").is_none() && (200..300).contains(&status) {
621        out["error"] = json!("a response without data was sent with a 2xx status");
622        return Ok(out);
623    }
624    for k in ["data", "errors", "extensions"] {
625        if let Some(v) = body.get(k) {
626            out[k] = v.clone();
627        }
628    }
629    Ok(out)
630}