1use std::sync::LazyLock;
10
11use anyhow::{bail, ensure, Context, Result};
12use serde::{Deserialize, Serialize};
13use serde_json::{json, Map, Value};
14
15use crate::llm::actions::client_trait::{Client, ClientActionResult};
16use crate::llm::actions::protocol_trait::Protocol;
17use crate::llm::actions::{ActionDefinition, Parameter, ParameterDefinition, StartupExamples};
18use crate::protocol::log_template::LogTemplate;
19use crate::protocol::{ConnectContext, EventType};
20use crate::state::app_state::AppState;
21
22pub struct SparkClientProtocol;
23
24impl SparkClientProtocol {
25 pub fn new() -> Self {
26 Self
27 }
28}
29
30impl Default for SparkClientProtocol {
31 fn default() -> Self {
32 Self::new()
33 }
34}
35
36fn p(name: &str, type_hint: &str, description: &str, required: bool) -> Parameter {
37 Parameter {
38 name: name.to_string(),
39 type_hint: type_hint.to_string(),
40 description: description.to_string(),
41 required,
42 }
43}
44
45fn action(
46 name: &str,
47 description: &str,
48 parameters: Vec<Parameter>,
49 example: Value,
50) -> ActionDefinition {
51 ActionDefinition {
52 name: name.to_string(),
53 description: description.to_string(),
54 parameters,
55 example,
56 log_template: Some(LogTemplate::new().with_info(format!("-> Spark {name} {{app_id}}"))),
57 }
58}
59
60fn app_id_param() -> Parameter {
61 p(
62 "app_id",
63 "string",
64 "The application id exactly as Spark reports it (local-1792044029340, \
65 app-20261015060028-0000, application_1700000000000_0001)",
66 true,
67 )
68}
69
70fn app_attempt_param() -> Parameter {
71 p(
72 "app_attempt_id",
73 "string",
74 "The application attempt, only for an application the History Server lists with an \
75 attemptId (YARN cluster mode); omit it otherwise",
76 false,
77 )
78}
79
80fn max_items_param() -> Parameter {
81 p(
82 "max_items",
83 "number",
84 &format!(
85 "How many entries to deliver at most (default {}, at most {}); truncated says \
86 whether more existed",
87 super::DEFAULT_MAX_ITEMS,
88 super::MAX_ITEMS_LIMIT
89 ),
90 false,
91 )
92}
93
94pub fn actions() -> Vec<ActionDefinition> {
95 vec![
96 action(
97 "list_applications",
98 "List the applications the server knows (GET /api/v1/applications). A live driver UI \
99 knows its own application; a History Server every one in its event log directory. \
100 Answered by spark_applications.",
101 vec![
102 p("status", "string", "completed or running", false),
103 p("min_date", "string", "Started on or after: 2026-10-15 or 2026-10-15T06:00:00.000GMT", false),
104 p("max_date", "string", "Started on or before, same format", false),
105 p("min_end_date", "string", "Ended on or after, same format", false),
106 p("max_end_date", "string", "Ended on or before, same format", false),
107 p("limit", "number", "How many applications the server returns at most (1..=1000)", false),
108 max_items_param(),
109 ],
110 json!({"type": "list_applications", "status": "completed", "limit": 20}),
111 ),
112 action(
113 "get_application",
114 "One application with its attempts (GET /api/v1/applications/{app_id}). Answered by spark_applications.",
115 vec![app_id_param()],
116 json!({"type": "get_application", "app_id": "local-1792044029340"}),
117 ),
118 action(
119 "list_jobs",
120 "An application's jobs, newest first (GET .../{app_id}/jobs). Answered by spark_jobs.",
121 vec![
122 app_id_param(),
123 app_attempt_param(),
124 p("status", "string|array", "Only jobs in these states (one, or a list): running, succeeded, failed, unknown", false),
125 max_items_param(),
126 ],
127 json!({"type": "list_jobs", "app_id": "local-1792044029340", "status": "failed"}),
128 ),
129 action(
130 "get_job",
131 "One job (GET .../{app_id}/jobs/{job_id}). Answered by spark_jobs.",
132 vec![app_id_param(), app_attempt_param(), p("job_id", "number", "The jobId, as list_jobs reports it", true)],
133 json!({"type": "get_job", "app_id": "local-1792044029340", "job_id": 1}),
134 ),
135 action(
136 "list_stages",
137 "An application's stages, every attempt of each (GET .../{app_id}/stages). Answered by spark_stages.",
138 vec![
139 app_id_param(),
140 app_attempt_param(),
141 p("status", "string|array", "Only stages in these states (one, or a list): active, complete, pending, failed, skipped", false),
142 max_items_param(),
143 ],
144 json!({"type": "list_stages", "app_id": "local-1792044029340", "status": "failed"}),
145 ),
146 action(
147 "get_stage",
148 "One stage: every attempt of it, or the one attempt named (GET .../{app_id}/stages/{stage_id}[/{stage_attempt_id}]). Task lists are never fetched. Answered by spark_stages.",
149 vec![
150 app_id_param(),
151 app_attempt_param(),
152 p("stage_id", "number", "The stageId, as list_jobs (stageIds) or list_stages reports it", true),
153 p("stage_attempt_id", "number", "One attempt of the stage (0 is the first)", false),
154 ],
155 json!({"type": "get_stage", "app_id": "local-1792044029340", "stage_id": 2, "stage_attempt_id": 0}),
156 ),
157 action(
158 "get_task_summary",
159 "The distribution of task metrics in one stage attempt at the quantiles asked for (GET .../stages/{stage_id}/{stage_attempt_id}/taskSummary). Answered by spark_task_summary.",
160 vec![
161 app_id_param(),
162 app_attempt_param(),
163 p("stage_id", "number", "The stageId, as list_jobs (stageIds) or list_stages reports it", true),
164 p("stage_attempt_id", "number", "The stage attempt (0 is the first)", true),
165 p("quantiles", "array", "Up to 10 numbers in 0..=1 (default 0.05, 0.25, 0.5, 0.75, 0.95)", false),
166 ],
167 json!({"type": "get_task_summary", "app_id": "local-1792044029340", "stage_id": 2, "stage_attempt_id": 0, "quantiles": [0.5, 0.95]}),
168 ),
169 action(
170 "list_executors",
171 "An application's executors, the driver included (GET .../{app_id}/executors, or .../allexecutors with include_dead). Answered by spark_executors.",
172 vec![
173 app_id_param(),
174 app_attempt_param(),
175 p("include_dead", "boolean", "Also executors that have been removed", false),
176 max_items_param(),
177 ],
178 json!({"type": "list_executors", "app_id": "local-1792044029340"}),
179 ),
180 action(
181 "get_environment",
182 "An application's runtime (Java, Scala) and its Spark properties, as Spark reports them with secrets already redacted (GET .../{app_id}/environment). Answered by spark_environment.",
183 vec![app_id_param(), app_attempt_param(), max_items_param()],
184 json!({"type": "get_environment", "app_id": "local-1792044029340"}),
185 ),
186 action(
187 "list_sql_executions",
188 "An application's SQL/DataFrame executions with the jobs each ran (GET .../{app_id}/sql). Answered by spark_sql_executions.",
189 vec![
190 app_id_param(),
191 app_attempt_param(),
192 p("offset", "number", "Skip this many executions", false),
193 max_items_param(),
194 ],
195 json!({"type": "list_sql_executions", "app_id": "local-1792044029340"}),
196 ),
197 action(
198 "get_sql_execution",
199 "One SQL execution with its physical plan text and plan nodes (GET .../{app_id}/sql/{execution_id}). Answered by spark_sql_executions.",
200 vec![
201 app_id_param(),
202 app_attempt_param(),
203 p("execution_id", "number", "The SQL execution id", true),
204 ],
205 json!({"type": "get_sql_execution", "app_id": "local-1792044029340", "execution_id": 0}),
206 ),
207 action(
208 "get_version",
209 "The Spark version of the server (GET /api/v1/version). Answered by spark_version.",
210 vec![],
211 json!({"type": "get_version"}),
212 ),
213 action(
214 "disconnect",
215 "Stop this client. The REST API keeps no session, so nothing is sent.",
216 vec![],
217 json!({"type": "disconnect"}),
218 ),
219 ]
220}
221
222fn ev(id: &str, description: &str, params: Vec<Parameter>) -> EventType {
223 EventType::new(
224 id,
225 description,
226 json!({"type": "list_jobs", "app_id": "local-1792044029340"}),
227 )
228 .with_parameters(params)
229 .with_actions(actions())
230}
231
232fn common(mut params: Vec<Parameter>) -> Vec<Parameter> {
233 params.insert(
234 0,
235 p(
236 "operation",
237 "string",
238 "The action this answers, e.g. list_jobs or get_job",
239 true,
240 ),
241 );
242 params.extend([
243 p("count", "number", "How many entries are delivered", true),
244 p(
245 "total",
246 "number",
247 "How many entries the server returned",
248 true,
249 ),
250 p(
251 "truncated",
252 "boolean",
253 "True when the server returned more than max_items",
254 true,
255 ),
256 ]);
257 params
258}
259
260fn app_param() -> Parameter {
261 p(
262 "app_id",
263 "string",
264 "The application the entries belong to",
265 true,
266 )
267}
268
269pub static CONNECTED_EVENT: LazyLock<EventType> = LazyLock::new(|| {
270 ev(
271 "spark_connected",
272 "The client reached a Spark monitoring REST API (a live driver UI or a History Server) and read its version; nothing else has been asked.",
273 vec![
274 p("endpoint", "string", "The base URL every request goes to", true),
275 p("spark_version", "string", "The version the server reported (GET /api/v1/version)", true),
276 ],
277 )
278});
279
280pub static APPLICATIONS_EVENT: LazyLock<EventType> = LazyLock::new(|| {
281 ev(
282 "spark_applications",
283 "Applications, from list_applications or get_application.",
284 common(vec![
285 p("applications", "array", "One {id, name, attempts: [{attemptId, startTime, endTime, lastUpdated, duration, sparkUser, completed, appSparkVersion}], coresGranted, maxCores, memoryPerExecutorMB} per application", true),
286 p("filters", "object", "The filters list_applications sent", false),
287 ]),
288 )
289});
290
291pub static JOBS_EVENT: LazyLock<EventType> = LazyLock::new(|| {
292 ev(
293 "spark_jobs",
294 "Jobs, from list_jobs or get_job.",
295 common(vec![
296 app_param(),
297 p("jobs", "array", "One {jobId, name, description, status (RUNNING/SUCCEEDED/FAILED/UNKNOWN), stageIds, jobGroup, submissionTime, completionTime, numTasks, numCompletedTasks, numFailedTasks, numSkippedTasks, numKilledTasks, numActiveStages, numCompletedStages, numSkippedStages, numFailedStages} per job", true),
298 ]),
299 )
300});
301
302pub static STAGES_EVENT: LazyLock<EventType> = LazyLock::new(|| {
303 ev(
304 "spark_stages",
305 "Stage attempts, from list_stages or get_stage.",
306 common(vec![
307 app_param(),
308 p("stages", "array", "One {stageId, attemptId, status (ACTIVE/COMPLETE/PENDING/FAILED/SKIPPED), name, description, numTasks, numCompleteTasks, numFailedTasks, numKilledTasks, submissionTime, completionTime, failureReason, executorRunTime, inputBytes, outputBytes, shuffleReadBytes, shuffleWriteBytes, memoryBytesSpilled, diskBytesSpilled, rddIds, ...} per stage attempt", true),
309 ]),
310 )
311});
312
313pub static TASK_SUMMARY_EVENT: LazyLock<EventType> = LazyLock::new(|| {
314 ev(
315 "spark_task_summary",
316 "The task-metric distribution of one stage attempt, from get_task_summary.",
317 common(vec![
318 app_param(),
319 p("stage_id", "number", "The stage the summary describes", true),
320 p("stage_attempt_id", "number", "The stage attempt", true),
321 p("quantiles", "array", "The quantiles every metric list is given at, in order", true),
322 p("metrics", "object", "Metric name -> one number per quantile (duration, executorRunTime, jvmGcTime, resultSize, schedulerDelay, peakExecutionMemory, memoryBytesSpilled, diskBytesSpilled, ...); inputMetrics/outputMetrics/shuffleReadMetrics/shuffleWriteMetrics are objects of the same", true),
323 ]),
324 )
325});
326
327pub static EXECUTORS_EVENT: LazyLock<EventType> = LazyLock::new(|| {
328 ev(
329 "spark_executors",
330 "Executors, from list_executors.",
331 common(vec![
332 app_param(),
333 p("executors", "array", "One {id (\"driver\" or a number as text), hostPort, isActive, totalCores, maxTasks, activeTasks, failedTasks, completedTasks, totalTasks, totalDuration, totalGCTime, totalInputBytes, totalShuffleRead, totalShuffleWrite, memoryUsed, maxMemory, diskUsed, rddBlocks, addTime, removeTime, removeReason} per executor", true),
334 ]),
335 )
336});
337
338pub static ENVIRONMENT_EVENT: LazyLock<EventType> = LazyLock::new(|| {
339 ev(
340 "spark_environment",
341 "An application's environment, from get_environment.",
342 common(vec![
343 app_param(),
344 p("runtime", "object", "{javaVersion, javaHome, scalaVersion}", true),
345 p("spark_properties", "object", "Spark property name -> value, as Spark reports them (it redacts secrets itself); count/total/truncated describe this map", true),
346 p("system_property_count", "number", "How many JVM system properties Spark reported (not delivered)", true),
347 p("hadoop_property_count", "number", "How many Hadoop properties Spark reported (not delivered)", true),
348 p("classpath_entry_count", "number", "How many classpath entries Spark reported (not delivered)", true),
349 ]),
350 )
351});
352
353pub static SQL_EVENT: LazyLock<EventType> = LazyLock::new(|| {
354 ev(
355 "spark_sql_executions",
356 "SQL/DataFrame executions, from list_sql_executions or get_sql_execution.",
357 common(vec![
358 app_param(),
359 p("executions", "array", "One {id, status (RUNNING/COMPLETED/FAILED), description, submissionTime, duration, runningJobIds, successJobIds, failedJobIds, queryId, errorMessage} per execution; get_sql_execution adds planDescription and nodes [{nodeId, nodeName}]", true),
360 ]),
361 )
362});
363
364pub static VERSION_EVENT: LazyLock<EventType> = LazyLock::new(|| {
365 ev(
366 "spark_version",
367 "The server's Spark version, from get_version.",
368 vec![
369 p(
370 "operation",
371 "string",
372 "The action this answers: get_version",
373 true,
374 ),
375 p(
376 "spark_version",
377 "string",
378 "The version the server reported",
379 true,
380 ),
381 ],
382 )
383});
384
385pub static ERROR_EVENT: LazyLock<EventType> = LazyLock::new(|| {
386 ev(
387 "spark_error",
388 "A request failed: the server answered with an error status (404 unknown app or job, 400 bad filter), or the answer could not be read (not JSON, too large, too deep, not the shape Spark sends). The client is still connected.",
389 vec![
390 p("operation", "string", "The action that failed", true),
391 p("app_id", "string", "The application it asked about, if any", false),
392 p("http_status", "number", "The HTTP status of the response, when there was one", false),
393 p("message", "string", "What went wrong, in the server's words (its plain-text error body) or this client's", true),
394 ],
395 )
396});
397
398#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
402pub struct ApiCall {
403 pub operation: String,
405 pub segments: Vec<String>,
407 pub query: Vec<(String, String)>,
409 pub max_items: usize,
411 pub context: Map<String, Value>,
413}
414
415pub fn app_id(raw: &str) -> Result<String> {
417 let id = raw.trim();
418 ensure!(!id.is_empty(), "app_id is empty");
419 ensure!(
420 id.len() <= super::MAX_ID_CHARS,
421 "app_id is {} characters; at most {} are accepted",
422 id.len(),
423 super::MAX_ID_CHARS
424 );
425 ensure!(
426 id.chars()
427 .all(|c| c.is_ascii_alphanumeric() || matches!(c, '-' | '_' | '.')),
428 "app_id {id:?}: only letters, digits, '-', '_' and '.' are accepted (Spark ids look \
429 like local-1792044029340 or application_1700000000000_0001)"
430 );
431 ensure!(!id.starts_with('.'), "app_id {id:?} may not start with '.'");
432 Ok(id.to_string())
433}
434
435fn attempt_id(raw: &str) -> Result<String> {
436 let id = raw.trim();
437 ensure!(
438 !id.is_empty()
439 && id.len() <= 32
440 && !id.starts_with('.')
441 && id
442 .chars()
443 .all(|c| c.is_ascii_alphanumeric() || matches!(c, '-' | '_' | '.')),
444 "app_attempt_id {id:?}: at most 32 letters, digits, '-', '_' and '.'"
445 );
446 Ok(id.to_string())
447}
448
449fn numeric_id(v: &Value, name: &str) -> Result<u64> {
451 let n = match v {
452 Value::Number(n) => n
453 .as_u64()
454 .with_context(|| format!("{name} is a whole number >= 0, not {n}"))?,
455 Value::String(s)
456 if !s.is_empty() && s.len() <= 19 && s.bytes().all(|b| b.is_ascii_digit()) =>
457 {
458 s.parse::<u64>()?
459 }
460 other => bail!("{name} is a whole number >= 0, not {other}"),
461 };
462 ensure!(
463 n <= i64::MAX as u64,
464 "{name} {n} is larger than any id Spark issues"
465 );
466 Ok(n)
467}
468
469fn get<'a>(v: &'a Value, key: &str) -> Option<&'a Value> {
470 v.get(key).filter(|x| !x.is_null())
471}
472
473fn required_id(v: &Value, key: &str) -> Result<u64> {
474 numeric_id(
475 get(v, key).with_context(|| format!("{} needs {key}", op_name(v)))?,
476 key,
477 )
478}
479
480fn op_name(v: &Value) -> &str {
481 v.get("type")
482 .and_then(Value::as_str)
483 .unwrap_or("the action")
484}
485
486fn date_filter(v: &Value, key: &str) -> Result<String> {
489 let s = v
490 .as_str()
491 .with_context(|| format!("{key} is a date as text"))?
492 .trim();
493 let b = s.as_bytes();
494 let date_ok = b.len() >= 10
495 && b[..10].iter().enumerate().all(|(i, c)| match i {
496 4 | 7 => *c == b'-',
497 _ => c.is_ascii_digit(),
498 });
499 let rest_ok = b.len() == 10
500 || (b.len() >= 23
501 && b.len() <= 32
502 && b[10] == b'T'
503 && b[11..23].iter().enumerate().all(|(i, c)| match i {
504 2 | 5 => *c == b':',
505 8 => *c == b'.',
506 _ => c.is_ascii_digit(),
507 })
508 && b[23..]
509 .iter()
510 .all(|c| c.is_ascii_alphanumeric() || *c == b'+' || *c == b'-'));
511 ensure!(
512 date_ok && rest_ok,
513 "{key} {s:?}: Spark reads 2026-10-15 or 2026-10-15T06:00:00.000GMT"
514 );
515 Ok(s.to_string())
516}
517
518fn statuses(v: &Value, allowed: &[&str]) -> Result<Vec<String>> {
520 let list: Vec<&Value> = match v {
521 Value::Array(items) => items.iter().collect(),
522 other => vec![other],
523 };
524 ensure!(
525 !list.is_empty() && list.len() <= allowed.len(),
526 "status: one to {} of {}",
527 allowed.len(),
528 allowed.join(", ")
529 );
530 let mut out = Vec::new();
531 for item in list {
532 let s = item
533 .as_str()
534 .context("status is text")?
535 .trim()
536 .to_ascii_lowercase();
537 ensure!(
538 allowed.contains(&s.as_str()),
539 "status {s:?} is not one of {}",
540 allowed.join(", ")
541 );
542 if !out.contains(&s) {
543 out.push(s);
544 }
545 }
546 Ok(out)
547}
548
549fn max_items(v: &Value) -> Result<usize> {
550 let n = match get(v, "max_items") {
551 None => super::DEFAULT_MAX_ITEMS as u64,
552 Some(n) => n.as_u64().context("max_items is a whole number")?,
553 };
554 ensure!(
555 (1..=super::MAX_ITEMS_LIMIT as u64).contains(&n),
556 "max_items is 1..={}",
557 super::MAX_ITEMS_LIMIT
558 );
559 Ok(n as usize)
560}
561
562fn app_prefix(v: &Value, ctx: &mut Map<String, Value>) -> Result<Vec<String>> {
564 let id = app_id(
565 get(v, "app_id")
566 .and_then(Value::as_str)
567 .with_context(|| format!("{} needs app_id (text)", op_name(v)))?,
568 )?;
569 ctx.insert("app_id".into(), json!(id));
570 let mut segments = vec!["applications".to_string(), id];
571 if let Some(attempt) = get(v, "app_attempt_id") {
572 let attempt = attempt_id(&match attempt {
573 Value::String(s) => s.clone(),
574 Value::Number(n) => n.to_string(),
575 other => bail!("app_attempt_id is text, not {other}"),
576 })?;
577 ctx.insert("app_attempt_id".into(), json!(attempt));
578 segments.push(attempt);
579 }
580 Ok(segments)
581}
582
583pub fn request(v: &Value) -> Result<ApiCall> {
585 let kind = v
586 .get("type")
587 .and_then(Value::as_str)
588 .context("an action needs a type")?;
589 let mut context = Map::new();
590 let mut query: Vec<(String, String)> = Vec::new();
591 let mut items = super::DEFAULT_MAX_ITEMS;
592 let segments: Vec<String> = match kind {
593 "list_applications" => {
594 let mut filters = Map::new();
595 if let Some(s) = get(v, "status") {
596 let s = statuses(s, &["completed", "running"])?;
597 ensure!(s.len() == 1, "list_applications takes one status");
598 filters.insert("status".into(), json!(s[0]));
599 query.push(("status".into(), s[0].clone()));
600 }
601 for (key, param) in [
602 ("min_date", "minDate"),
603 ("max_date", "maxDate"),
604 ("min_end_date", "minEndDate"),
605 ("max_end_date", "maxEndDate"),
606 ] {
607 if let Some(d) = get(v, key) {
608 let d = date_filter(d, key)?;
609 filters.insert(key.into(), json!(d));
610 query.push((param.into(), d));
611 }
612 }
613 if let Some(limit) = get(v, "limit") {
614 let limit = limit.as_u64().context("limit is a whole number")?;
615 ensure!((1..=1000).contains(&limit), "limit is 1..=1000");
616 filters.insert("limit".into(), json!(limit));
617 query.push(("limit".into(), limit.to_string()));
618 }
619 context.insert("filters".into(), Value::Object(filters));
620 items = max_items(v)?;
621 vec!["applications".into()]
622 }
623 "get_application" => {
624 let id = app_id(
625 get(v, "app_id")
626 .and_then(Value::as_str)
627 .context("get_application needs app_id (text)")?,
628 )?;
629 context.insert("app_id".into(), json!(id));
630 vec!["applications".into(), id]
631 }
632 "list_jobs" => {
633 let mut s = app_prefix(v, &mut context)?;
634 if let Some(st) = get(v, "status") {
635 let st = statuses(st, &["running", "succeeded", "failed", "unknown"])?;
636 context.insert("status".into(), json!(st));
637 query.extend(st.into_iter().map(|x| ("status".to_string(), x)));
638 }
639 items = max_items(v)?;
640 s.push("jobs".into());
641 s
642 }
643 "get_job" => {
644 let mut s = app_prefix(v, &mut context)?;
645 let job = required_id(v, "job_id")?;
646 context.insert("job_id".into(), json!(job));
647 s.extend(["jobs".to_string(), job.to_string()]);
648 s
649 }
650 "list_stages" => {
651 let mut s = app_prefix(v, &mut context)?;
652 if let Some(st) = get(v, "status") {
653 let st = statuses(st, &["active", "complete", "pending", "failed", "skipped"])?;
654 context.insert("status".into(), json!(st));
655 query.extend(st.into_iter().map(|x| ("status".to_string(), x)));
656 }
657 query.push(("details".into(), "false".into()));
659 items = max_items(v)?;
660 s.push("stages".into());
661 s
662 }
663 "get_stage" => {
664 let mut s = app_prefix(v, &mut context)?;
665 let stage = required_id(v, "stage_id")?;
666 context.insert("stage_id".into(), json!(stage));
667 s.extend(["stages".to_string(), stage.to_string()]);
668 if let Some(a) = get(v, "stage_attempt_id") {
669 let a = numeric_id(a, "stage_attempt_id")?;
670 context.insert("stage_attempt_id".into(), json!(a));
671 s.push(a.to_string());
672 }
673 query.push(("details".into(), "false".into()));
674 s
675 }
676 "get_task_summary" => {
677 let mut s = app_prefix(v, &mut context)?;
678 let stage = required_id(v, "stage_id")?;
679 let attempt = required_id(v, "stage_attempt_id")?;
680 context.insert("stage_id".into(), json!(stage));
681 context.insert("stage_attempt_id".into(), json!(attempt));
682 let quantiles: Vec<f64> = match get(v, "quantiles") {
683 None => vec![0.05, 0.25, 0.5, 0.75, 0.95],
684 Some(Value::Array(q)) => q
685 .iter()
686 .map(|x| x.as_f64().context("quantiles are numbers"))
687 .collect::<Result<_>>()?,
688 Some(other) => bail!("quantiles is an array of numbers, not {other}"),
689 };
690 ensure!(
691 !quantiles.is_empty() && quantiles.len() <= super::MAX_QUANTILES,
692 "quantiles: 1 to {} numbers",
693 super::MAX_QUANTILES
694 );
695 ensure!(
696 quantiles.iter().all(|q| (0.0..=1.0).contains(q)),
697 "quantiles are in 0..=1"
698 );
699 query.push((
700 "quantiles".into(),
701 quantiles
702 .iter()
703 .map(|q| q.to_string())
704 .collect::<Vec<_>>()
705 .join(","),
706 ));
707 s.extend([
708 "stages".to_string(),
709 stage.to_string(),
710 attempt.to_string(),
711 "taskSummary".to_string(),
712 ]);
713 s
714 }
715 "list_executors" => {
716 let mut s = app_prefix(v, &mut context)?;
717 let dead = match get(v, "include_dead") {
718 None => false,
719 Some(b) => b.as_bool().context("include_dead is true or false")?,
720 };
721 context.insert("include_dead".into(), json!(dead));
722 items = max_items(v)?;
723 s.push(if dead { "allexecutors" } else { "executors" }.into());
724 s
725 }
726 "get_environment" => {
727 let mut s = app_prefix(v, &mut context)?;
728 items = max_items(v)?;
729 s.push("environment".into());
730 s
731 }
732 "list_sql_executions" => {
733 let mut s = app_prefix(v, &mut context)?;
734 items = max_items(v)?;
735 let offset = match get(v, "offset") {
736 None => 0,
737 Some(o) => o.as_u64().context("offset is a whole number")?,
738 };
739 ensure!(offset <= i32::MAX as u64, "offset is too large");
740 context.insert("offset".into(), json!(offset));
741 query.extend([
743 ("details".to_string(), "false".to_string()),
744 ("planDescription".to_string(), "false".to_string()),
745 ("offset".to_string(), offset.to_string()),
746 ("length".to_string(), items.to_string()),
747 ]);
748 s.push("sql".into());
749 s
750 }
751 "get_sql_execution" => {
752 let mut s = app_prefix(v, &mut context)?;
753 let id = required_id(v, "execution_id")?;
754 context.insert("execution_id".into(), json!(id));
755 query.extend([
756 ("details".to_string(), "true".to_string()),
757 ("planDescription".to_string(), "true".to_string()),
758 ]);
759 s.extend(["sql".to_string(), id.to_string()]);
760 s
761 }
762 "get_version" => vec!["version".into()],
763 "disconnect" => Vec::new(),
764 other => bail!("Unknown Spark client action: {other}"),
765 };
766 Ok(ApiCall {
767 operation: kind.to_string(),
768 segments,
769 query,
770 max_items: items,
771 context,
772 })
773}
774
775fn param(
776 name: &str,
777 type_hint: &str,
778 description: &str,
779 required: bool,
780 example: Value,
781 default: Option<Value>,
782) -> ParameterDefinition {
783 ParameterDefinition {
784 name: name.into(),
785 type_hint: type_hint.into(),
786 description: description.into(),
787 required,
788 example,
789 default,
790 }
791}
792
793impl Protocol for SparkClientProtocol {
794 fn protocol_name(&self) -> &'static str {
795 "Spark"
796 }
797 fn stack_name(&self) -> &'static str {
798 "ETH>IP>TCP>HTTP>SPARK"
799 }
800 fn keywords(&self) -> Vec<&'static str> {
801 vec![
802 "spark client",
803 "spark history server",
804 "spark rest api",
805 "spark jobs",
806 "spark ui",
807 ]
808 }
809 fn description(&self) -> &'static str {
810 "Apache Spark monitoring REST API client: reads applications, jobs, stages, task summaries, executors, environment and SQL executions from a live driver UI or a History Server"
811 }
812 fn get_async_actions(&self, _: &AppState) -> Vec<ActionDefinition> {
813 actions()
814 }
815 fn get_sync_actions(&self) -> Vec<ActionDefinition> {
816 actions()
817 }
818 fn get_event_types(&self) -> Vec<EventType> {
819 vec![
820 CONNECTED_EVENT.clone(),
821 APPLICATIONS_EVENT.clone(),
822 JOBS_EVENT.clone(),
823 STAGES_EVENT.clone(),
824 TASK_SUMMARY_EVENT.clone(),
825 EXECUTORS_EVENT.clone(),
826 ENVIRONMENT_EVENT.clone(),
827 SQL_EVENT.clone(),
828 VERSION_EVENT.clone(),
829 ERROR_EVENT.clone(),
830 ]
831 }
832 fn get_startup_parameters(&self) -> Vec<ParameterDefinition> {
833 vec![
834 param(
835 "endpoint_url",
836 "string",
837 "The REST API's base URL, used exactly as given: http://127.0.0.1:4040 (a live driver UI) or http://127.0.0.1:18080 (a History Server), or a proxy path such as https://gateway/sparkhistory. Without it the client's remote_addr is the target (host:port means http://host:port). There is no default target",
838 false,
839 json!("http://127.0.0.1:18080"),
840 None,
841 ),
842 param(
843 "request_timeout_secs",
844 "number",
845 "Seconds one request may take (1..=600)",
846 false,
847 json!(30),
848 Some(json!(super::REQUEST_TIMEOUT.as_secs())),
849 ),
850 ]
851 }
852 fn metadata(&self) -> crate::protocol::metadata::ProtocolMetadataV2 {
853 use crate::protocol::metadata::*;
854 ProtocolMetadataV2::builder()
855 .state(DevelopmentState::Beta)
856 .privilege_requirement(PrivilegeRequirement::None)
857 .implementation("Apache Spark monitoring REST API (/api/v1) over HTTP(S), reqwest natively and the hyper transport in the browser build: GET version, applications (status/minDate/maxDate/minEndDate/maxEndDate/limit), an application, its jobs (status), a job, stages (status, details=false), a stage's attempts or one attempt (details=false), a stage attempt's taskSummary (quantiles), executors/allexecutors, environment, sql (details=false, planDescription=false, offset/length) and one SQL execution with its plan. Paths are built from validated ids, one percent-encoded segment each, on the base URL fixed at connect; redirects are not followed. Responses are bounded in size and JSON depth before parsing and projected onto a fixed set of fields per resource")
858 .llm_control("Which application, job, stage, stage attempt or SQL execution to read, with which filters and how many entries; it reads the decoded entries and the server's own error text")
859 .e2e_testing("tests/client/spark/real_server_test.rs runs the official Apache Spark 4.2.0 distribution (SHA-512 pinned by install_peers.sh) with the model mocked: a PySpark application with a one-stage job, a shuffled two-stage job, a DataFrame aggregation and a job that fails on purpose, read through its live driver UI (spark.ui.port) and, once it has stopped, through Spark's own History Server over the event log it wrote. Every job id, status and stage id, the stage the model picked out of a job with its task summary, the failed stage and its failure reason, the SQL execution and its jobs, the driver's task counts, the environment, the application's duration and every stage attempt (skipped ones included) are checked against what Spark itself recorded: SparkContext.statusTracker() for the live UI, the SparkListener events of the event log for the History Server. A recording TCP relay between NetGet and Spark asserts the exact request lines Spark received for the model's actions and that nothing else was sent. Spark's 404 for an unknown job, the follow-up bound and an LLM failure (Spark sees only the connect probe) run against the same peers. Fails rather than skips without NETGET_SPARK_HOME")
860 .notes("Beta against the client bar: the peer is Apache Spark itself, the reference implementation, in both of the roles that serve this API (a live driver UI and the History Server); the suite fails rather than skips; each session is a real chain of reads whose later requests are built from ids the model took out of earlier answers; and that the client acts on the model's answer is asserted on the wire, from the request lines Spark received. The API is read-only, so the effect asserted is the request and Spark's answer to it; Spark keeps no request log, which is why the test relays the connection. Only Spark 4.2.0 is tested; older releases' REST API is the same documented v1 surface, and a field one of them lacks is simply absent from the projection. Not implemented: the REST submission API (/v1/submissions on a standalone master), which runs arbitrary code on a cluster and is not something a monitoring client should do on a model's say-so; task lists, RDD storage, streaming statistics, event-log download and thread dumps")
861 .build()
862 }
863 fn example_prompt(&self) -> &'static str {
864 "Connect to the Spark History Server at http://127.0.0.1:18080 and tell me which jobs failed in the newest application"
865 }
866 fn get_startup_examples(&self) -> StartupExamples {
867 let llm = json!({
868 "type": "open_client",
869 "protocol": "Spark",
870 "remote_addr": "127.0.0.1:18080",
871 "startup_params": {"request_timeout_secs": 30},
872 "instruction": "Find the newest completed application and list its failed jobs"
873 });
874 let mut static_example = llm.clone();
875 static_example["event_handlers"] = json!([
876 {"event_pattern": "spark_connected", "handler": {"type": "static", "actions": [
877 {"type": "list_applications", "status": "completed", "limit": 5}
878 ]}},
879 {"event_pattern": "*", "handler": {"type": "static", "actions": []}}
880 ]);
881 let mut scripted = static_example.clone();
882 scripted["event_handlers"][1] = json!({"event_pattern": "spark_applications", "handler": {
883 "type": "script", "language": "python",
884 "code": "import json,sys\ne=json.load(sys.stdin)['event']\na=[{'type':'list_jobs','app_id':x['id'],'status':'failed'} for x in e.get('applications',[])[:3]]\nprint(json.dumps({'actions':a}))"
885 }});
886 StartupExamples::new(llm, scripted, static_example)
887 }
888 fn group_name(&self) -> &'static str {
889 "Monitoring"
890 }
891}
892
893impl Client for SparkClientProtocol {
894 fn connect(
895 &self,
896 ctx: ConnectContext,
897 ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<std::net::SocketAddr>> + Send>>
898 {
899 Box::pin(super::connect(ctx))
900 }
901 fn execute_action(&self, v: Value) -> Result<ClientActionResult> {
902 let call = request(&v)?;
903 if call.operation == "disconnect" {
904 return Ok(ClientActionResult::Disconnect);
905 }
906 Ok(ClientActionResult::Custom {
907 name: call.operation.clone(),
908 data: serde_json::to_value(&call)?,
909 })
910 }
911}