1#![allow(
13 clippy::expect_used,
14 clippy::unwrap_used,
15 clippy::panic,
16 clippy::missing_panics_doc
17)]
18
19pub mod a2a;
21pub mod budgets;
23pub mod costs;
25pub mod diagnostics;
27pub mod endpoints;
29pub mod mcp_client;
31pub mod mcp_server;
33pub mod reporting;
35pub mod runner;
37pub mod scenario;
39pub mod selectors;
41pub mod vision;
43
44#[cfg(feature = "a2a-server")]
46pub mod a2a_server;
47
48#[cfg(feature = "macros")]
50pub mod macros;
51
52use std::collections::HashMap;
53use std::time::Duration;
54
55use serde_json::{json, Value};
56
57pub use costs::LlmResponse;
58pub use costs::LlmUsage;
59
60#[derive(Debug, Clone)]
63pub struct LlmConfig {
64 pub url: String,
66 pub model: String,
68 pub api_key: Option<String>,
70 pub headers: HashMap<String, String>,
72 pub timeout: Duration,
74 pub temperature: f64,
76 pub thinking: Option<bool>,
79 pub model_params: HashMap<String, Value>,
82 pub max_attempts: u32,
87}
88
89impl LlmConfig {
90 #[must_use]
93 pub fn from_env() -> Self {
94 Self {
95 url: llm_base_url(),
96 model: llm_model(),
97 api_key: std::env::var("HARNESS_LLM_API_KEY").ok(),
98 headers: parse_headers_env(),
99 timeout: Duration::from_secs(60),
100 temperature: 0.0,
101 thinking: None,
102 model_params: HashMap::new(),
103 max_attempts: default_llm_attempts(),
104 }
105 }
106}
107
108#[must_use]
111pub fn default_llm_attempts() -> u32 {
112 std::env::var("HARNESS_LLM_CALL_ATTEMPTS")
113 .ok()
114 .and_then(|v| v.parse().ok())
115 .filter(|n| *n >= 1)
116 .unwrap_or(3)
117}
118
119#[must_use]
123pub fn parse_headers_env() -> HashMap<String, String> {
124 let Ok(raw) = std::env::var("HARNESS_LLM_HEADERS") else {
125 return HashMap::new();
126 };
127 let Ok(json) = serde_json::from_str::<Value>(&raw) else {
128 return HashMap::new();
129 };
130 let Some(obj) = json.as_object() else {
131 return HashMap::new();
132 };
133 obj.iter()
134 .filter_map(|(k, v)| v.as_str().map(|s| (k.clone(), s.to_owned())))
135 .collect()
136}
137
138#[must_use]
141pub fn base_url() -> String {
142 std::env::var("HARNESS_BROWSER_BASE_URL").unwrap_or_else(|_| "http://localhost:4200".to_owned())
143}
144
145#[must_use]
148pub fn llm_base_url() -> String {
149 std::env::var("HARNESS_LLM_TEST_URL")
150 .unwrap_or_else(|_| "http://localhost:8080".to_owned())
151 .trim_end_matches('/')
152 .to_owned()
153}
154
155#[must_use]
158pub fn llm_model() -> String {
159 std::env::var("HARNESS_LLM_TEST_MODEL").unwrap_or_else(|_| "deepseek".to_owned())
160}
161
162#[must_use]
165pub fn browser_headless() -> bool {
166 std::env::var("HARNESS_BROWSER_HEADLESS")
167 .map_or(true, |v| v != "0" && v.to_lowercase() != "false")
168}
169
170#[must_use]
172pub fn http_client(timeout: Duration) -> reqwest::Client {
173 reqwest::Client::builder()
174 .timeout(timeout)
175 .build()
176 .expect("build reqwest client")
177}
178
179#[must_use]
185pub async fn llm_chat(llm: &LlmConfig, system: &str, user: &str) -> Option<String> {
186 llm_chat_with_usage(llm, system, user)
187 .await
188 .map(|r| r.content)
189 .ok()
190}
191
192pub async fn llm_chat_with_usage(
212 llm: &LlmConfig,
213 system: &str,
214 user: &str,
215) -> Result<LlmResponse, String> {
216 chat_with_retry(llm, system, user, None).await
217}
218
219pub async fn llm_chat_vision_with_usage(
231 llm: &LlmConfig,
232 system: &str,
233 user: &str,
234 image_data_url: &str,
235) -> Result<LlmResponse, String> {
236 chat_with_retry(llm, system, user, Some(image_data_url)).await
237}
238
239pub async fn llm_chat_with_usage_chain(
251 primary: &LlmConfig,
252 fallbacks: &[LlmConfig],
253 system: &str,
254 user: &str,
255) -> Result<(LlmResponse, usize), String> {
256 chat_chain_with_retry(primary, fallbacks, system, user, None).await
257}
258
259pub async fn llm_chat_vision_with_usage_chain(
265 primary: &LlmConfig,
266 fallbacks: &[LlmConfig],
267 system: &str,
268 user: &str,
269 image_data_url: &str,
270) -> Result<(LlmResponse, usize), String> {
271 chat_chain_with_retry(primary, fallbacks, system, user, Some(image_data_url)).await
272}
273
274async fn chat_chain_with_retry(
277 primary: &LlmConfig,
278 fallbacks: &[LlmConfig],
279 system: &str,
280 user: &str,
281 image_data_url: Option<&str>,
282) -> Result<(LlmResponse, usize), String> {
283 let mut failures: Vec<String> = Vec::new();
284 for (i, llm) in std::iter::once(primary).chain(fallbacks.iter()).enumerate() {
285 match chat_with_retry(llm, system, user, image_data_url).await {
286 Ok(resp) => return Ok((resp, i)),
287 Err(e) => failures.push(format!("endpoint '{}' ({:?}): {e}", llm.url, llm.model)),
288 }
289 }
290 let details = failures.iter().fold(String::new(), |mut acc, f| {
291 use std::fmt::Write as _;
292 let _ = writeln!(acc, " - {f}");
293 acc
294 });
295 Err(format!(
296 "LLM call failed on all {} endpoint(s):\n{details}",
297 failures.len()
298 ))
299}
300
301async fn chat_with_retry(
303 llm: &LlmConfig,
304 system: &str,
305 user: &str,
306 image_data_url: Option<&str>,
307) -> Result<LlmResponse, String> {
308 let client = http_client(llm.timeout);
309 let mut last_err = String::from("LLM call failed");
310 let mut attempts: u32 = 0;
311
312 while attempts < llm.max_attempts {
313 attempts += 1;
314 match llm_chat_once(&client, llm, system, user, image_data_url).await {
315 Ok(resp) => return Ok(resp),
316 Err(err) => {
317 let backoff = match &err {
321 LlmCallError::EmptyBody { .. } => Duration::from_secs(3),
322 _ => Duration::from_millis(500 * u64::from(attempts)),
323 };
324 last_err = err.to_string();
325 if attempts >= llm.max_attempts || !err.is_retryable() {
326 break;
327 }
328 tokio::time::sleep(backoff).await;
329 }
330 }
331 }
332
333 Err(format!(
334 "LLM call failed after {attempts} attempt(s) (endpoint {url}): {last_err}",
335 url = llm.url
336 ))
337}
338
339#[must_use]
343fn build_messages(system: &str, user: &str, image_data_url: Option<&str>) -> Value {
344 let user_content = image_data_url.map_or_else(
345 || Value::String(user.to_owned()),
346 |url| {
347 json!([
348 {"type": "text", "text": user},
349 {"type": "image_url", "image_url": {"url": url}}
350 ])
351 },
352 );
353 json!([
354 {"role": "system", "content": system},
355 {"role": "user", "content": user_content}
356 ])
357}
358
359enum LlmCallError {
363 Transport { message: String },
365 Http { status: u16, body: String },
367 InvalidJson {
369 status: u16,
370 detail: String,
371 body: String,
372 },
373 EmptyBody { status: u16 },
377 MissingContent { json: String },
379}
380
381impl std::fmt::Display for LlmCallError {
382 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
383 match self {
384 Self::Transport { message } => write!(f, "LLM HTTP request failed: {message}"),
385 Self::Http { status, body } => {
386 write!(
387 f,
388 "LLM endpoint returned HTTP {status}: {}",
389 truncate(body, 300)
390 )
391 }
392 Self::InvalidJson {
393 status,
394 detail,
395 body,
396 } => write!(
397 f,
398 "LLM endpoint returned HTTP {status} with non-JSON body ({detail}): {}",
399 truncate(body, 300)
400 ),
401 Self::EmptyBody { status } => write!(
402 f,
403 "LLM endpoint returned HTTP {status} with an empty response (likely gateway warm-up)"
404 ),
405 Self::MissingContent { json } => write!(
406 f,
407 "LLM response missing choices[0].message.content: {}",
408 truncate(json, 300)
409 ),
410 }
411 }
412}
413
414impl LlmCallError {
415 #[must_use]
419 fn is_retryable(&self) -> bool {
420 match self {
421 Self::Transport { .. } | Self::MissingContent { .. } | Self::EmptyBody { .. } => true,
422 Self::Http { status, .. } => {
423 *status == 408 || *status == 429 || (500..600).contains(status)
424 }
425 Self::InvalidJson { status, .. } => {
426 *status == 200 || *status == 408 || *status == 429 || (500..600).contains(status)
427 }
428 }
429 }
430}
431
432async fn llm_chat_once(
434 client: &reqwest::Client,
435 llm: &LlmConfig,
436 system: &str,
437 user: &str,
438 image_data_url: Option<&str>,
439) -> Result<LlmResponse, LlmCallError> {
440 let mut payload = serde_json::json!({
441 "model": llm.model,
442 "messages": build_messages(system, user, image_data_url),
443 "max_tokens": 4096,
444 "temperature": llm.temperature
445 });
446 if let Some(think) = llm.thinking {
447 if think {
448 payload["thinking"] = serde_json::json!({"type": "enabled"});
449 } else {
450 payload["thinking"] = serde_json::json!({"type": "disabled"});
451 }
452 }
453 if !llm.model_params.is_empty() {
455 if let Value::Object(ref mut map) = payload {
456 for (key, val) in &llm.model_params {
457 map.insert(key.clone(), val.clone());
458 }
459 }
460 }
461
462 let mut req = client
463 .post(format!("{}/v1/chat/completions", llm.url))
464 .header("Content-Type", "application/json");
465
466 if let Some(ref key) = llm.api_key {
467 req = req.header("Authorization", format!("Bearer {key}"));
468 }
469 for (name, value) in &llm.headers {
470 req = req.header(name.as_str(), value.as_str());
471 }
472
473 let resp = req
474 .json(&payload)
475 .send()
476 .await
477 .map_err(|e| LlmCallError::Transport {
478 message: e.to_string(),
479 })?;
480 let status = resp.status();
481 let status_u16 = status.as_u16();
482 let body = resp.text().await.unwrap_or_default();
483 if !status.is_success() {
484 return Err(LlmCallError::Http {
485 status: status_u16,
486 body,
487 });
488 }
489 if body.trim().is_empty() {
490 return Err(LlmCallError::EmptyBody { status: status_u16 });
493 }
494 let json: Value = match serde_json::from_str(&body) {
495 Ok(v) => v,
496 Err(e) => {
497 return Err(LlmCallError::InvalidJson {
498 status: status_u16,
499 detail: e.to_string(),
500 body,
501 });
502 }
503 };
504 let usage = costs::extract_usage(&json);
505 let content = json["choices"][0]["message"]["content"]
506 .as_str()
507 .map(String::from)
508 .ok_or_else(|| LlmCallError::MissingContent {
509 json: json.to_string(),
510 })?;
511
512 Ok(LlmResponse { content, usage })
513}
514
515pub const DOM_EXTRACT_JS: &str = r#"
518(() => {
519 const interactive = 'a, button, input, textarea, select, [role="button"], [onclick], [tabindex], [data-testid], [aria-label]';
520 const els = document.querySelectorAll(interactive);
521 const info = [];
522 const seen = new Set();
523 els.forEach((el, i) => {
524 const rect = el.getBoundingClientRect();
525 if (rect.width === 0 || rect.height === 0) return;
526 const tag = el.tagName.toLowerCase();
527 let selector = '';
528 if (el.id) selector = '#' + CSS.escape(el.id);
529 else if (el.getAttribute('data-testid')) selector = '[data-testid="' + el.getAttribute('data-testid') + '"]';
530 else if (el.name) selector = '[name="' + CSS.escape(el.name) + '"]';
531 else if (el.className && typeof el.className === 'string') {
532 const cls = el.className.trim().split(/\\s+/)[0];
533 if (cls) selector = tag + '.' + CSS.escape(cls);
534 }
535 if (!selector) selector = tag;
536 if (seen.has(selector)) return;
537 seen.add(selector);
538
539 let label = '';
540 const aria = el.getAttribute('aria-label');
541 if (aria) {
542 label = aria;
543 } else if (tag === 'input' || tag === 'textarea' || tag === 'select') {
544 label = el.placeholder || el.name || el.getAttribute('aria-label') || '';
545 if (el.type && !label) label = el.type;
546 } else {
547 label = (el.textContent || '').trim().substring(0, 80);
548 }
549
550 info.push(i + ': ' + selector + ' [' + tag + '] "' + label + '"');
551 });
552 return JSON.stringify(info);
553})()
554"#;
555
556#[must_use]
559pub fn truncate(s: &str, max_len: usize) -> String {
560 if s.len() <= max_len {
561 s.to_owned()
562 } else {
563 format!("{}...<truncated>", &s[..max_len])
564 }
565}
566
567#[cfg(test)]
568mod tests {
569 use crate::costs::extract_usage;
570 use crate::truncate;
571 use crate::{default_llm_attempts, llm_base_url, llm_model, parse_headers_env, LlmConfig};
572
573 fn mock_llm_server(status: u16, body: &'static str) -> String {
576 use std::io::{Read, Write};
577 use std::net::TcpListener;
578 let listener = TcpListener::bind("127.0.0.1:0").unwrap();
579 let addr = listener.local_addr().unwrap();
580 std::thread::spawn(move || {
581 for stream in listener.incoming() {
582 let Ok(mut stream) = stream else { break };
583 let mut buf = [0u8; 4096];
584 let _ = stream.read(&mut buf);
585 let resp = format!(
586 "HTTP/1.1 {status} {}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
587 if status == 200 { "OK" } else { "ERROR" },
588 body.len(),
589 body
590 );
591 let _ = stream.write_all(resp.as_bytes());
592 }
593 });
594 format!("http://{addr}")
595 }
596
597 const PASS_BODY: &str = r#"{"choices":[{"message":{"content":"PASS"}}],"usage":{"prompt_tokens":7,"completion_tokens":2}}"#;
598
599 fn cfg(url: &str, attempts: u32) -> LlmConfig {
600 LlmConfig {
601 url: url.to_owned(),
602 model: "mock".to_owned(),
603 api_key: None,
604 headers: std::collections::HashMap::new(),
605 timeout: std::time::Duration::from_secs(10),
606 temperature: 0.0,
607 thinking: None,
608 model_params: std::collections::HashMap::new(),
609 max_attempts: attempts,
610 }
611 }
612
613 #[tokio::test]
614 async fn test_chain_primary_success_returns_index_zero() {
615 let good = mock_llm_server(200, PASS_BODY);
616 let (resp, idx) = crate::llm_chat_with_usage_chain(&cfg(&good, 2), &[], "s", "u")
617 .await
618 .expect("primary endpoint should answer");
619 assert_eq!(idx, 0);
620 assert_eq!(resp.content, "PASS");
621 assert_eq!(resp.usage.prompt_tokens, 7);
622 }
623
624 #[tokio::test]
625 async fn test_chain_falls_back_on_empty_200() {
626 let broken = mock_llm_server(200, "");
630 let good = mock_llm_server(200, PASS_BODY);
631 let (resp, idx) =
632 crate::llm_chat_with_usage_chain(&cfg(&broken, 2), &[cfg(&good, 2)], "s", "u")
633 .await
634 .expect("fallback endpoint should answer");
635 assert_eq!(idx, 1);
636 assert_eq!(resp.content, "PASS");
637 }
638
639 #[tokio::test]
640 async fn test_chain_reports_all_endpoints_on_total_failure() {
641 let broken1 = mock_llm_server(200, "");
642 let broken2 = mock_llm_server(503, "unavailable");
643 let err =
644 crate::llm_chat_with_usage_chain(&cfg(&broken1, 2), &[cfg(&broken2, 2)], "s", "u")
645 .await
646 .expect_err("both endpoints fail");
647 assert!(err.contains("all 2 endpoint(s)"), "got: {err}");
648 assert!(err.contains(&broken1), "primary URL missing: {err}");
649 assert!(err.contains(&broken2), "fallback URL missing: {err}");
650 }
651
652 #[test]
653 fn test_default_llm_attempts_env() {
654 std::env::set_var("HARNESS_LLM_CALL_ATTEMPTS", "7");
655 assert_eq!(default_llm_attempts(), 7);
656 std::env::set_var("HARNESS_LLM_CALL_ATTEMPTS", "0");
657 assert_eq!(default_llm_attempts(), 3, "0 must fall back to default");
658 std::env::set_var("HARNESS_LLM_CALL_ATTEMPTS", "junk");
659 assert_eq!(default_llm_attempts(), 3, "non-numeric must fall back");
660 std::env::remove_var("HARNESS_LLM_CALL_ATTEMPTS");
661 assert_eq!(default_llm_attempts(), 3);
662 }
663
664 #[test]
665 fn test_truncate_short() {
666 assert_eq!(truncate("hello", 10), "hello");
667 }
668
669 #[test]
670 fn test_truncate_long() {
671 let result = truncate("hello world", 5);
672 assert!(result.contains("<truncated>"));
673 assert!(result.starts_with("hello"));
674 }
675
676 #[test]
677 fn test_truncate_exact_length() {
678 assert_eq!(truncate("abcde", 5), "abcde");
679 }
680
681 #[test]
682 fn test_truncate_empty() {
683 assert_eq!(truncate("", 5), "");
684 }
685
686 #[test]
687 fn test_parse_headers_env_empty() {
688 std::env::remove_var("HARNESS_LLM_HEADERS");
689 let h = parse_headers_env();
690 assert!(h.is_empty());
691 }
692
693 #[test]
694 fn test_parse_headers_env_valid() {
695 std::env::set_var("HARNESS_LLM_HEADERS", r#"{"X-Org":"acme","X-Version":"1"}"#);
696 let h = parse_headers_env();
697 assert_eq!(h.get("X-Org").map(String::as_str), Some("acme"));
698 assert_eq!(h.get("X-Version").map(String::as_str), Some("1"));
699 std::env::remove_var("HARNESS_LLM_HEADERS");
700 }
701
702 #[test]
703 fn test_parse_headers_env_invalid_json() {
704 std::env::set_var("HARNESS_LLM_HEADERS", "not-json");
705 let h = parse_headers_env();
706 assert!(h.is_empty());
707 std::env::remove_var("HARNESS_LLM_HEADERS");
708 }
709
710 #[test]
711 fn test_llm_config_from_env_defaults() {
712 #[allow(clippy::float_cmp)]
713 {
714 let config = LlmConfig::from_env();
715 assert_eq!(config.temperature, 0.0);
716 assert!(config.thinking.is_none());
717 assert!(config.model_params.is_empty());
718 }
719 }
720
721 #[test]
722 fn test_extract_usage_full() {
723 let json = serde_json::json!({
724 "usage": {
725 "prompt_tokens": 100,
726 "completion_tokens": 200,
727 "total_tokens": 300
728 }
729 });
730 let usage = extract_usage(&json);
731 assert_eq!(usage.prompt_tokens, 100);
732 assert_eq!(usage.completion_tokens, 200);
733 assert_eq!(usage.total_tokens, 300);
734 }
735
736 #[test]
737 fn test_extract_usage_empty() {
738 let json = serde_json::json!({});
739 let usage = extract_usage(&json);
740 assert_eq!(usage.prompt_tokens, 0);
741 assert_eq!(usage.completion_tokens, 0);
742 assert_eq!(usage.total_tokens, 0);
743 }
744
745 #[test]
746 fn test_truncate_unicode() {
747 assert_eq!(truncate("héllo", 3), "hé...<truncated>");
750 assert_eq!(truncate("hello", 5), "hello");
752 }
753
754 #[test]
755 fn test_parse_headers_env_non_object() {
756 std::env::set_var("HARNESS_LLM_HEADERS", "[1, 2, 3]");
757 let h = parse_headers_env();
758 assert!(h.is_empty());
759 std::env::remove_var("HARNESS_LLM_HEADERS");
760 }
761
762 #[test]
763 fn test_parse_headers_env_nested_values_filtered() {
764 std::env::set_var(
765 "HARNESS_LLM_HEADERS",
766 r#"{"str":"val","num":42,"bool":true}"#,
767 );
768 let h = parse_headers_env();
769 assert_eq!(h.get("str").map(String::as_str), Some("val"));
770 assert!(!h.contains_key("num"));
771 assert!(!h.contains_key("bool"));
772 std::env::remove_var("HARNESS_LLM_HEADERS");
773 }
774
775 #[test]
776 fn test_llm_config_has_default_model() {
777 let config = LlmConfig::from_env();
778 assert!(!config.model.is_empty());
779 }
780
781 #[test]
782 fn test_llm_base_url_default() {
783 std::env::remove_var("HARNESS_LLM_TEST_URL");
784 let url = llm_base_url();
785 assert_eq!(url, "http://localhost:8080");
786 }
787
788 #[test]
789 fn test_llm_base_url_custom() {
790 std::env::set_var("HARNESS_LLM_TEST_URL", "https://custom.api.com/v1");
791 let url = llm_base_url();
792 assert_eq!(url, "https://custom.api.com/v1");
793 std::env::remove_var("HARNESS_LLM_TEST_URL");
794 }
795
796 #[test]
797 fn test_llm_base_url_trailing_slash() {
798 std::env::set_var("HARNESS_LLM_TEST_URL", "https://api.com/");
799 let url = llm_base_url();
800 assert_eq!(url, "https://api.com");
801 std::env::remove_var("HARNESS_LLM_TEST_URL");
802 }
803
804 #[test]
805 fn test_llm_model_default() {
806 std::env::remove_var("HARNESS_LLM_TEST_MODEL");
807 assert_eq!(llm_model(), "deepseek");
808 }
809
810 #[test]
811 fn test_llm_model_custom() {
812 std::env::set_var("HARNESS_LLM_TEST_MODEL", "gpt-4o");
813 assert_eq!(llm_model(), "gpt-4o");
814 std::env::remove_var("HARNESS_LLM_TEST_MODEL");
815 }
816
817 #[test]
818 fn test_extract_usage_partial() {
819 let json = serde_json::json!({
820 "usage": {
821 "prompt_tokens": 50
822 }
823 });
824 let usage = extract_usage(&json);
825 assert_eq!(usage.prompt_tokens, 50);
826 assert_eq!(usage.completion_tokens, 0);
827 assert_eq!(usage.total_tokens, 0);
828 }
829
830 #[test]
831 fn test_browser_headless_default() {
832 std::env::remove_var("HARNESS_BROWSER_HEADLESS");
833 assert!(crate::browser_headless());
834 }
835}