custom headers via --headers (#30)

* add custom headers via --headers

* add tests

* better parsing

cli/src/flags.rs

@@ -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"));