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}
83
84impl LlmConfig {
85 #[must_use]
88 pub fn from_env() -> Self {
89 Self {
90 url: llm_base_url(),
91 model: llm_model(),
92 api_key: std::env::var("HARNESS_LLM_API_KEY").ok(),
93 headers: parse_headers_env(),
94 timeout: Duration::from_secs(60),
95 temperature: 0.0,
96 thinking: None,
97 model_params: HashMap::new(),
98 }
99 }
100}
101
102#[must_use]
106pub fn parse_headers_env() -> HashMap<String, String> {
107 let Ok(raw) = std::env::var("HARNESS_LLM_HEADERS") else {
108 return HashMap::new();
109 };
110 let Ok(json) = serde_json::from_str::<Value>(&raw) else {
111 return HashMap::new();
112 };
113 let Some(obj) = json.as_object() else {
114 return HashMap::new();
115 };
116 obj.iter()
117 .filter_map(|(k, v)| v.as_str().map(|s| (k.clone(), s.to_owned())))
118 .collect()
119}
120
121#[must_use]
124pub fn base_url() -> String {
125 std::env::var("HARNESS_BROWSER_BASE_URL").unwrap_or_else(|_| "http://localhost:4200".to_owned())
126}
127
128#[must_use]
131pub fn llm_base_url() -> String {
132 std::env::var("HARNESS_LLM_TEST_URL")
133 .unwrap_or_else(|_| "http://localhost:8080".to_owned())
134 .trim_end_matches('/')
135 .to_owned()
136}
137
138#[must_use]
141pub fn llm_model() -> String {
142 std::env::var("HARNESS_LLM_TEST_MODEL").unwrap_or_else(|_| "deepseek".to_owned())
143}
144
145#[must_use]
148pub fn browser_headless() -> bool {
149 std::env::var("HARNESS_BROWSER_HEADLESS")
150 .map_or(true, |v| v != "0" && v.to_lowercase() != "false")
151}
152
153#[must_use]
155pub fn http_client(timeout: Duration) -> reqwest::Client {
156 reqwest::Client::builder()
157 .timeout(timeout)
158 .build()
159 .expect("build reqwest client")
160}
161
162#[must_use]
168pub async fn llm_chat(llm: &LlmConfig, system: &str, user: &str) -> Option<String> {
169 llm_chat_with_usage(llm, system, user)
170 .await
171 .map(|r| r.content)
172 .ok()
173}
174
175pub async fn llm_chat_with_usage(
195 llm: &LlmConfig,
196 system: &str,
197 user: &str,
198) -> Result<LlmResponse, String> {
199 chat_with_retry(llm, system, user, None).await
200}
201
202pub async fn llm_chat_vision_with_usage(
214 llm: &LlmConfig,
215 system: &str,
216 user: &str,
217 image_data_url: &str,
218) -> Result<LlmResponse, String> {
219 chat_with_retry(llm, system, user, Some(image_data_url)).await
220}
221
222async fn chat_with_retry(
224 llm: &LlmConfig,
225 system: &str,
226 user: &str,
227 image_data_url: Option<&str>,
228) -> Result<LlmResponse, String> {
229 let client = http_client(llm.timeout);
230 let mut last_err = String::from("LLM call failed");
231 let mut attempts: u32 = 0;
232
233 while attempts < LLM_CALL_ATTEMPTS {
234 attempts += 1;
235 match llm_chat_once(&client, llm, system, user, image_data_url).await {
236 Ok(resp) => return Ok(resp),
237 Err(err) => {
238 let backoff = match &err {
242 LlmCallError::EmptyBody { .. } => Duration::from_secs(3),
243 _ => Duration::from_millis(500 * u64::from(attempts)),
244 };
245 last_err = err.to_string();
246 if attempts >= LLM_CALL_ATTEMPTS || !err.is_retryable() {
247 break;
248 }
249 tokio::time::sleep(backoff).await;
250 }
251 }
252 }
253
254 Err(format!(
255 "LLM call failed after {attempts} attempt(s) (endpoint {url}): {last_err}",
256 url = llm.url
257 ))
258}
259
260const LLM_CALL_ATTEMPTS: u32 = 3;
262
263#[must_use]
267fn build_messages(system: &str, user: &str, image_data_url: Option<&str>) -> Value {
268 let user_content = image_data_url.map_or_else(
269 || Value::String(user.to_owned()),
270 |url| {
271 json!([
272 {"type": "text", "text": user},
273 {"type": "image_url", "image_url": {"url": url}}
274 ])
275 },
276 );
277 json!([
278 {"role": "system", "content": system},
279 {"role": "user", "content": user_content}
280 ])
281}
282
283enum LlmCallError {
287 Transport { message: String },
289 Http { status: u16, body: String },
291 InvalidJson {
293 status: u16,
294 detail: String,
295 body: String,
296 },
297 EmptyBody { status: u16 },
301 MissingContent { json: String },
303}
304
305impl std::fmt::Display for LlmCallError {
306 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
307 match self {
308 Self::Transport { message } => write!(f, "LLM HTTP request failed: {message}"),
309 Self::Http { status, body } => {
310 write!(
311 f,
312 "LLM endpoint returned HTTP {status}: {}",
313 truncate(body, 300)
314 )
315 }
316 Self::InvalidJson {
317 status,
318 detail,
319 body,
320 } => write!(
321 f,
322 "LLM endpoint returned HTTP {status} with non-JSON body ({detail}): {}",
323 truncate(body, 300)
324 ),
325 Self::EmptyBody { status } => write!(
326 f,
327 "LLM endpoint returned HTTP {status} with an empty response (likely gateway warm-up)"
328 ),
329 Self::MissingContent { json } => write!(
330 f,
331 "LLM response missing choices[0].message.content: {}",
332 truncate(json, 300)
333 ),
334 }
335 }
336}
337
338impl LlmCallError {
339 #[must_use]
343 fn is_retryable(&self) -> bool {
344 match self {
345 Self::Transport { .. } | Self::MissingContent { .. } | Self::EmptyBody { .. } => true,
346 Self::Http { status, .. } => {
347 *status == 408 || *status == 429 || (500..600).contains(status)
348 }
349 Self::InvalidJson { status, .. } => {
350 *status == 200 || *status == 408 || *status == 429 || (500..600).contains(status)
351 }
352 }
353 }
354}
355
356async fn llm_chat_once(
358 client: &reqwest::Client,
359 llm: &LlmConfig,
360 system: &str,
361 user: &str,
362 image_data_url: Option<&str>,
363) -> Result<LlmResponse, LlmCallError> {
364 let mut payload = serde_json::json!({
365 "model": llm.model,
366 "messages": build_messages(system, user, image_data_url),
367 "max_tokens": 4096,
368 "temperature": llm.temperature
369 });
370 if let Some(think) = llm.thinking {
371 if think {
372 payload["thinking"] = serde_json::json!({"type": "enabled"});
373 } else {
374 payload["thinking"] = serde_json::json!({"type": "disabled"});
375 }
376 }
377 if !llm.model_params.is_empty() {
379 if let Value::Object(ref mut map) = payload {
380 for (key, val) in &llm.model_params {
381 map.insert(key.clone(), val.clone());
382 }
383 }
384 }
385
386 let mut req = client
387 .post(format!("{}/v1/chat/completions", llm.url))
388 .header("Content-Type", "application/json");
389
390 if let Some(ref key) = llm.api_key {
391 req = req.header("Authorization", format!("Bearer {key}"));
392 }
393 for (name, value) in &llm.headers {
394 req = req.header(name.as_str(), value.as_str());
395 }
396
397 let resp = req
398 .json(&payload)
399 .send()
400 .await
401 .map_err(|e| LlmCallError::Transport {
402 message: e.to_string(),
403 })?;
404 let status = resp.status();
405 let status_u16 = status.as_u16();
406 let body = resp.text().await.unwrap_or_default();
407 if !status.is_success() {
408 return Err(LlmCallError::Http {
409 status: status_u16,
410 body,
411 });
412 }
413 if body.trim().is_empty() {
414 return Err(LlmCallError::EmptyBody { status: status_u16 });
417 }
418 let json: Value = match serde_json::from_str(&body) {
419 Ok(v) => v,
420 Err(e) => {
421 return Err(LlmCallError::InvalidJson {
422 status: status_u16,
423 detail: e.to_string(),
424 body,
425 });
426 }
427 };
428 let usage = costs::extract_usage(&json);
429 let content = json["choices"][0]["message"]["content"]
430 .as_str()
431 .map(String::from)
432 .ok_or_else(|| LlmCallError::MissingContent {
433 json: json.to_string(),
434 })?;
435
436 Ok(LlmResponse { content, usage })
437}
438
439pub const DOM_EXTRACT_JS: &str = r#"
442(() => {
443 const interactive = 'a, button, input, textarea, select, [role="button"], [onclick], [tabindex], [data-testid], [aria-label]';
444 const els = document.querySelectorAll(interactive);
445 const info = [];
446 const seen = new Set();
447 els.forEach((el, i) => {
448 const rect = el.getBoundingClientRect();
449 if (rect.width === 0 || rect.height === 0) return;
450 const tag = el.tagName.toLowerCase();
451 let selector = '';
452 if (el.id) selector = '#' + CSS.escape(el.id);
453 else if (el.getAttribute('data-testid')) selector = '[data-testid="' + el.getAttribute('data-testid') + '"]';
454 else if (el.name) selector = '[name="' + CSS.escape(el.name) + '"]';
455 else if (el.className && typeof el.className === 'string') {
456 const cls = el.className.trim().split(/\\s+/)[0];
457 if (cls) selector = tag + '.' + CSS.escape(cls);
458 }
459 if (!selector) selector = tag;
460 if (seen.has(selector)) return;
461 seen.add(selector);
462
463 let label = '';
464 const aria = el.getAttribute('aria-label');
465 if (aria) {
466 label = aria;
467 } else if (tag === 'input' || tag === 'textarea' || tag === 'select') {
468 label = el.placeholder || el.name || el.getAttribute('aria-label') || '';
469 if (el.type && !label) label = el.type;
470 } else {
471 label = (el.textContent || '').trim().substring(0, 80);
472 }
473
474 info.push(i + ': ' + selector + ' [' + tag + '] "' + label + '"');
475 });
476 return JSON.stringify(info);
477})()
478"#;
479
480#[must_use]
483pub fn truncate(s: &str, max_len: usize) -> String {
484 if s.len() <= max_len {
485 s.to_owned()
486 } else {
487 format!("{}...<truncated>", &s[..max_len])
488 }
489}
490
491#[cfg(test)]
492mod tests {
493 use crate::costs::extract_usage;
494 use crate::truncate;
495 use crate::{llm_base_url, llm_model, parse_headers_env, LlmConfig};
496
497 #[test]
498 fn test_truncate_short() {
499 assert_eq!(truncate("hello", 10), "hello");
500 }
501
502 #[test]
503 fn test_truncate_long() {
504 let result = truncate("hello world", 5);
505 assert!(result.contains("<truncated>"));
506 assert!(result.starts_with("hello"));
507 }
508
509 #[test]
510 fn test_truncate_exact_length() {
511 assert_eq!(truncate("abcde", 5), "abcde");
512 }
513
514 #[test]
515 fn test_truncate_empty() {
516 assert_eq!(truncate("", 5), "");
517 }
518
519 #[test]
520 fn test_parse_headers_env_empty() {
521 std::env::remove_var("HARNESS_LLM_HEADERS");
522 let h = parse_headers_env();
523 assert!(h.is_empty());
524 }
525
526 #[test]
527 fn test_parse_headers_env_valid() {
528 std::env::set_var("HARNESS_LLM_HEADERS", r#"{"X-Org":"acme","X-Version":"1"}"#);
529 let h = parse_headers_env();
530 assert_eq!(h.get("X-Org").map(String::as_str), Some("acme"));
531 assert_eq!(h.get("X-Version").map(String::as_str), Some("1"));
532 std::env::remove_var("HARNESS_LLM_HEADERS");
533 }
534
535 #[test]
536 fn test_parse_headers_env_invalid_json() {
537 std::env::set_var("HARNESS_LLM_HEADERS", "not-json");
538 let h = parse_headers_env();
539 assert!(h.is_empty());
540 std::env::remove_var("HARNESS_LLM_HEADERS");
541 }
542
543 #[test]
544 fn test_llm_config_from_env_defaults() {
545 #[allow(clippy::float_cmp)]
546 {
547 let config = LlmConfig::from_env();
548 assert_eq!(config.temperature, 0.0);
549 assert!(config.thinking.is_none());
550 assert!(config.model_params.is_empty());
551 }
552 }
553
554 #[test]
555 fn test_extract_usage_full() {
556 let json = serde_json::json!({
557 "usage": {
558 "prompt_tokens": 100,
559 "completion_tokens": 200,
560 "total_tokens": 300
561 }
562 });
563 let usage = extract_usage(&json);
564 assert_eq!(usage.prompt_tokens, 100);
565 assert_eq!(usage.completion_tokens, 200);
566 assert_eq!(usage.total_tokens, 300);
567 }
568
569 #[test]
570 fn test_extract_usage_empty() {
571 let json = serde_json::json!({});
572 let usage = extract_usage(&json);
573 assert_eq!(usage.prompt_tokens, 0);
574 assert_eq!(usage.completion_tokens, 0);
575 assert_eq!(usage.total_tokens, 0);
576 }
577
578 #[test]
579 fn test_truncate_unicode() {
580 assert_eq!(truncate("héllo", 3), "hé...<truncated>");
583 assert_eq!(truncate("hello", 5), "hello");
585 }
586
587 #[test]
588 fn test_parse_headers_env_non_object() {
589 std::env::set_var("HARNESS_LLM_HEADERS", "[1, 2, 3]");
590 let h = parse_headers_env();
591 assert!(h.is_empty());
592 std::env::remove_var("HARNESS_LLM_HEADERS");
593 }
594
595 #[test]
596 fn test_parse_headers_env_nested_values_filtered() {
597 std::env::set_var(
598 "HARNESS_LLM_HEADERS",
599 r#"{"str":"val","num":42,"bool":true}"#,
600 );
601 let h = parse_headers_env();
602 assert_eq!(h.get("str").map(String::as_str), Some("val"));
603 assert!(!h.contains_key("num"));
604 assert!(!h.contains_key("bool"));
605 std::env::remove_var("HARNESS_LLM_HEADERS");
606 }
607
608 #[test]
609 fn test_llm_config_has_default_model() {
610 let config = LlmConfig::from_env();
611 assert!(!config.model.is_empty());
612 }
613
614 #[test]
615 fn test_llm_base_url_default() {
616 std::env::remove_var("HARNESS_LLM_TEST_URL");
617 let url = llm_base_url();
618 assert_eq!(url, "http://localhost:8080");
619 }
620
621 #[test]
622 fn test_llm_base_url_custom() {
623 std::env::set_var("HARNESS_LLM_TEST_URL", "https://custom.api.com/v1");
624 let url = llm_base_url();
625 assert_eq!(url, "https://custom.api.com/v1");
626 std::env::remove_var("HARNESS_LLM_TEST_URL");
627 }
628
629 #[test]
630 fn test_llm_base_url_trailing_slash() {
631 std::env::set_var("HARNESS_LLM_TEST_URL", "https://api.com/");
632 let url = llm_base_url();
633 assert_eq!(url, "https://api.com");
634 std::env::remove_var("HARNESS_LLM_TEST_URL");
635 }
636
637 #[test]
638 fn test_llm_model_default() {
639 std::env::remove_var("HARNESS_LLM_TEST_MODEL");
640 assert_eq!(llm_model(), "deepseek");
641 }
642
643 #[test]
644 fn test_llm_model_custom() {
645 std::env::set_var("HARNESS_LLM_TEST_MODEL", "gpt-4o");
646 assert_eq!(llm_model(), "gpt-4o");
647 std::env::remove_var("HARNESS_LLM_TEST_MODEL");
648 }
649
650 #[test]
651 fn test_extract_usage_partial() {
652 let json = serde_json::json!({
653 "usage": {
654 "prompt_tokens": 50
655 }
656 });
657 let usage = extract_usage(&json);
658 assert_eq!(usage.prompt_tokens, 50);
659 assert_eq!(usage.completion_tokens, 0);
660 assert_eq!(usage.total_tokens, 0);
661 }
662
663 #[test]
664 fn test_browser_headless_default() {
665 std::env::remove_var("HARNESS_BROWSER_HEADLESS");
666 assert!(crate::browser_headless());
667 }
668}