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(
192 llm: &LlmConfig,
193 system: &str,
194 user: &str,
195) -> Result<LlmResponse, String> {
196 chat_with_retry(llm, system, user, None).await
197}
198
199pub async fn llm_chat_vision_with_usage(
211 llm: &LlmConfig,
212 system: &str,
213 user: &str,
214 image_data_url: &str,
215) -> Result<LlmResponse, String> {
216 chat_with_retry(llm, system, user, Some(image_data_url)).await
217}
218
219async fn chat_with_retry(
221 llm: &LlmConfig,
222 system: &str,
223 user: &str,
224 image_data_url: Option<&str>,
225) -> Result<LlmResponse, String> {
226 let client = http_client(llm.timeout);
227 let mut last_err = String::from("LLM call failed");
228 let mut attempts: u32 = 0;
229
230 while attempts < LLM_CALL_ATTEMPTS {
231 attempts += 1;
232 match llm_chat_once(&client, llm, system, user, image_data_url).await {
233 Ok(resp) => return Ok(resp),
234 Err(err) => {
235 last_err = err.to_string();
236 if attempts >= LLM_CALL_ATTEMPTS || !err.is_retryable() {
237 break;
238 }
239 tokio::time::sleep(Duration::from_millis(500 * u64::from(attempts))).await;
240 }
241 }
242 }
243
244 Err(format!(
245 "LLM call failed after {attempts} attempt(s) (endpoint {url}): {last_err}",
246 url = llm.url
247 ))
248}
249
250const LLM_CALL_ATTEMPTS: u32 = 3;
252
253#[must_use]
257fn build_messages(system: &str, user: &str, image_data_url: Option<&str>) -> Value {
258 let user_content = image_data_url.map_or_else(
259 || Value::String(user.to_owned()),
260 |url| {
261 json!([
262 {"type": "text", "text": user},
263 {"type": "image_url", "image_url": {"url": url}}
264 ])
265 },
266 );
267 json!([
268 {"role": "system", "content": system},
269 {"role": "user", "content": user_content}
270 ])
271}
272
273enum LlmCallError {
277 Transport { message: String },
279 Http { status: u16, body: String },
281 InvalidJson {
283 status: u16,
284 detail: String,
285 body: String,
286 },
287 MissingContent { json: String },
289}
290
291impl std::fmt::Display for LlmCallError {
292 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
293 match self {
294 Self::Transport { message } => write!(f, "LLM HTTP request failed: {message}"),
295 Self::Http { status, body } => {
296 write!(
297 f,
298 "LLM endpoint returned HTTP {status}: {}",
299 truncate(body, 300)
300 )
301 }
302 Self::InvalidJson {
303 status,
304 detail,
305 body,
306 } => write!(
307 f,
308 "LLM endpoint returned HTTP {status} with non-JSON body ({detail}): {}",
309 truncate(body, 300)
310 ),
311 Self::MissingContent { json } => write!(
312 f,
313 "LLM response missing choices[0].message.content: {}",
314 truncate(json, 300)
315 ),
316 }
317 }
318}
319
320impl LlmCallError {
321 #[must_use]
325 fn is_retryable(&self) -> bool {
326 match self {
327 Self::Transport { .. } | Self::MissingContent { .. } => true,
328 Self::Http { status, .. } => {
329 *status == 408 || *status == 429 || (500..600).contains(status)
330 }
331 Self::InvalidJson { status, .. } => {
332 *status == 200 || *status == 408 || *status == 429 || (500..600).contains(status)
333 }
334 }
335 }
336}
337
338async fn llm_chat_once(
340 client: &reqwest::Client,
341 llm: &LlmConfig,
342 system: &str,
343 user: &str,
344 image_data_url: Option<&str>,
345) -> Result<LlmResponse, LlmCallError> {
346 let mut payload = serde_json::json!({
347 "model": llm.model,
348 "messages": build_messages(system, user, image_data_url),
349 "max_tokens": 4096,
350 "temperature": llm.temperature
351 });
352 if let Some(think) = llm.thinking {
353 if think {
354 payload["thinking"] = serde_json::json!({"type": "enabled"});
355 } else {
356 payload["thinking"] = serde_json::json!({"type": "disabled"});
357 }
358 }
359 if !llm.model_params.is_empty() {
361 if let Value::Object(ref mut map) = payload {
362 for (key, val) in &llm.model_params {
363 map.insert(key.clone(), val.clone());
364 }
365 }
366 }
367
368 let mut req = client
369 .post(format!("{}/v1/chat/completions", llm.url))
370 .header("Content-Type", "application/json");
371
372 if let Some(ref key) = llm.api_key {
373 req = req.header("Authorization", format!("Bearer {key}"));
374 }
375 for (name, value) in &llm.headers {
376 req = req.header(name.as_str(), value.as_str());
377 }
378
379 let resp = req
380 .json(&payload)
381 .send()
382 .await
383 .map_err(|e| LlmCallError::Transport {
384 message: e.to_string(),
385 })?;
386 let status = resp.status();
387 let status_u16 = status.as_u16();
388 let body = resp.text().await.unwrap_or_default();
389 if !status.is_success() {
390 return Err(LlmCallError::Http {
391 status: status_u16,
392 body,
393 });
394 }
395 let json: Value = match serde_json::from_str(&body) {
396 Ok(v) => v,
397 Err(e) => {
398 return Err(LlmCallError::InvalidJson {
399 status: status_u16,
400 detail: e.to_string(),
401 body,
402 });
403 }
404 };
405 let usage = costs::extract_usage(&json);
406 let content = json["choices"][0]["message"]["content"]
407 .as_str()
408 .map(String::from)
409 .ok_or_else(|| LlmCallError::MissingContent {
410 json: json.to_string(),
411 })?;
412
413 Ok(LlmResponse { content, usage })
414}
415
416pub const DOM_EXTRACT_JS: &str = r#"
419(() => {
420 const interactive = 'a, button, input, textarea, select, [role="button"], [onclick], [tabindex], [data-testid], [aria-label]';
421 const els = document.querySelectorAll(interactive);
422 const info = [];
423 const seen = new Set();
424 els.forEach((el, i) => {
425 const rect = el.getBoundingClientRect();
426 if (rect.width === 0 || rect.height === 0) return;
427 const tag = el.tagName.toLowerCase();
428 let selector = '';
429 if (el.id) selector = '#' + CSS.escape(el.id);
430 else if (el.getAttribute('data-testid')) selector = '[data-testid="' + el.getAttribute('data-testid') + '"]';
431 else if (el.name) selector = '[name="' + CSS.escape(el.name) + '"]';
432 else if (el.className && typeof el.className === 'string') {
433 const cls = el.className.trim().split(/\\s+/)[0];
434 if (cls) selector = tag + '.' + CSS.escape(cls);
435 }
436 if (!selector) selector = tag;
437 if (seen.has(selector)) return;
438 seen.add(selector);
439
440 let label = '';
441 const aria = el.getAttribute('aria-label');
442 if (aria) {
443 label = aria;
444 } else if (tag === 'input' || tag === 'textarea' || tag === 'select') {
445 label = el.placeholder || el.name || el.getAttribute('aria-label') || '';
446 if (el.type && !label) label = el.type;
447 } else {
448 label = (el.textContent || '').trim().substring(0, 80);
449 }
450
451 info.push(i + ': ' + selector + ' [' + tag + '] "' + label + '"');
452 });
453 return JSON.stringify(info);
454})()
455"#;
456
457#[must_use]
460pub fn truncate(s: &str, max_len: usize) -> String {
461 if s.len() <= max_len {
462 s.to_owned()
463 } else {
464 format!("{}...<truncated>", &s[..max_len])
465 }
466}
467
468#[cfg(test)]
469mod tests {
470 use crate::costs::extract_usage;
471 use crate::truncate;
472 use crate::{llm_base_url, llm_model, parse_headers_env, LlmConfig};
473
474 #[test]
475 fn test_truncate_short() {
476 assert_eq!(truncate("hello", 10), "hello");
477 }
478
479 #[test]
480 fn test_truncate_long() {
481 let result = truncate("hello world", 5);
482 assert!(result.contains("<truncated>"));
483 assert!(result.starts_with("hello"));
484 }
485
486 #[test]
487 fn test_truncate_exact_length() {
488 assert_eq!(truncate("abcde", 5), "abcde");
489 }
490
491 #[test]
492 fn test_truncate_empty() {
493 assert_eq!(truncate("", 5), "");
494 }
495
496 #[test]
497 fn test_parse_headers_env_empty() {
498 std::env::remove_var("HARNESS_LLM_HEADERS");
499 let h = parse_headers_env();
500 assert!(h.is_empty());
501 }
502
503 #[test]
504 fn test_parse_headers_env_valid() {
505 std::env::set_var("HARNESS_LLM_HEADERS", r#"{"X-Org":"acme","X-Version":"1"}"#);
506 let h = parse_headers_env();
507 assert_eq!(h.get("X-Org").map(String::as_str), Some("acme"));
508 assert_eq!(h.get("X-Version").map(String::as_str), Some("1"));
509 std::env::remove_var("HARNESS_LLM_HEADERS");
510 }
511
512 #[test]
513 fn test_parse_headers_env_invalid_json() {
514 std::env::set_var("HARNESS_LLM_HEADERS", "not-json");
515 let h = parse_headers_env();
516 assert!(h.is_empty());
517 std::env::remove_var("HARNESS_LLM_HEADERS");
518 }
519
520 #[test]
521 fn test_llm_config_from_env_defaults() {
522 #[allow(clippy::float_cmp)]
523 {
524 let config = LlmConfig::from_env();
525 assert_eq!(config.temperature, 0.0);
526 assert!(config.thinking.is_none());
527 assert!(config.model_params.is_empty());
528 }
529 }
530
531 #[test]
532 fn test_extract_usage_full() {
533 let json = serde_json::json!({
534 "usage": {
535 "prompt_tokens": 100,
536 "completion_tokens": 200,
537 "total_tokens": 300
538 }
539 });
540 let usage = extract_usage(&json);
541 assert_eq!(usage.prompt_tokens, 100);
542 assert_eq!(usage.completion_tokens, 200);
543 assert_eq!(usage.total_tokens, 300);
544 }
545
546 #[test]
547 fn test_extract_usage_empty() {
548 let json = serde_json::json!({});
549 let usage = extract_usage(&json);
550 assert_eq!(usage.prompt_tokens, 0);
551 assert_eq!(usage.completion_tokens, 0);
552 assert_eq!(usage.total_tokens, 0);
553 }
554
555 #[test]
556 fn test_truncate_unicode() {
557 assert_eq!(truncate("héllo", 3), "hé...<truncated>");
560 assert_eq!(truncate("hello", 5), "hello");
562 }
563
564 #[test]
565 fn test_parse_headers_env_non_object() {
566 std::env::set_var("HARNESS_LLM_HEADERS", "[1, 2, 3]");
567 let h = parse_headers_env();
568 assert!(h.is_empty());
569 std::env::remove_var("HARNESS_LLM_HEADERS");
570 }
571
572 #[test]
573 fn test_parse_headers_env_nested_values_filtered() {
574 std::env::set_var(
575 "HARNESS_LLM_HEADERS",
576 r#"{"str":"val","num":42,"bool":true}"#,
577 );
578 let h = parse_headers_env();
579 assert_eq!(h.get("str").map(String::as_str), Some("val"));
580 assert!(!h.contains_key("num"));
581 assert!(!h.contains_key("bool"));
582 std::env::remove_var("HARNESS_LLM_HEADERS");
583 }
584
585 #[test]
586 fn test_llm_config_has_default_model() {
587 let config = LlmConfig::from_env();
588 assert!(!config.model.is_empty());
589 }
590
591 #[test]
592 fn test_llm_base_url_default() {
593 std::env::remove_var("HARNESS_LLM_TEST_URL");
594 let url = llm_base_url();
595 assert_eq!(url, "http://localhost:8080");
596 }
597
598 #[test]
599 fn test_llm_base_url_custom() {
600 std::env::set_var("HARNESS_LLM_TEST_URL", "https://custom.api.com/v1");
601 let url = llm_base_url();
602 assert_eq!(url, "https://custom.api.com/v1");
603 std::env::remove_var("HARNESS_LLM_TEST_URL");
604 }
605
606 #[test]
607 fn test_llm_base_url_trailing_slash() {
608 std::env::set_var("HARNESS_LLM_TEST_URL", "https://api.com/");
609 let url = llm_base_url();
610 assert_eq!(url, "https://api.com");
611 std::env::remove_var("HARNESS_LLM_TEST_URL");
612 }
613
614 #[test]
615 fn test_llm_model_default() {
616 std::env::remove_var("HARNESS_LLM_TEST_MODEL");
617 assert_eq!(llm_model(), "deepseek");
618 }
619
620 #[test]
621 fn test_llm_model_custom() {
622 std::env::set_var("HARNESS_LLM_TEST_MODEL", "gpt-4o");
623 assert_eq!(llm_model(), "gpt-4o");
624 std::env::remove_var("HARNESS_LLM_TEST_MODEL");
625 }
626
627 #[test]
628 fn test_extract_usage_partial() {
629 let json = serde_json::json!({
630 "usage": {
631 "prompt_tokens": 50
632 }
633 });
634 let usage = extract_usage(&json);
635 assert_eq!(usage.prompt_tokens, 50);
636 assert_eq!(usage.completion_tokens, 0);
637 assert_eq!(usage.total_tokens, 0);
638 }
639
640 #[test]
641 fn test_browser_headless_default() {
642 std::env::remove_var("HARNESS_BROWSER_HEADLESS");
643 assert!(crate::browser_headless());
644 }
645}