1pub 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
164fn 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
172struct Socket {
174 sink: futures::stream::SplitSink<WsStream, Message>,
175 stream: futures::stream::SplitStream<WsStream>,
176 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
186async 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
364async 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 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
548async 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(¶ms)
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 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}