netget/client/yarn/
mod.rs

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

netget/client/yarn/
mod.rs
