Skip to main content

netget/client/spark/
mod.rs

1//! Spark client: NetGet reads Apache Spark's monitoring REST API (`/api/v1`, served by a live
2//! driver's UI and by the History Server) and the model decides what to read. See `AGENTS.md`
3//! beside this file.
4//!
5//! Rust owns the URL, the HTTP exchange, every bound and the decoding: an action names a
6//! resource by structured ids (`actions::request` validates them), the path is assembled here
7//! one percent-encoded segment at a time on the base URL fixed at connect, and the response is
8//! measured, parsed and projected (`decode`) before the model sees it.
9//!
10//! Two tasks per client, both registered so `stop_client` ends them:
11//!
12//! * the **session** runs one request at a time — the model's, or one injected from the
13//!   dashboard or MCP — and turns the answer into an event.
14//! * the **turns** task asks the model about each event and hands its actions back. Each event
15//!   carries the number of model turns that led to it; past [`MAX_FOLLOWUP_DEPTH`] the model is
16//!   not asked again.
17//!
18//! **The target is mandatory.** The client talks to `endpoint_url` exactly as given, or to its
19//! `remote_addr`; without one it refuses to start, and redirects are not followed. `connect`
20//! reads `GET /api/v1/version` first, so a target that is not a Spark REST API fails the
21//! client before the model is asked anything.
22//!
23//! Fail closed: an unreachable or non-Spark target fails `connect`; an LLM failure sends
24//! nothing; an error status, or an answer that cannot be read, becomes a `spark_error` naming
25//! why, never an invented result.
26pub mod actions;
27pub mod decode;
28
29use std::net::SocketAddr;
30use std::time::Duration;
31
32use anyhow::{anyhow, ensure, Context, Result};
33use serde_json::{json, Value};
34use tokio::sync::mpsc;
35
36use crate::client::http_fetch::FetchClient;
37use crate::client::llm_budget::call_llm_for_client;
38use crate::llm::actions::client_trait::{Client, ClientActionResult};
39use crate::logging::emit::Log;
40use crate::protocol::{ConnectContext, Event, EventType};
41use crate::state::client_handles::{ClientCommand, ClientSendOutcome};
42use crate::state::{AccessLogOwner, ClientStatus};
43pub use actions::{ApiCall, SparkClientProtocol};
44use decode::Resource;
45
46/// How many model turns may follow from one another before the chain stops.
47pub const MAX_FOLLOWUP_DEPTH: u32 = 8;
48/// How long one request may take.
49pub const REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
50const MAX_REQUEST_TIMEOUT_SECS: u64 = 600;
51/// The largest response body read. A History Server's stage list for a large application is
52/// a few MiB; anything past this is refused and the model told to narrow the request.
53pub const MAX_RESPONSE_BYTES: usize = 8 * 1024 * 1024;
54/// Entries delivered when the action names no `max_items`.
55pub const DEFAULT_MAX_ITEMS: usize = 50;
56/// The most entries one event carries.
57pub const MAX_ITEMS_LIMIT: usize = 500;
58/// The longest application id accepted.
59pub const MAX_ID_CHARS: usize = 128;
60/// The most quantiles one task summary asks for.
61pub const MAX_QUANTILES: usize = 10;
62/// The most of a server's own error text kept.
63pub const MAX_ERROR_BYTES: usize = 1024;
64/// Events waiting for the model.
65const TURNS: usize = 64;
66
67fn user_agent() -> String {
68    format!("NetGet/{}", env!("CARGO_PKG_VERSION"))
69}
70
71/// The base URL this client talks to: `endpoint_url` as given, else `remote_addr` (a bare
72/// `host:port` meaning `http://host:port`). There is no third option. A path is allowed, for a
73/// History Server behind a proxy; a query, a fragment or credentials are not.
74pub fn resolve_endpoint(remote_addr: &str, endpoint_url: Option<&str>) -> Result<url::Url> {
75    let given = match endpoint_url.map(str::trim).filter(|u| !u.is_empty()) {
76        Some(url) => url.to_string(),
77        None => {
78            let raw = remote_addr.trim();
79            ensure!(
80                !raw.is_empty(),
81                "the Spark client needs a target: a remote_addr (host:port) or an endpoint_url. \
82                 Refusing to start without one"
83            );
84            if raw.contains("://") {
85                raw.to_string()
86            } else {
87                format!("http://{raw}")
88            }
89        }
90    };
91    let parsed = url::Url::parse(&given).with_context(|| format!("{given:?} is not a URL"))?;
92    ensure!(
93        matches!(parsed.scheme(), "http" | "https") && parsed.host_str().is_some(),
94        "{given:?} is not an http(s) URL with a host"
95    );
96    ensure!(
97        parsed.query().is_none() && parsed.fragment().is_none(),
98        "{given:?}: the endpoint is a base URL, without a query or fragment"
99    );
100    ensure!(
101        parsed.username().is_empty() && parsed.password().is_none(),
102        "{given:?}: credentials in the URL are not accepted"
103    );
104    ensure!(
105        !parsed.path().trim_end_matches('/').ends_with("/api/v1"),
106        "{given:?}: give the server's base URL; the client adds /api/v1 itself"
107    );
108    Ok(parsed)
109}
110
111/// `base` + `/api/v1/` + each segment, percent-encoded as one segment, + the query. The host,
112/// port and scheme are the base's; nothing in `segments` or `query` can change them.
113pub fn api_url(base: &url::Url, segments: &[String], query: &[(String, String)]) -> Result<String> {
114    let mut url = base.clone();
115    {
116        let mut path = url
117            .path_segments_mut()
118            .map_err(|_| anyhow!("{base} cannot carry a path"))?;
119        path.pop_if_empty().push("api").push("v1");
120        for segment in segments {
121            ensure!(
122                !segment.is_empty() && segment != "." && segment != "..",
123                "an empty or dot path segment is never sent"
124            );
125            path.push(segment);
126        }
127    }
128    if !query.is_empty() {
129        url.query_pairs_mut().extend_pairs(query);
130    }
131    ensure!(
132        url.origin() == base.origin(),
133        "{url} is not on the client's target {base}"
134    );
135    Ok(url.to_string())
136}
137
138#[cfg(not(target_arch = "wasm32"))]
139fn http_client(base: &str, timeout: Duration) -> Result<FetchClient> {
140    let builder = reqwest::Client::builder()
141        .timeout(timeout)
142        .redirect(reqwest::redirect::Policy::none())
143        .user_agent(user_agent());
144    let client = crate::llm::ollama_client::configured_for_endpoint(builder, base)
145        .build()
146        .context("building the HTTP client")?;
147    Ok(FetchClient::from_reqwest(client).with_max_body(MAX_RESPONSE_BYTES))
148}
149
150#[cfg(target_arch = "wasm32")]
151fn http_client(_base: &str, timeout: Duration) -> Result<FetchClient> {
152    Ok(FetchClient::transport(timeout)
153        .with_max_body(MAX_RESPONSE_BYTES)
154        .with_user_agent(&user_agent()))
155}
156
157/// What one request produced: the event to raise.
158pub type Done = (&'static EventType, Value);
159
160/// The event a resource is answered with.
161fn event_for(resource: Resource) -> &'static EventType {
162    match resource {
163        Resource::Applications => &actions::APPLICATIONS_EVENT,
164        Resource::Jobs => &actions::JOBS_EVENT,
165        Resource::Stages => &actions::STAGES_EVENT,
166        Resource::TaskSummary => &actions::TASK_SUMMARY_EVENT,
167        Resource::Executors => &actions::EXECUTORS_EVENT,
168        Resource::Environment => &actions::ENVIRONMENT_EVENT,
169        Resource::Sql => &actions::SQL_EVENT,
170        Resource::Version => &actions::VERSION_EVENT,
171    }
172}
173
174fn error_event(call: &ApiCall, status: Option<u16>, message: String) -> Done {
175    let mut data = json!({"operation": call.operation, "message": message});
176    if let Some(app) = call.context.get("app_id") {
177        data["app_id"] = app.clone();
178    }
179    if let Some(status) = status {
180        data["http_status"] = json!(status);
181    }
182    (&actions::ERROR_EVENT, data)
183}
184
185/// The server's own words for an error: its plain-text body, or the title of an HTML page.
186fn error_text(body: &[u8]) -> String {
187    let text = String::from_utf8_lossy(body);
188    let trimmed = text.trim();
189    let message = if trimmed.starts_with('<') {
190        let lower = trimmed.to_ascii_lowercase();
191        match (lower.find("<title>"), lower.find("</title>")) {
192            (Some(a), Some(b)) if b > a + 7 => trimmed[a + 7..b].to_string(),
193            _ => "an HTML error page".to_string(),
194        }
195    } else {
196        trimmed.to_string()
197    };
198    crate::utils::truncate::truncate_with_suffix(
199        &crate::utils::sanitize::line_field(&message),
200        MAX_ERROR_BYTES,
201        "...",
202    )
203}
204
205/// The REST API this client reads.
206pub struct Api {
207    http: FetchClient,
208    pub base: url::Url,
209    timeout: Duration,
210}
211
212impl Api {
213    pub fn new(base: url::Url, timeout: Duration) -> Result<Api> {
214        Ok(Api {
215            http: http_client(base.as_str(), timeout)?,
216            base,
217            timeout,
218        })
219    }
220
221    /// One GET, bounded in time and size: `(status, body)`, or why there was no answer.
222    async fn get(&self, url: &str) -> Result<(u16, Option<String>, bytes::Bytes)> {
223        let send = async {
224            let resp = self
225                .http
226                .get(url)
227                .header("Accept", "application/json")
228                .send()
229                .await
230                .map_err(|e| anyhow!("no answer from {}: {e:#}", self.base))?;
231            let status = resp.status().as_u16();
232            let location = resp
233                .headers()
234                .get("location")
235                .and_then(|v| v.to_str().ok())
236                .map(str::to_string);
237            let body = resp.bytes().await.map_err(|e| {
238                anyhow!("reading the answer (at most {MAX_RESPONSE_BYTES} bytes are read): {e:#}")
239            })?;
240            ensure!(
241                body.len() <= MAX_RESPONSE_BYTES,
242                "the answer is {} bytes; at most {MAX_RESPONSE_BYTES} are read — narrow the \
243                 request with a status filter or max_items",
244                body.len()
245            );
246            Ok((status, location, body))
247        };
248        tokio::time::timeout(self.timeout, send)
249            .await
250            .map_err(|_| anyhow!("no answer within {}s", self.timeout.as_secs()))?
251    }
252
253    /// Run one validated request. Always an event: the resource's, or `spark_error`.
254    pub async fn run(&self, call: &ApiCall) -> Done {
255        let Some(resource) = Resource::of(&call.operation) else {
256            return error_event(call, None, format!("{} reads nothing", call.operation));
257        };
258        let url = match api_url(&self.base, &call.segments, &call.query) {
259            Ok(u) => u,
260            Err(e) => return error_event(call, None, format!("{e:#}")),
261        };
262        let (status, location, body) = match self.get(&url).await {
263            Ok(r) => r,
264            Err(e) => return error_event(call, None, format!("{e:#}")),
265        };
266        if (300..400).contains(&status) {
267            let to = location
268                .map(|l| crate::utils::sanitize::line_field(&l))
269                .unwrap_or_else(|| "nowhere".into());
270            return error_event(
271                call,
272                Some(status),
273                format!(
274                    "the server redirected to {to}; redirects are not followed — open a client \
275                     on that address if it is the one meant"
276                ),
277            );
278        }
279        if !(200..300).contains(&status) {
280            return error_event(call, Some(status), error_text(&body));
281        }
282        let value = match decode::parse(&body) {
283            Ok(v) => v,
284            Err(e) => return error_event(call, Some(status), format!("{e:#}")),
285        };
286        match decode::decode(
287            resource,
288            &call.operation,
289            &value,
290            call.max_items,
291            &call.context,
292        ) {
293            Ok(payload) => (event_for(resource), payload),
294            Err(e) => error_event(
295                call,
296                Some(status),
297                format!("the answer is not what Spark sends: {e:#}"),
298            ),
299        }
300    }
301
302    /// `GET /api/v1/version`: the server's Spark version, or why this is not a Spark REST API.
303    pub async fn version(&self) -> Result<String> {
304        let call = actions::request(&json!({"type": "get_version"}))?;
305        match self.run(&call).await {
306            (event, payload) if event.id == actions::VERSION_EVENT.id => Ok(payload
307                ["spark_version"]
308                .as_str()
309                .unwrap_or_default()
310                .to_string()),
311            (_, payload) => Err(anyhow!(
312                "{} did not answer GET /api/v1/version as a Spark REST API does: {}",
313                self.base,
314                payload["message"].as_str().unwrap_or("no reason given")
315            )),
316        }
317    }
318}
319
320pub async fn connect(ctx: ConnectContext) -> Result<SocketAddr> {
321    let params = ctx.startup_params.as_ref();
322    let endpoint_url = params
323        .map(|p| p.get_optional_string("endpoint_url"))
324        .transpose()?
325        .flatten();
326    let base = resolve_endpoint(&ctx.remote_addr, endpoint_url.as_deref())?;
327    crate::client::http_fetch::check_url(base.as_str())?;
328    let timeout_secs = params
329        .map(|p| p.get_optional_u64("request_timeout_secs"))
330        .transpose()?
331        .flatten()
332        .unwrap_or(REQUEST_TIMEOUT.as_secs());
333    ensure!(
334        (1..=MAX_REQUEST_TIMEOUT_SECS).contains(&timeout_secs),
335        "request_timeout_secs is 1..={MAX_REQUEST_TIMEOUT_SECS}"
336    );
337    let log = Log::new(Some(&ctx.status_tx));
338    let api = Api::new(base.clone(), Duration::from_secs(timeout_secs))?;
339    let version = match api.version().await {
340        Ok(v) => v,
341        Err(e) => {
342            log.warn(format!(
343                "Spark client {}: {e:#} decision=fail_closed_connect",
344                ctx.client_id
345            ));
346            return Err(e);
347        }
348    };
349    ctx.state
350        .update_client_status(ctx.client_id, ClientStatus::Connected)
351        .await;
352    log.info(format!(
353        "Spark client {} reached Spark {version} at {base}",
354        ctx.client_id
355    ));
356    let _ = ctx.status_tx.send("__UPDATE_UI__".into());
357
358    // Registered before the connected turn can run, so the operator can reach the client
359    // while that turn is parked on a manual rule.
360    let external =
361        crate::client::command_support::register_command_channel(&ctx.state, ctx.client_id).await;
362    let (event_tx, event_rx) = mpsc::channel::<(Event, u32)>(TURNS);
363    let (internal_tx, internal) = mpsc::channel::<(Value, u32)>(TURNS);
364    let connected = json!({"endpoint": base.as_str(), "spark_version": version});
365    let _ = event_tx.try_send((Event::new(&actions::CONNECTED_EVENT, connected), 0));
366    let dispatcher = tokio::spawn(run_turns(ctx.clone(), event_rx, internal_tx));
367    let dispatcher_abort = dispatcher.abort_handle();
368    ctx.state
369        .register_client_task(ctx.client_id, dispatcher)
370        .await;
371    let session_ctx = ctx.clone();
372    let task = tokio::spawn(async move {
373        run_session(&api, &session_ctx, external, internal, event_tx).await;
374        dispatcher_abort.abort();
375        session_ctx
376            .state
377            .update_client_status(session_ctx.client_id, ClientStatus::Disconnected)
378            .await;
379        session_ctx
380            .state
381            .remove_client_handle(session_ctx.client_id)
382            .await;
383        let _ = session_ctx.status_tx.send("__UPDATE_UI__".into());
384    });
385    ctx.state.register_client_task(ctx.client_id, task).await;
386    // HTTP holds no one socket for the session; there is no local address to report.
387    Ok(SocketAddr::from(([0, 0, 0, 0], 0)))
388}
389
390fn reply(command: Option<ClientCommand>, outcome: ClientSendOutcome) {
391    if let Some(c) = command {
392        crate::client::command_support::reply(c, Ok(outcome));
393    }
394}
395
396/// Run requests one at a time until the client disconnects or is stopped.
397async fn run_session(
398    api: &Api,
399    ctx: &ConnectContext,
400    mut external: mpsc::Receiver<ClientCommand>,
401    mut internal: mpsc::Receiver<(Value, u32)>,
402    events: mpsc::Sender<(Event, u32)>,
403) {
404    let protocol = SparkClientProtocol;
405    loop {
406        let (action, depth, command) = tokio::select! {
407            c = external.recv() => match c {
408                Some(c) => {
409                    ctx.state.record_access_log(AccessLogOwner::Client(ctx.client_id.as_u32()), "Spark",
410                        None, "injected_action", c.action.clone(), vec![]).await;
411                    (c.action.clone(), 0, Some(c))
412                }
413                None => return,
414            },
415            a = internal.recv() => match a {
416                Some((a, d)) => (a, d, None),
417                None => return,
418            },
419        };
420        let call = match protocol.execute_action(action) {
421            Ok(ClientActionResult::Custom { data, .. }) => {
422                match serde_json::from_value::<ApiCall>(data) {
423                    Ok(call) => call,
424                    Err(e) => {
425                        reply(
426                            command,
427                            ClientSendOutcome::Rejected {
428                                error: e.to_string(),
429                            },
430                        );
431                        continue;
432                    }
433                }
434            }
435            Ok(ClientActionResult::Disconnect) => {
436                Log::new(Some(&ctx.status_tx)).info(format!(
437                    "Spark client {} disconnecting from {}",
438                    ctx.client_id, api.base
439                ));
440                reply(command, ClientSendOutcome::Disconnected);
441                return;
442            }
443            Ok(_) => {
444                reply(
445                    command,
446                    ClientSendOutcome::Executed {
447                        detail: "nothing to send".into(),
448                    },
449                );
450                continue;
451            }
452            Err(e) => {
453                Log::new(Some(&ctx.status_tx)).warn(format!(
454                    "Spark client {}: action refused: {e:#}",
455                    ctx.client_id
456                ));
457                reply(
458                    command,
459                    ClientSendOutcome::Rejected {
460                        error: e.to_string(),
461                    },
462                );
463                continue;
464            }
465        };
466        let (event_type, payload) = api.run(&call).await;
467        let request_bytes = api_url(&api.base, &call.segments, &call.query)
468            .map(|u| u.len())
469            .unwrap_or(0);
470        reply(
471            command,
472            ClientSendOutcome::Sent {
473                bytes_sent: request_bytes,
474            },
475        );
476        tracing::info!(
477            "Spark client {} {} -> {} (depth {depth})",
478            ctx.client_id,
479            call.operation,
480            event_type.id
481        );
482        if events
483            .try_send((Event::new(event_type, payload), depth))
484            .is_err()
485        {
486            Log::new(Some(&ctx.status_tx)).warn(format!(
487                "Spark client {}: {} events are waiting for the model; this one is dropped \
488                 decision=turn_queue_full",
489                ctx.client_id, TURNS
490            ));
491        }
492        let _ = ctx.status_tx.send("__UPDATE_UI__".into());
493    }
494}
495
496/// Ask the model about each event, in order, and hand its actions to the session.
497async fn run_turns(
498    ctx: ConnectContext,
499    mut events: mpsc::Receiver<(Event, u32)>,
500    internal: mpsc::Sender<(Value, u32)>,
501) {
502    let protocol = SparkClientProtocol;
503    while let Some((event, depth)) = events.recv().await {
504        if depth >= MAX_FOLLOWUP_DEPTH {
505            Log::new(Some(&ctx.status_tx)).warn(format!(
506                "Spark client {} not asking the model about {}: {depth} turns deep \
507                 decision=followup_depth",
508                ctx.client_id,
509                event.id()
510            ));
511            continue;
512        }
513        let instruction = ctx
514            .state
515            .get_instruction_for_client(ctx.client_id)
516            .await
517            .unwrap_or_default();
518        let memory = ctx
519            .state
520            .get_memory_for_client(ctx.client_id)
521            .await
522            .unwrap_or_default();
523        match call_llm_for_client(
524            &ctx.llm_client,
525            &ctx.state,
526            ctx.client_id.to_string(),
527            &instruction,
528            &memory,
529            Some(&event),
530            &protocol,
531            &ctx.status_tx,
532        )
533        .await
534        {
535            Ok(result) => {
536                if let Some(memory) = result.memory_updates {
537                    ctx.state.set_memory_for_client(ctx.client_id, memory).await;
538                }
539                tracing::info!(
540                    "Spark client {} {} decision={} actions={}",
541                    ctx.client_id,
542                    event.id(),
543                    if result.actions.is_empty() {
544                        "model_silent"
545                    } else {
546                        "model_actions"
547                    },
548                    result.actions.len()
549                );
550                for action in result.actions {
551                    if internal.send((action, depth + 1)).await.is_err() {
552                        return;
553                    }
554                }
555            }
556            // Fail closed: nothing is read on the model's behalf when it cannot be asked.
557            Err(e) => Log::new(Some(&ctx.status_tx)).warn(format!(
558                "Spark client {} {}: {e:#} decision=fail_closed_llm_error",
559                ctx.client_id,
560                event.id()
561            )),
562        }
563    }
564}