Skip to main content

llm_browser_testkit/
lib.rs

1//! LLM-driven browser test framework.
2//!
3//! Provides reusable building blocks for browser-based test scenarios:
4//! - Browser client via Chrome `DevTools` Protocol (headless)
5//! - LLM client for natural language element targeting and assertions
6//! - A2A agent integration for agent-to-agent communication
7//! - MCP client/server integration for tool-calling
8//! - Cost tracking, token counting, and budget enforcement
9//! - Declarative TOML scenario runner
10//! - `#[browser_test]` macros for `cargo test` integration
11
12#![allow(
13    clippy::expect_used,
14    clippy::unwrap_used,
15    clippy::panic,
16    clippy::missing_panics_doc
17)]
18
19/// A2A agent protocol client.
20pub mod a2a;
21/// Budget tracking and enforcement.
22pub mod budgets;
23/// Cost calculation, usage tracking, and pricing.
24pub mod costs;
25/// Failure diagnostics — page-state capture and artifact writing.
26pub mod diagnostics;
27/// Endpoint registry and routing resolver.
28pub mod endpoints;
29/// MCP client for connecting to external MCP servers.
30pub mod mcp_client;
31/// MCP server for exposing the framework as an MCP server.
32pub mod mcp_server;
33/// Cost and token report printer.
34pub mod reporting;
35/// Step-by-step scenario executor (navigate, click, type, wait, assert).
36pub mod runner;
37/// Declarative TOML-based test scenario types.
38pub mod scenario;
39/// CSS selector sanitization for LLM-generated selectors.
40pub mod selectors;
41/// Vision support — screenshot capture/downscale/encode for visual asserts.
42pub mod vision;
43
44/// A2A agent server for accepting agent tasks.
45#[cfg(feature = "a2a-server")]
46pub mod a2a_server;
47
48/// `#[browser_test]` macros for `cargo test` integration.
49#[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/// Configuration for the LLM client — bundles URL, model, auth, timeouts,
61/// and provider-specific options into a single struct passed everywhere.
62#[derive(Debug, Clone)]
63pub struct LlmConfig {
64    /// OpenAI-compatible API base URL (without trailing `/v1/…`).
65    pub url: String,
66    /// Model name (e.g. `gpt-4o-mini`, `deepseek`).
67    pub model: String,
68    /// API key sent as `Authorization: Bearer <key>`.
69    pub api_key: Option<String>,
70    /// Custom headers appended to every LLM request.
71    pub headers: HashMap<String, String>,
72    /// HTTP timeout.
73    pub timeout: Duration,
74    /// Sampling temperature (0.0–1.0).
75    pub temperature: f64,
76    /// Enable extended thinking / reasoning tokens.
77    /// `None` = don't send any thinking key (provider default).
78    pub thinking: Option<bool>,
79    /// Provider-specific parameters merged into the request body
80    /// (e.g. `effort = "high"` for Anthropic).
81    pub model_params: HashMap<String, Value>,
82}
83
84impl LlmConfig {
85    /// Build a config from environment defaults, falling back to safe
86    /// values when no env vars are set.
87    #[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/// Parses `HARNESS_LLM_HEADERS` env var (JSON object) into a header map.
103///
104/// Exposed for use by `endpoints.rs` and tests.
105#[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/// Returns the target base URL from `HARNESS_BROWSER_BASE_URL` env,
122/// defaulting to `http://localhost:4200`.
123#[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/// Returns the LLM server base URL from `HARNESS_LLM_TEST_URL` env,
129/// defaulting to `http://localhost:8080`.
130#[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/// Returns the LLM model name from `HARNESS_LLM_TEST_MODEL` env,
139/// defaulting to `deepseek`.
140#[must_use]
141pub fn llm_model() -> String {
142    std::env::var("HARNESS_LLM_TEST_MODEL").unwrap_or_else(|_| "deepseek".to_owned())
143}
144
145/// Returns whether to run the browser in headless mode from
146/// `HARNESS_BROWSER_HEADLESS` env, defaulting to `true`.
147#[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/// Builds a `reqwest::Client` with the given timeout.
154#[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/// Sends a chat completion request to the LLM.
163///
164/// Returns `Some(content)` on success, `None` on any error.
165///
166/// Prefer `llm_chat_with_usage` if you need token counting.
167#[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
175/// Sends a chat completion request to the LLM and returns both the content
176/// and token usage from the API response.
177///
178/// Retries transient failures (network errors, HTTP 429/5xx, invalid
179/// responses) with a short backoff, and returns the last underlying error
180/// instead of collapsing everything into a generic "server down" message.
181/// The error text includes the HTTP status and a truncated response-body
182/// snippet, so a gateway that answers with an HTML error page is
183/// identifiable in CI logs instead of surfacing as a bare JSON decode
184/// error.
185///
186/// # Errors
187///
188/// Returns the last underlying error as a human-readable string when every
189/// attempt fails (transport error, non-success HTTP status, response that is
190/// not valid JSON, or a response missing `choices[0].message.content`).
191pub 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
199/// Sends a vision-enabled chat completion request: the user message carries
200/// both the text prompt and a screenshot (JPEG/PNG data URL) as an
201/// OpenAI-compatible `image_url` content part.
202///
203/// Retries and error reporting behave like [`llm_chat_with_usage`].
204///
205/// # Errors
206///
207/// Returns the last underlying error as a human-readable string when every
208/// attempt fails (transport error, non-success HTTP status, response that is
209/// not valid JSON, or a response missing `choices[0].message.content`).
210pub 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
219/// Shared retry loop for text-only and vision chat completions.
220async 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
250/// Number of attempts for a single chat completion call.
251const LLM_CALL_ATTEMPTS: u32 = 3;
252
253/// Builds the chat messages array. Text-only messages keep the plain
254/// string `content` shape (maximum provider compatibility); vision calls
255/// use the OpenAI-compatible content-part array with a `data:` image URL.
256#[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
273/// Internal error type for a single LLM request attempt, distinguishing
274/// transient failures (worth retrying) from deterministic configuration
275/// errors (fail immediately).
276enum LlmCallError {
277    /// Transport-level failure (connect, timeout, TLS, …).
278    Transport { message: String },
279    /// Non-success HTTP status with a body snippet.
280    Http { status: u16, body: String },
281    /// Success status but the body is not valid JSON.
282    InvalidJson {
283        status: u16,
284        detail: String,
285        body: String,
286    },
287    /// Valid JSON but missing `choices[0].message.content`.
288    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    /// Whether another attempt may succeed. Network errors, rate limits,
322    /// server errors, and 200-with-garbage responses can be transient on
323    /// flaky gateways; auth/not-found errors are deterministic.
324    #[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
338/// Single LLM chat request attempt; returns the underlying error as text.
339async 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    // Merge provider-specific parameters into the request body.
360    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
416/// JavaScript to extract interactive elements from the current page.
417/// Returns a JSON array of objects with tag, selector, and label.
418pub 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/// Truncates a string to the given maximum length, appending a marker
458/// if truncation occurred.
459#[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        // truncate uses byte-level slicing, so max_len refers to byte count
558        // 'é' is 2 bytes, so index 3 captures "hé" (h=0, é=bytes 1-2)
559        assert_eq!(truncate("héllo", 3), "hé...<truncated>");
560        // Length 5 captures full string (5 bytes)
561        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}