custom headers via --headers (#30)

* add custom headers via --headers * add tests * better parsing
2026-01-12 12:01:19 -06:00
parent 4f6fd8ec5c
commit 1a88d7f585
8 changed files with 347 additions and 7 deletions
@@ -80,7 +80,14 @@ pub fn parse_command(args: &[String], flags: &Flags) -> Result<Value, ParseError
            } else {
                format!("https://{}", url)
            };
-            Ok(json!({ "id": id, "action": "navigate", "url": url }))
+            let mut nav_cmd = json!({ "id": id, "action": "navigate", "url": url });
+            // If --headers flag is set, include headers (scoped to this origin)
+            if let Some(ref headers_json) = flags.headers {
+                if let Ok(headers) = serde_json::from_str::<serde_json::Value>(headers_json) {
+                    nav_cmd["headers"] = headers;
+                }
+            }
+            Ok(nav_cmd)
        }
        "back" => Ok(json!({ "id": id, "action": "back" })),
        "forward" => Ok(json!({ "id": id, "action": "forward" })),
@@ -766,7 +773,13 @@ fn parse_set(rest: &[&str], id: &str) -> Result<Value, ParseError> {
                context: "set headers".to_string(),
                usage: "set headers <json>",
            })?;
-            Ok(json!({ "id": id, "action": "headers", "headers": headers_json }))
+            // Parse the JSON string into an object
+            let headers: serde_json::Value = serde_json::from_str(headers_json)
+                .map_err(|_| ParseError::MissingArguments {
+                    context: "set headers".to_string(),
+                    usage: "set headers <json> (must be valid JSON object)",
+                })?;
+            Ok(json!({ "id": id, "action": "headers", "headers": headers }))
        }
        Some("credentials") | Some("auth") => {
            let user = rest.get(1).ok_or_else(|| ParseError::MissingArguments {
@@ -886,6 +899,7 @@ mod tests {
            full: false,
            headed: false,
            debug: false,
+            headers: None,
            executable_path: None,
        }
    }
@@ -1013,6 +1027,81 @@ mod tests {
        assert_eq!(cmd["url"], "https://example.com");
    }

+    #[test]
+    fn test_navigate_with_headers() {
+        let mut flags = default_flags();
+        flags.headers = Some(r#"{"Authorization": "Bearer token"}"#.to_string());
+        let cmd = parse_command(&args("open api.example.com"), &flags).unwrap();
+        assert_eq!(cmd["action"], "navigate");
+        assert_eq!(cmd["url"], "https://api.example.com");
+        assert_eq!(cmd["headers"]["Authorization"], "Bearer token");
+    }
+
+    #[test]
+    fn test_navigate_with_multiple_headers() {
+        let mut flags = default_flags();
+        flags.headers = Some(r#"{"Authorization": "Bearer token", "X-Custom": "value"}"#.to_string());
+        let cmd = parse_command(&args("open api.example.com"), &flags).unwrap();
+        assert_eq!(cmd["headers"]["Authorization"], "Bearer token");
+        assert_eq!(cmd["headers"]["X-Custom"], "value");
+    }
+
+    #[test]
+    fn test_navigate_without_headers_flag() {
+        let cmd = parse_command(&args("open example.com"), &default_flags()).unwrap();
+        assert_eq!(cmd["action"], "navigate");
+        // headers should not be present when flag is not set
+        assert!(cmd.get("headers").is_none());
+    }
+
+    #[test]
+    fn test_navigate_with_invalid_headers_json() {
+        let mut flags = default_flags();
+        flags.headers = Some("not valid json".to_string());
+        let cmd = parse_command(&args("open api.example.com"), &flags).unwrap();
+        // Invalid JSON should result in no headers field (graceful handling)
+        assert!(cmd.get("headers").is_none());
+    }
+
+    // === Set Headers Tests ===
+
+    #[test]
+    fn test_set_headers_parses_json() {
+        let input: Vec<String> = vec![
+            "set".to_string(),
+            "headers".to_string(),
+            r#"{"Authorization":"Bearer token"}"#.to_string(),
+        ];
+        let cmd = parse_command(&input, &default_flags()).unwrap();
+        assert_eq!(cmd["action"], "headers");
+        // Headers should be an object, not a string
+        assert!(cmd["headers"].is_object());
+        assert_eq!(cmd["headers"]["Authorization"], "Bearer token");
+    }
+
+    #[test]
+    fn test_set_headers_with_multiple_values() {
+        let input: Vec<String> = vec![
+            "set".to_string(),
+            "headers".to_string(),
+            r#"{"Authorization": "Bearer token", "X-Custom": "value"}"#.to_string(),
+        ];
+        let cmd = parse_command(&input, &default_flags()).unwrap();
+        assert_eq!(cmd["headers"]["Authorization"], "Bearer token");
+        assert_eq!(cmd["headers"]["X-Custom"], "value");
+    }
+
+    #[test]
+    fn test_set_headers_invalid_json_error() {
+        let input: Vec<String> = vec![
+            "set".to_string(),
+            "headers".to_string(),
+            "not-valid-json".to_string(),
+        ];
+        let result = parse_command(&input, &default_flags());
+        assert!(result.is_err());
+    }
+
    #[test]
    fn test_back() {
        let cmd = parse_command(&args("back"), &default_flags()).unwrap();
@@ -6,6 +6,7 @@ pub struct Flags {
    pub headed: bool,
    pub debug: bool,
    pub session: String,
+    pub headers: Option<String>,
    pub executable_path: Option<String>,
 }

@@ -16,6 +17,7 @@ pub fn parse_flags(args: &[String]) -> Flags {
        headed: false,
        debug: false,
        session: env::var("AGENT_BROWSER_SESSION").unwrap_or_else(|_| "default".to_string()),
+        headers: None,
        executable_path: env::var("AGENT_BROWSER_EXECUTABLE_PATH").ok(),
    };

@@ -32,6 +34,12 @@ pub fn parse_flags(args: &[String]) -> Flags {
                    i += 1;
                }
            }
+            "--headers" => {
+                if let Some(h) = args.get(i + 1) {
+                    flags.headers = Some(h.clone());
+                    i += 1;
+                }
+            }
            "--executable-path" => {
                if let Some(s) = args.get(i + 1) {
                    flags.executable_path = Some(s.clone());
@@ -51,15 +59,15 @@ pub fn clean_args(args: &[String]) -> Vec<String> {

    // Global flags that should be stripped from command args
    const GLOBAL_FLAGS: &[&str] = &["--json", "--full", "--headed", "--debug"];
-    // Flags that take a value (skip both the flag and the next arg)
-    const VALUE_FLAGS: &[&str] = &["--session", "--executable-path"];
+    // Global flags that take a value (need to skip the next arg too)
+    const GLOBAL_FLAGS_WITH_VALUE: &[&str] = &["--session", "--headers", "--executable-path"];

    for arg in args.iter() {
        if skip_next {
            skip_next = false;
            continue;
        }
-        if VALUE_FLAGS.contains(&arg.as_str()) {
+        if GLOBAL_FLAGS_WITH_VALUE.contains(&arg.as_str()) {
            skip_next = true;
            continue;
        }
@@ -80,6 +88,74 @@ mod tests {
        s.split_whitespace().map(String::from).collect()
    }

+    #[test]
+    fn test_parse_headers_flag() {
+        let flags = parse_flags(&args(r#"open example.com --headers {"Auth":"token"}"#));
+        assert_eq!(flags.headers, Some(r#"{"Auth":"token"}"#.to_string()));
+    }
+
+    #[test]
+    fn test_parse_headers_flag_with_spaces() {
+        // Headers JSON is passed as a single quoted argument in shell
+        let input: Vec<String> = vec![
+            "open".to_string(),
+            "example.com".to_string(),
+            "--headers".to_string(),
+            r#"{"Authorization": "Bearer token"}"#.to_string(),
+        ];
+        let flags = parse_flags(&input);
+        assert_eq!(flags.headers, Some(r#"{"Authorization": "Bearer token"}"#.to_string()));
+    }
+
+    #[test]
+    fn test_parse_no_headers_flag() {
+        let flags = parse_flags(&args("open example.com"));
+        assert!(flags.headers.is_none());
+    }
+
+    #[test]
+    fn test_clean_args_removes_headers() {
+        let input: Vec<String> = vec![
+            "open".to_string(),
+            "example.com".to_string(),
+            "--headers".to_string(),
+            r#"{"Auth":"token"}"#.to_string(),
+        ];
+        let clean = clean_args(&input);
+        assert_eq!(clean, vec!["open", "example.com"]);
+    }
+
+    #[test]
+    fn test_clean_args_removes_headers_at_start() {
+        let input: Vec<String> = vec![
+            "--headers".to_string(),
+            r#"{"Auth":"token"}"#.to_string(),
+            "open".to_string(),
+            "example.com".to_string(),
+        ];
+        let clean = clean_args(&input);
+        assert_eq!(clean, vec!["open", "example.com"]);
+    }
+
+    #[test]
+    fn test_headers_with_other_flags() {
+        let input: Vec<String> = vec![
+            "open".to_string(),
+            "example.com".to_string(),
+            "--headers".to_string(),
+            r#"{"Auth":"token"}"#.to_string(),
+            "--json".to_string(),
+            "--headed".to_string(),
+        ];
+        let flags = parse_flags(&input);
+        assert_eq!(flags.headers, Some(r#"{"Auth":"token"}"#.to_string()));
+        assert!(flags.json);
+        assert!(flags.headed);
+        
+        let clean = clean_args(&input);
+        assert_eq!(clean, vec!["open", "example.com"]);
+    }
+
    #[test]
    fn test_parse_executable_path_flag() {
        let flags = parse_flags(&args("--executable-path /path/to/chromium open example.com"));
@@ -162,12 +162,15 @@ Aliases: goto, navigate
 Global Options:
  --json               Output as JSON
  --session <name>     Use specific session
+  --headers <json>     Set HTTP headers (scoped to this origin)
  --headed             Show browser window

 Examples:
  agent-browser open example.com
  agent-browser open https://github.com
  agent-browser open localhost:3000
+  agent-browser open api.example.com --headers '{"Authorization": "Bearer token"}'
+    # ^ Headers only sent to api.example.com, not other domains
 "##,
        "back" => r##"
 agent-browser back - Navigate back in history
@@ -1186,6 +1189,7 @@ Snapshot Options:

 Options:
  --session <name>           Isolated session (or AGENT_BROWSER_SESSION env)
+  --headers <json>           HTTP headers scoped to URL's origin (for auth)
  --executable-path <path>   Custom browser executable (or AGENT_BROWSER_EXECUTABLE_PATH)
  --json                     JSON output
  --full, -f                 Full page screenshot