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/// Endpoint registry and routing resolver.
26pub mod endpoints;
27/// MCP client for connecting to external MCP servers.
28pub mod mcp_client;
29/// MCP server for exposing the framework as an MCP server.
30pub mod mcp_server;
31/// Cost and token report printer.
32pub mod reporting;
33/// Step-by-step scenario executor (navigate, click, type, wait, assert).
34pub mod runner;
35/// Declarative TOML-based test scenario types.
36pub mod scenario;
37
38/// A2A agent server for accepting agent tasks.
39#[cfg(feature = "a2a-server")]
40pub mod a2a_server;
41
42/// `#[browser_test]` macros for `cargo test` integration.
43#[cfg(feature = "macros")]
44pub mod macros;
45
46use std::collections::HashMap;
47use std::time::Duration;
48
49use serde_json::Value;
50
51pub use costs::LlmResponse;
52pub use costs::LlmUsage;
53
54/// Configuration for the LLM client — bundles URL, model, auth, timeouts,
55/// and provider-specific options into a single struct passed everywhere.
56#[derive(Debug, Clone)]
57pub struct LlmConfig {
58    /// OpenAI-compatible API base URL (without trailing `/v1/…`).
59    pub url: String,
60    /// Model name (e.g. `gpt-4o-mini`, `deepseek`).
61    pub model: String,
62    /// API key sent as `Authorization: Bearer <key>`.
63    pub api_key: Option<String>,
64    /// Custom headers appended to every LLM request.
65    pub headers: HashMap<String, String>,
66    /// HTTP timeout.
67    pub timeout: Duration,
68    /// Sampling temperature (0.0–1.0).
69    pub temperature: f64,
70    /// Enable extended thinking / reasoning tokens.
71    /// `None` = don't send any thinking key (provider default).
72    pub thinking: Option<bool>,
73    /// Provider-specific parameters merged into the request body
74    /// (e.g. `effort = "high"` for Anthropic).
75    pub model_params: HashMap<String, Value>,
76}
77
78impl LlmConfig {
79    /// Build a config from environment defaults, falling back to safe
80    /// values when no env vars are set.
81    #[must_use]
82    pub fn from_env() -> Self {
83        Self {
84            url: llm_base_url(),
85            model: llm_model(),
86            api_key: std::env::var("HARNESS_LLM_API_KEY").ok(),
87            headers: parse_headers_env(),
88            timeout: Duration::from_secs(60),
89            temperature: 0.0,
90            thinking: None,
91            model_params: HashMap::new(),
92        }
93    }
94}
95
96/// Parses `HARNESS_LLM_HEADERS` env var (JSON object) into a header map.
97///
98/// Exposed for use by `endpoints.rs` and tests.
99#[must_use]
100pub fn parse_headers_env() -> HashMap<String, String> {
101    let Ok(raw) = std::env::var("HARNESS_LLM_HEADERS") else {
102        return HashMap::new();
103    };
104    let Ok(json) = serde_json::from_str::<Value>(&raw) else {
105        return HashMap::new();
106    };
107    let Some(obj) = json.as_object() else {
108        return HashMap::new();
109    };
110    obj.iter()
111        .filter_map(|(k, v)| v.as_str().map(|s| (k.clone(), s.to_owned())))
112        .collect()
113}
114
115/// Returns the target base URL from `HARNESS_BROWSER_BASE_URL` env,
116/// defaulting to `http://localhost:4200`.
117#[must_use]
118pub fn base_url() -> String {
119    std::env::var("HARNESS_BROWSER_BASE_URL").unwrap_or_else(|_| "http://localhost:4200".to_owned())
120}
121
122/// Returns the LLM server base URL from `HARNESS_LLM_TEST_URL` env,
123/// defaulting to `http://localhost:8080`.
124#[must_use]
125pub fn llm_base_url() -> String {
126    std::env::var("HARNESS_LLM_TEST_URL")
127        .unwrap_or_else(|_| "http://localhost:8080".to_owned())
128        .trim_end_matches('/')
129        .to_owned()
130}
131
132/// Returns the LLM model name from `HARNESS_LLM_TEST_MODEL` env,
133/// defaulting to `deepseek`.
134#[must_use]
135pub fn llm_model() -> String {
136    std::env::var("HARNESS_LLM_TEST_MODEL").unwrap_or_else(|_| "deepseek".to_owned())
137}
138
139/// Returns whether to run the browser in headless mode from
140/// `HARNESS_BROWSER_HEADLESS` env, defaulting to `true`.
141#[must_use]
142pub fn browser_headless() -> bool {
143    std::env::var("HARNESS_BROWSER_HEADLESS")
144        .map_or(true, |v| v != "0" && v.to_lowercase() != "false")
145}
146
147/// Builds a `reqwest::Client` with the given timeout.
148#[must_use]
149pub fn http_client(timeout: Duration) -> reqwest::Client {
150    reqwest::Client::builder()
151        .timeout(timeout)
152        .build()
153        .expect("build reqwest client")
154}
155
156/// Sends a chat completion request to the LLM.
157///
158/// Returns `Some(content)` on success, `None` on any error.
159///
160/// Prefer `llm_chat_with_usage` if you need token counting.
161#[must_use]
162pub async fn llm_chat(llm: &LlmConfig, system: &str, user: &str) -> Option<String> {
163    llm_chat_with_usage(llm, system, user)
164        .await
165        .map(|r| r.content)
166        .ok()
167}
168
169/// Sends a chat completion request to the LLM and returns both the content
170/// and token usage from the API response.
171///
172/// Retries transient failures (network errors, HTTP 429/5xx, invalid
173/// responses) with a short backoff, and returns the last underlying error
174/// instead of collapsing everything into a generic "server down" message.
175///
176/// # Errors
177///
178/// Returns the last underlying error as a human-readable string when every
179/// attempt fails (transport error, non-success HTTP status, response that is
180/// not valid JSON, or a response missing `choices[0].message.content`).
181pub async fn llm_chat_with_usage(
182    llm: &LlmConfig,
183    system: &str,
184    user: &str,
185) -> Result<LlmResponse, String> {
186    let client = http_client(llm.timeout);
187    let mut last_err = String::from("LLM call failed");
188
189    for attempt in 0..3u32 {
190        match llm_chat_once(&client, llm, system, user).await {
191            Ok(resp) => return Ok(resp),
192            Err(err) => {
193                last_err = err;
194                if attempt < 2 {
195                    tokio::time::sleep(Duration::from_millis(500 * u64::from(attempt + 1))).await;
196                }
197            }
198        }
199    }
200
201    Err(last_err)
202}
203
204/// Single LLM chat request attempt; returns the underlying error as text.
205async fn llm_chat_once(
206    client: &reqwest::Client,
207    llm: &LlmConfig,
208    system: &str,
209    user: &str,
210) -> Result<LlmResponse, String> {
211    let mut payload = serde_json::json!({
212        "model": llm.model,
213        "messages": [
214            {"role": "system", "content": system},
215            {"role": "user", "content": user}
216        ],
217        "max_tokens": 4096,
218        "temperature": llm.temperature
219    });
220    if let Some(think) = llm.thinking {
221        if think {
222            payload["thinking"] = serde_json::json!({"type": "enabled"});
223        } else {
224            payload["thinking"] = serde_json::json!({"type": "disabled"});
225        }
226    }
227    // Merge provider-specific parameters into the request body.
228    if !llm.model_params.is_empty() {
229        if let Value::Object(ref mut map) = payload {
230            for (key, val) in &llm.model_params {
231                map.insert(key.clone(), val.clone());
232            }
233        }
234    }
235
236    let mut req = client
237        .post(format!("{}/v1/chat/completions", llm.url))
238        .header("Content-Type", "application/json");
239
240    if let Some(ref key) = llm.api_key {
241        req = req.header("Authorization", format!("Bearer {key}"));
242    }
243    for (name, value) in &llm.headers {
244        req = req.header(name.as_str(), value.as_str());
245    }
246
247    let resp = req
248        .json(&payload)
249        .send()
250        .await
251        .map_err(|e| format!("LLM HTTP request failed: {e}"))?;
252    let status = resp.status();
253    if !status.is_success() {
254        let body = resp.text().await.unwrap_or_default();
255        return Err(format!("LLM endpoint returned HTTP {status}: {body}"));
256    }
257    let json: Value = resp
258        .json()
259        .await
260        .map_err(|e| format!("LLM response not valid JSON: {e}"))?;
261    let usage = costs::extract_usage(&json);
262    let content = json["choices"][0]["message"]["content"]
263        .as_str()
264        .map(String::from)
265        .ok_or_else(|| format!("LLM response missing choices[0].message.content: {json}"))?;
266
267    Ok(LlmResponse { content, usage })
268}
269
270/// JavaScript to extract interactive elements from the current page.
271/// Returns a JSON array of objects with tag, selector, and label.
272pub const DOM_EXTRACT_JS: &str = r#"
273(() => {
274  const interactive = 'a, button, input, textarea, select, [role="button"], [onclick], [tabindex], [data-testid], [aria-label]';
275  const els = document.querySelectorAll(interactive);
276  const info = [];
277  const seen = new Set();
278  els.forEach((el, i) => {
279    const rect = el.getBoundingClientRect();
280    if (rect.width === 0 || rect.height === 0) return;
281    const tag = el.tagName.toLowerCase();
282    let selector = '';
283    if (el.id) selector = '#' + CSS.escape(el.id);
284    else if (el.getAttribute('data-testid')) selector = '[data-testid="' + el.getAttribute('data-testid') + '"]';
285    else if (el.name) selector = '[name="' + CSS.escape(el.name) + '"]';
286    else if (el.className && typeof el.className === 'string') {
287      const cls = el.className.trim().split(/\\s+/)[0];
288      if (cls) selector = tag + '.' + CSS.escape(cls);
289    }
290    if (!selector) selector = tag;
291    if (seen.has(selector)) return;
292    seen.add(selector);
293
294    let label = '';
295    const aria = el.getAttribute('aria-label');
296    if (aria) {
297      label = aria;
298    } else if (tag === 'input' || tag === 'textarea' || tag === 'select') {
299      label = el.placeholder || el.name || el.getAttribute('aria-label') || '';
300      if (el.type && !label) label = el.type;
301    } else {
302      label = (el.textContent || '').trim().substring(0, 80);
303    }
304
305    info.push(i + ': ' + selector + ' [' + tag + '] "' + label + '"');
306  });
307  return JSON.stringify(info);
308})()
309"#;
310
311/// Truncates a string to the given maximum length, appending a marker
312/// if truncation occurred.
313#[must_use]
314pub fn truncate(s: &str, max_len: usize) -> String {
315    if s.len() <= max_len {
316        s.to_owned()
317    } else {
318        format!("{}...<truncated>", &s[..max_len])
319    }
320}
321
322#[cfg(test)]
323mod tests {
324    use crate::costs::extract_usage;
325    use crate::truncate;
326    use crate::{llm_base_url, llm_model, parse_headers_env, LlmConfig};
327
328    #[test]
329    fn test_truncate_short() {
330        assert_eq!(truncate("hello", 10), "hello");
331    }
332
333    #[test]
334    fn test_truncate_long() {
335        let result = truncate("hello world", 5);
336        assert!(result.contains("<truncated>"));
337        assert!(result.starts_with("hello"));
338    }
339
340    #[test]
341    fn test_truncate_exact_length() {
342        assert_eq!(truncate("abcde", 5), "abcde");
343    }
344
345    #[test]
346    fn test_truncate_empty() {
347        assert_eq!(truncate("", 5), "");
348    }
349
350    #[test]
351    fn test_parse_headers_env_empty() {
352        std::env::remove_var("HARNESS_LLM_HEADERS");
353        let h = parse_headers_env();
354        assert!(h.is_empty());
355    }
356
357    #[test]
358    fn test_parse_headers_env_valid() {
359        std::env::set_var("HARNESS_LLM_HEADERS", r#"{"X-Org":"acme","X-Version":"1"}"#);
360        let h = parse_headers_env();
361        assert_eq!(h.get("X-Org").map(String::as_str), Some("acme"));
362        assert_eq!(h.get("X-Version").map(String::as_str), Some("1"));
363        std::env::remove_var("HARNESS_LLM_HEADERS");
364    }
365
366    #[test]
367    fn test_parse_headers_env_invalid_json() {
368        std::env::set_var("HARNESS_LLM_HEADERS", "not-json");
369        let h = parse_headers_env();
370        assert!(h.is_empty());
371        std::env::remove_var("HARNESS_LLM_HEADERS");
372    }
373
374    #[test]
375    fn test_llm_config_from_env_defaults() {
376        #[allow(clippy::float_cmp)]
377        {
378            let config = LlmConfig::from_env();
379            assert_eq!(config.temperature, 0.0);
380            assert!(config.thinking.is_none());
381            assert!(config.model_params.is_empty());
382        }
383    }
384
385    #[test]
386    fn test_extract_usage_full() {
387        let json = serde_json::json!({
388            "usage": {
389                "prompt_tokens": 100,
390                "completion_tokens": 200,
391                "total_tokens": 300
392            }
393        });
394        let usage = extract_usage(&json);
395        assert_eq!(usage.prompt_tokens, 100);
396        assert_eq!(usage.completion_tokens, 200);
397        assert_eq!(usage.total_tokens, 300);
398    }
399
400    #[test]
401    fn test_extract_usage_empty() {
402        let json = serde_json::json!({});
403        let usage = extract_usage(&json);
404        assert_eq!(usage.prompt_tokens, 0);
405        assert_eq!(usage.completion_tokens, 0);
406        assert_eq!(usage.total_tokens, 0);
407    }
408
409    #[test]
410    fn test_truncate_unicode() {
411        // truncate uses byte-level slicing, so max_len refers to byte count
412        // 'é' is 2 bytes, so index 3 captures "hé" (h=0, é=bytes 1-2)
413        assert_eq!(truncate("héllo", 3), "hé...<truncated>");
414        // Length 5 captures full string (5 bytes)
415        assert_eq!(truncate("hello", 5), "hello");
416    }
417
418    #[test]
419    fn test_parse_headers_env_non_object() {
420        std::env::set_var("HARNESS_LLM_HEADERS", "[1, 2, 3]");
421        let h = parse_headers_env();
422        assert!(h.is_empty());
423        std::env::remove_var("HARNESS_LLM_HEADERS");
424    }
425
426    #[test]
427    fn test_parse_headers_env_nested_values_filtered() {
428        std::env::set_var(
429            "HARNESS_LLM_HEADERS",
430            r#"{"str":"val","num":42,"bool":true}"#,
431        );
432        let h = parse_headers_env();
433        assert_eq!(h.get("str").map(String::as_str), Some("val"));
434        assert!(!h.contains_key("num"));
435        assert!(!h.contains_key("bool"));
436        std::env::remove_var("HARNESS_LLM_HEADERS");
437    }
438
439    #[test]
440    fn test_llm_config_has_default_model() {
441        let config = LlmConfig::from_env();
442        assert!(!config.model.is_empty());
443    }
444
445    #[test]
446    fn test_llm_base_url_default() {
447        std::env::remove_var("HARNESS_LLM_TEST_URL");
448        let url = llm_base_url();
449        assert_eq!(url, "http://localhost:8080");
450    }
451
452    #[test]
453    fn test_llm_base_url_custom() {
454        std::env::set_var("HARNESS_LLM_TEST_URL", "https://custom.api.com/v1");
455        let url = llm_base_url();
456        assert_eq!(url, "https://custom.api.com/v1");
457        std::env::remove_var("HARNESS_LLM_TEST_URL");
458    }
459
460    #[test]
461    fn test_llm_base_url_trailing_slash() {
462        std::env::set_var("HARNESS_LLM_TEST_URL", "https://api.com/");
463        let url = llm_base_url();
464        assert_eq!(url, "https://api.com");
465        std::env::remove_var("HARNESS_LLM_TEST_URL");
466    }
467
468    #[test]
469    fn test_llm_model_default() {
470        std::env::remove_var("HARNESS_LLM_TEST_MODEL");
471        assert_eq!(llm_model(), "deepseek");
472    }
473
474    #[test]
475    fn test_llm_model_custom() {
476        std::env::set_var("HARNESS_LLM_TEST_MODEL", "gpt-4o");
477        assert_eq!(llm_model(), "gpt-4o");
478        std::env::remove_var("HARNESS_LLM_TEST_MODEL");
479    }
480
481    #[test]
482    fn test_extract_usage_partial() {
483        let json = serde_json::json!({
484            "usage": {
485                "prompt_tokens": 50
486            }
487        });
488        let usage = extract_usage(&json);
489        assert_eq!(usage.prompt_tokens, 50);
490        assert_eq!(usage.completion_tokens, 0);
491        assert_eq!(usage.total_tokens, 0);
492    }
493
494    #[test]
495    fn test_browser_headless_default() {
496        std::env::remove_var("HARNESS_BROWSER_HEADLESS");
497        assert!(crate::browser_headless());
498    }
499}