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, and HTTP 200 + empty body — a gateway warm-up signature,
180/// retried with a 3s backoff) with a short backoff, and returns the last
181/// underlying error instead of collapsing everything into a generic
182/// "server down" message. The error text includes the HTTP status and a
183/// truncated response-body snippet, so a gateway that answers with an
184/// HTML error page is identifiable in CI logs instead of surfacing as a
185/// bare JSON decode error. Deterministic client errors (401/403/404) are
186/// not retried.
187///
188/// # Errors
189///
190/// Returns the last underlying error as a human-readable string when every
191/// attempt fails (transport error, non-success HTTP status, response that is
192/// not valid JSON, an empty HTTP 200 body, or a response missing
193/// `choices[0].message.content`).
194pub 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
202/// Sends a vision-enabled chat completion request: the user message carries
203/// both the text prompt and a screenshot (JPEG/PNG data URL) as an
204/// OpenAI-compatible `image_url` content part.
205///
206/// Retries and error reporting behave like [`llm_chat_with_usage`].
207///
208/// # Errors
209///
210/// Returns the last underlying error as a human-readable string when every
211/// attempt fails (transport error, non-success HTTP status, response that is
212/// not valid JSON, or a response missing `choices[0].message.content`).
213pub 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
222/// Shared retry loop for text-only and vision chat completions.
223async 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                // An HTTP 200 with an empty body is a gateway warm-up
239                // signature: give it a longer window to finish booting
240                // instead of hammering it with short retries.
241                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
260/// Number of attempts for a single chat completion call.
261const LLM_CALL_ATTEMPTS: u32 = 3;
262
263/// Builds the chat messages array. Text-only messages keep the plain
264/// string `content` shape (maximum provider compatibility); vision calls
265/// use the OpenAI-compatible content-part array with a `data:` image URL.
266#[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
283/// Internal error type for a single LLM request attempt, distinguishing
284/// transient failures (worth retrying) from deterministic configuration
285/// errors (fail immediately).
286enum LlmCallError {
287    /// Transport-level failure (connect, timeout, TLS, …).
288    Transport { message: String },
289    /// Non-success HTTP status with a body snippet.
290    Http { status: u16, body: String },
291    /// Success status but the body is not valid JSON.
292    InvalidJson {
293        status: u16,
294        detail: String,
295        body: String,
296    },
297    /// Success status with an EMPTY body — the classic transient gateway
298    /// warm-up signature (HTTP 200, zero bytes). Retried with a longer
299    /// backoff than other errors.
300    EmptyBody { status: u16 },
301    /// Valid JSON but missing `choices[0].message.content`.
302    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    /// Whether another attempt may succeed. Network errors, rate limits,
340    /// server errors, and 200-with-garbage responses can be transient on
341    /// flaky gateways; auth/not-found errors are deterministic.
342    #[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
356/// Single LLM chat request attempt; returns the underlying error as text.
357async 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    // Merge provider-specific parameters into the request body.
378    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        // HTTP 200 + empty body: a transient gateway hiccup (cold-start /
415        // warm-up), not a client error. Retried with a longer backoff.
416        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
439/// JavaScript to extract interactive elements from the current page.
440/// Returns a JSON array of objects with tag, selector, and label.
441pub 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/// Truncates a string to the given maximum length, appending a marker
481/// if truncation occurred.
482#[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        // truncate uses byte-level slicing, so max_len refers to byte count
581        // 'é' is 2 bytes, so index 3 captures "hé" (h=0, é=bytes 1-2)
582        assert_eq!(truncate("héllo", 3), "hé...<truncated>");
583        // Length 5 captures full string (5 bytes)
584        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}