feat: add AWS Bedrock AgentCore browser provider (native Rust) (#397)

* feat: add AWS Bedrock AgentCore browser provider (native Rust) - Add agentcore provider with SigV4 authentication - AWS SDK deps are optional behind 'agentcore' feature flag - Build with: cargo build --features agentcore - Supports AGENTCORE_REGION, AGENTCORE_PROFILE_ID, AGENTCORE_BROWSER_ID env vars - Returns session ID and Live View URL in launch response - Add connect_cdp_with_headers for signed WebSocket connections * test: add unit tests for AgentCore provider * refactor: use lightweight manual SigV4 signing instead of AWS SDK - Replace aws-sigv4/aws-config with manual HMAC-SHA256 signing - Removes ~60s compile time and significant binary size - Credentials read from AWS_ACCESS_KEY_ID/AWS_SECRET_ACCESS_KEY env vars - Supports AWS_SESSION_TOKEN for temporary credentials * fix: correct AgentCore API endpoints - Host: bedrock-agentcore.{region}.amazonaws.com - Start session: PUT /browsers/{id}/sessions/start - Stop session: PUT /browsers/{id}/sessions/stop - Add urlencoding for browser ID in path - Add AWS_DEFAULT_REGION fallback * fix: use profileConfiguration.profileIdentifier for AgentCore profile The AWS Bedrock AgentCore API expects profile configuration in the format: { "profileConfiguration": { "profileIdentifier": "<profile-id>" } } Not the flat "profileId" field that was previously used. * feat: support AWS credential provider chain via AWS CLI - Try env vars first (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY) - Fall back to 'aws configure export-credentials --format env' - Honor AWS_PROFILE environment variable - Works with SSO, IAM roles, credential files, etc. --------- Co-authored-by: Chris Tate <chris@ctate.dev>
2026-04-02 18:33:37 -05:00
parent 89595836c6
commit 8561a755ef
7 changed files with 673 additions and 16 deletions
@@ -1619,8 +1619,18 @@ async fn handle_launch(cmd: &Value, state: &mut DaemonState) -> Result<Value, St
            }
            _ => {
                let conn = providers::connect_provider(provider).await?;
+
+                let ws_headers = if provider.eq_ignore_ascii_case("agentcore") {
+                    providers::take_agentcore_ws_headers()
+                } else {
+                    None
+                };
+
                let connect_result = if conn.direct_page {
                    BrowserManager::connect_cdp_direct(&conn.ws_url).await
+                } else if ws_headers.is_some() {
+                    BrowserManager::connect_cdp_with_headers(&conn.ws_url, ws_headers)
+                        .await
                } else {
                    BrowserManager::connect_cdp(&conn.ws_url).await
                };
@@ -1633,6 +1643,17 @@ async fn handle_launch(cmd: &Value, state: &mut DaemonState) -> Result<Value, St
                        state.start_dialog_handler();
                        state.update_stream_client().await;
                        write_provider_file(&state.session_id, provider);
+
+                        #[cfg(feature = "agentcore")]
+                        if let Some(info) = providers::get_agentcore_info() {
+                            return Ok(json!({
+                                "launched": true,
+                                "provider": provider,
+                                "agentCoreSessionId": info.session_id,
+                                "agentCoreLiveViewUrl": info.live_view_url
+                            }));
+                        }
+
                        return Ok(json!({ "launched": true, "provider": provider }));
                    }
                    Err(e) => {
@@ -328,18 +328,29 @@ impl BrowserManager {
    }

    pub async fn connect_cdp(url: &str) -> Result<Self, String> {
-        Self::connect_cdp_inner(url, false).await
+        Self::connect_cdp_inner(url, false, None).await
    }

    /// Connect to a provider CDP proxy where the WebSocket IS the page session.
    /// Skips browser-level Target.* commands that most proxies don't support.
    pub async fn connect_cdp_direct(url: &str) -> Result<Self, String> {
-        Self::connect_cdp_inner(url, true).await
+        Self::connect_cdp_inner(url, true, None).await
    }

-    async fn connect_cdp_inner(url: &str, direct_page: bool) -> Result<Self, String> {
+    pub async fn connect_cdp_with_headers(
+        url: &str,
+        headers: Option<Vec<(String, String)>>,
+    ) -> Result<Self, String> {
+        Self::connect_cdp_inner(url, false, headers).await
+    }
+
+    async fn connect_cdp_inner(
+        url: &str,
+        direct_page: bool,
+        headers: Option<Vec<(String, String)>>,
+    ) -> Result<Self, String> {
        let ws_url = resolve_cdp_url(url).await?;
-        let client = Arc::new(CdpClient::connect(&ws_url).await?);
+        let client = Arc::new(CdpClient::connect_with_headers(&ws_url, headers).await?);
        let mut manager = Self {
            client,
            browser_process: None,
@@ -6,6 +6,7 @@ use std::sync::Arc;
 use futures_util::{SinkExt, StreamExt};
 use serde_json::Value;
 use tokio::sync::{broadcast, oneshot, Mutex};
+use tokio_tungstenite::tungstenite::client::IntoClientRequest;
 use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
 use tokio_tungstenite::tungstenite::Message;

@@ -46,9 +47,29 @@ pub struct CdpClient {

 impl CdpClient {
    pub async fn connect(url: &str) -> Result<Self, String> {
-        // Use unlimited message/frame sizes to handle large CDP responses
-        // (e.g. Accessibility.getFullAXTree) over remote WSS connections where
-        // proxies may produce frames exceeding the default 16 MiB limit.
+        Self::connect_with_headers(url, None).await
+    }
+
+    pub async fn connect_with_headers(
+        url: &str,
+        headers: Option<Vec<(String, String)>>,
+    ) -> Result<Self, String> {
+        let mut request = url
+            .into_client_request()
+            .map_err(|e| format!("Invalid WebSocket URL: {}", e))?;
+
+        if let Some(hdrs) = headers {
+            let req_headers = request.headers_mut();
+            for (key, value) in hdrs {
+                if let (Ok(name), Ok(val)) = (
+                    key.parse::<tokio_tungstenite::tungstenite::http::header::HeaderName>(),
+                    value.parse::<tokio_tungstenite::tungstenite::http::header::HeaderValue>(),
+                ) {
+                    req_headers.insert(name, val);
+                }
+            }
+        }
+
        let ws_config = WebSocketConfig {
            max_message_size: None,
            max_frame_size: None,
@@ -56,15 +77,10 @@ impl CdpClient {
        };

        let (ws_stream, _) =
-            tokio_tungstenite::connect_async_with_config(url, Some(ws_config), false)
+            tokio_tungstenite::connect_async_with_config(request, Some(ws_config), false)
                .await
                .map_err(|e| format!("CDP WebSocket connect failed: {}", e))?;

-        // Enable TCP SO_KEEPALIVE on the underlying socket. This matches the
-        // behavior of Playwright's WebSocket transport (pre-v0.20.0) which used
-        // Node.js HTTP agents with keepAlive: true. TCP-level keepalive probes
-        // maintain the connection at the transport layer, complementing the
-        // WebSocket-level Ping frames sent by the keepalive task below.
        enable_tcp_keepalive(ws_stream.get_ref());

        let (ws_tx, mut ws_rx) = ws_stream.split();
@@ -7,11 +7,13 @@ use serde_json::{json, Value};
 use std::env;

 /// Provider session info for cleanup on failure.
+#[derive(Debug)]
 pub struct ProviderSession {
    pub provider: String,
    pub session_id: String,
 }

+#[derive(Debug)]
 pub struct ProviderConnection {
    pub ws_url: String,
    pub session: Option<ProviderSession>,
@@ -55,8 +57,16 @@ pub async fn connect_provider(provider_name: &str) -> Result<ProviderConnection,
                direct_page: false,
            })
        }
+        "agentcore" => {
+            let (url, session) = connect_agentcore().await?;
+            Ok(ProviderConnection {
+                ws_url: url,
+                session,
+                direct_page: false,
+            })
+        }
        _ => Err(format!(
-            "Unknown provider '{}'. Supported: browserbase, browserless, browser-use, kernel",
+            "Unknown provider '{}'. Supported: browserbase, browserless, browser-use, kernel, agentcore",
            provider_name
        )),
    }
@@ -113,6 +123,10 @@ pub async fn close_provider_session(session: &ProviderSession) {
                    .await;
            }
        }
+        "agentcore" => {
+            // AgentCore session cleanup is handled via signed DELETE request
+            let _ = close_agentcore_session(&session.session_id).await;
+        }
        _ => {}
    }
 }
@@ -347,3 +361,457 @@ async fn connect_kernel() -> Result<(String, Option<ProviderSession>), String> {
        }),
    ))
 }
+
+// ============================================================================
+// AgentCore Provider (AWS Bedrock AgentCore Browser)
+// Requires: cargo build --features agentcore
+// ============================================================================
+
+#[cfg(feature = "agentcore")]
+mod agentcore {
+    use super::*;
+
+    /// AgentCore-specific session info for Live View URL
+    pub struct AgentCoreSessionInfo {
+        pub session_id: String,
+        pub browser_identifier: String,
+        pub region: String,
+        pub live_view_url: String,
+    }
+
+    thread_local! {
+        static AGENTCORE_INFO: std::cell::RefCell<Option<AgentCoreSessionInfo>> = const { std::cell::RefCell::new(None) };
+        static AGENTCORE_WS_HEADERS: std::cell::RefCell<Option<Vec<(String, String)>>> = const { std::cell::RefCell::new(None) };
+    }
+
+    pub fn set_agentcore_info(info: AgentCoreSessionInfo) {
+        AGENTCORE_INFO.with(|cell| *cell.borrow_mut() = Some(info));
+    }
+
+    pub fn get_agentcore_info() -> Option<AgentCoreSessionInfo> {
+        AGENTCORE_INFO.with(|cell| cell.borrow().as_ref().map(|i| AgentCoreSessionInfo {
+            session_id: i.session_id.clone(),
+            browser_identifier: i.browser_identifier.clone(),
+            region: i.region.clone(),
+            live_view_url: i.live_view_url.clone(),
+        }))
+    }
+
+    pub fn set_agentcore_ws_headers(headers: Vec<(String, String)>) {
+        AGENTCORE_WS_HEADERS.with(|cell| *cell.borrow_mut() = Some(headers));
+    }
+
+    pub fn take_agentcore_ws_headers() -> Option<Vec<(String, String)>> {
+        AGENTCORE_WS_HEADERS.with(|cell| cell.borrow_mut().take())
+    }
+
+    pub async fn connect() -> Result<(String, Option<ProviderSession>), String> {
+        let region = env::var("AGENTCORE_REGION")
+            .or_else(|_| env::var("AWS_REGION"))
+            .or_else(|_| env::var("AWS_DEFAULT_REGION"))
+            .unwrap_or_else(|_| "us-east-1".to_string());
+        let browser_id = env::var("AGENTCORE_BROWSER_ID")
+            .unwrap_or_else(|_| "aws.browser.v1".to_string());
+        let timeout_secs: u64 = env::var("AGENTCORE_SESSION_TIMEOUT")
+            .ok()
+            .and_then(|v| v.parse().ok())
+            .unwrap_or(3600);
+
+        let host = format!("bedrock-agentcore.{}.amazonaws.com", region);
+        let path = format!("/browsers/{}/sessions/start", urlencoding::encode(&browser_id));
+        let url = format!("https://{}{}", host, path);
+
+        // Generate a unique session name
+        let session_name = format!("agent-browser-{}", &uuid::Uuid::new_v4().to_string()[..8]);
+
+        let mut body_json = json!({
+            "name": session_name,
+            "sessionTimeoutSeconds": timeout_secs
+        });
+        if let Ok(profile_id) = env::var("AGENTCORE_PROFILE_ID") {
+            if !profile_id.is_empty() {
+                body_json.as_object_mut().unwrap().insert(
+                    "profileConfiguration".to_string(),
+                    json!({ "profileIdentifier": profile_id })
+                );
+            }
+        }
+        let body = serde_json::to_string(&body_json)
+            .map_err(|e| format!("Failed to serialize request body: {}", e))?;
+
+        let signed_headers = sign_request("PUT", &url, &region, Some(&body)).await?;
+
+        let client = reqwest::Client::new();
+        let mut req = client.put(&url).body(body.clone());
+        for (key, value) in &signed_headers {
+            req = req.header(key.as_str(), value.as_str());
+        }
+
+        let response = req.send().await
+            .map_err(|e| format!("AgentCore request failed: {}", e))?;
+
+        let status = response.status();
+        let resp_body = response.text().await
+            .map_err(|e| format!("Failed to read AgentCore response: {}", e))?;
+
+        if !status.is_success() {
+            return Err(format!("AgentCore API error ({}): {}", status.as_u16(), resp_body));
+        }
+
+        let json: Value = serde_json::from_str(&resp_body)
+            .map_err(|e| format!("Invalid AgentCore response: {}", e))?;
+
+        let session_id = json.get("sessionId")
+            .and_then(|v| v.as_str())
+            .ok_or_else(|| "AgentCore response missing sessionId".to_string())?
+            .to_string();
+
+        let browser_identifier = json.get("browserIdentifier")
+            .and_then(|v| v.as_str())
+            .unwrap_or(&browser_id)
+            .to_string();
+
+        let live_view_url = format!(
+            "https://{}.console.aws.amazon.com/bedrock-agentcore/browser/{}/session/{}#",
+            region, browser_identifier, session_id
+        );
+
+        set_agentcore_info(AgentCoreSessionInfo {
+            session_id: session_id.clone(),
+            browser_identifier: browser_identifier.clone(),
+            region: region.clone(),
+            live_view_url: live_view_url.clone(),
+        });
+
+        eprintln!("Session: {}", session_id);
+        eprintln!("Live View: {}", live_view_url);
+
+        let ws_path = format!("/browser-streams/{}/sessions/{}/automation", browser_identifier, session_id);
+        let ws_url = format!("wss://{}{}", host, ws_path);
+
+        let ws_headers = sign_request("GET", &format!("https://{}{}", host, ws_path), &region, None).await?;
+        set_agentcore_ws_headers(ws_headers);
+
+        Ok((
+            ws_url,
+            Some(ProviderSession {
+                provider: "agentcore".to_string(),
+                session_id,
+            }),
+        ))
+    }
+
+    /// Get AWS credentials from environment variables or AWS CLI
+    fn get_aws_credentials() -> Result<(String, String, Option<String>), String> {
+        // First try environment variables
+        if let (Ok(access_key), Ok(secret_key)) = (
+            env::var("AWS_ACCESS_KEY_ID"),
+            env::var("AWS_SECRET_ACCESS_KEY"),
+        ) {
+            return Ok((access_key, secret_key, env::var("AWS_SESSION_TOKEN").ok()));
+        }
+
+        // Fall back to AWS CLI
+        let mut cmd = std::process::Command::new("aws");
+        cmd.args(["configure", "export-credentials", "--format", "env"]);
+        
+        // Honor AWS_PROFILE
+        if let Ok(profile) = env::var("AWS_PROFILE") {
+            cmd.args(["--profile", &profile]);
+        }
+
+        let output = cmd.output()
+            .map_err(|e| format!("Failed to run aws CLI: {}. Install AWS CLI or set AWS_ACCESS_KEY_ID/AWS_SECRET_ACCESS_KEY", e))?;
+
+        if !output.status.success() {
+            let stderr = String::from_utf8_lossy(&output.stderr);
+            return Err(format!("AWS CLI failed: {}. Run 'aws sso login' or set credentials", stderr.trim()));
+        }
+
+        let stdout = String::from_utf8_lossy(&output.stdout);
+        let mut access_key = None;
+        let mut secret_key = None;
+        let mut session_token = None;
+
+        for line in stdout.lines() {
+            if let Some(val) = line.strip_prefix("export AWS_ACCESS_KEY_ID=") {
+                access_key = Some(val.to_string());
+            } else if let Some(val) = line.strip_prefix("export AWS_SECRET_ACCESS_KEY=") {
+                secret_key = Some(val.to_string());
+            } else if let Some(val) = line.strip_prefix("export AWS_SESSION_TOKEN=") {
+                session_token = Some(val.to_string());
+            }
+        }
+
+        match (access_key, secret_key) {
+            (Some(ak), Some(sk)) => Ok((ak, sk, session_token)),
+            _ => Err("Failed to parse credentials from AWS CLI output".to_string()),
+        }
+    }
+
+    async fn sign_request(
+        method: &str,
+        url: &str,
+        region: &str,
+        body: Option<&str>,
+    ) -> Result<Vec<(String, String)>, String> {
+        use hmac::{Hmac, Mac};
+        use sha2::{Sha256, Digest};
+
+        // Get credentials from environment or AWS CLI
+        let (access_key, secret_key, session_token) = get_aws_credentials()?;
+
+        let parsed_url = url::Url::parse(url)
+            .map_err(|e| format!("Invalid URL: {}", e))?;
+        let host = parsed_url.host_str().unwrap_or("");
+
+        // Get current time
+        let now = chrono::Utc::now();
+        let amz_date = now.format("%Y%m%dT%H%M%SZ").to_string();
+        let date_stamp = now.format("%Y%m%d").to_string();
+
+        // Create canonical request
+        let payload_hash = if let Some(b) = body {
+            let mut hasher = Sha256::new();
+            hasher.update(b.as_bytes());
+            hex::encode(hasher.finalize())
+        } else {
+            "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855".to_string() // empty string hash
+        };
+
+        let canonical_uri = parsed_url.path();
+        let canonical_querystring = parsed_url.query().unwrap_or("");
+
+        let mut signed_headers = "content-type;host;x-amz-date".to_string();
+        let mut canonical_headers = format!(
+            "content-type:application/json\nhost:{}\nx-amz-date:{}\n",
+            host, amz_date
+        );
+
+        if session_token.is_some() {
+            signed_headers = "content-type;host;x-amz-date;x-amz-security-token".to_string();
+            canonical_headers = format!(
+                "content-type:application/json\nhost:{}\nx-amz-date:{}\nx-amz-security-token:{}\n",
+                host, amz_date, session_token.as_ref().unwrap()
+            );
+        }
+
+        let canonical_request = format!(
+            "{}\n{}\n{}\n{}\n{}\n{}",
+            method, canonical_uri, canonical_querystring,
+            canonical_headers, signed_headers, payload_hash
+        );
+
+        // Create string to sign
+        let algorithm = "AWS4-HMAC-SHA256";
+        let credential_scope = format!("{}/{}/bedrock-agentcore/aws4_request", date_stamp, region);
+
+        let mut hasher = Sha256::new();
+        hasher.update(canonical_request.as_bytes());
+        let canonical_request_hash = hex::encode(hasher.finalize());
+
+        let string_to_sign = format!(
+            "{}\n{}\n{}\n{}",
+            algorithm, amz_date, credential_scope, canonical_request_hash
+        );
+
+        // Calculate signature
+        type HmacSha256 = Hmac<Sha256>;
+
+        let k_date = HmacSha256::new_from_slice(format!("AWS4{}", secret_key).as_bytes())
+            .unwrap()
+            .chain_update(date_stamp.as_bytes())
+            .finalize()
+            .into_bytes();
+
+        let k_region = HmacSha256::new_from_slice(&k_date)
+            .unwrap()
+            .chain_update(region.as_bytes())
+            .finalize()
+            .into_bytes();
+
+        let k_service = HmacSha256::new_from_slice(&k_region)
+            .unwrap()
+            .chain_update(b"bedrock-agentcore")
+            .finalize()
+            .into_bytes();
+
+        let k_signing = HmacSha256::new_from_slice(&k_service)
+            .unwrap()
+            .chain_update(b"aws4_request")
+            .finalize()
+            .into_bytes();
+
+        let signature = hex::encode(
+            HmacSha256::new_from_slice(&k_signing)
+                .unwrap()
+                .chain_update(string_to_sign.as_bytes())
+                .finalize()
+                .into_bytes()
+        );
+
+        // Build authorization header
+        let authorization = format!(
+            "{} Credential={}/{}, SignedHeaders={}, Signature={}",
+            algorithm, access_key, credential_scope, signed_headers, signature
+        );
+
+        let mut headers = vec![
+            ("host".to_string(), host.to_string()),
+            ("content-type".to_string(), "application/json".to_string()),
+            ("x-amz-date".to_string(), amz_date),
+            ("authorization".to_string(), authorization),
+        ];
+
+        if let Some(token) = session_token {
+            headers.push(("x-amz-security-token".to_string(), token));
+        }
+
+        Ok(headers)
+    }
+
+    pub async fn close_session(session_id: &str) -> Result<(), String> {
+        let info = get_agentcore_info();
+        let (region, browser_id) = match &info {
+            Some(i) => (i.region.clone(), i.browser_identifier.clone()),
+            None => {
+                let region = env::var("AGENTCORE_REGION")
+                    .or_else(|_| env::var("AWS_REGION"))
+                    .or_else(|_| env::var("AWS_DEFAULT_REGION"))
+                    .unwrap_or_else(|_| "us-east-1".to_string());
+                let browser_id = env::var("AGENTCORE_BROWSER_ID")
+                    .unwrap_or_else(|_| "aws.browser.v1".to_string());
+                (region, browser_id)
+            }
+        };
+
+        let host = format!("bedrock-agentcore.{}.amazonaws.com", region);
+        let path = format!("/browsers/{}/sessions/stop", urlencoding::encode(&browser_id));
+        let url = format!("https://{}{}", host, path);
+
+        let body = serde_json::to_string(&json!({ "sessionId": session_id }))
+            .map_err(|e| format!("Failed to serialize close request: {}", e))?;
+
+        let signed_headers = sign_request("PUT", &url, &region, Some(&body)).await?;
+
+        let client = reqwest::Client::new();
+        let mut req = client.put(&url).body(body);
+        for (key, value) in &signed_headers {
+            req = req.header(key.as_str(), value.as_str());
+        }
+
+        let _ = req.send().await;
+        Ok(())
+    }
+}
+
+#[cfg(feature = "agentcore")]
+pub use agentcore::{get_agentcore_info, take_agentcore_ws_headers};
+
+#[cfg(feature = "agentcore")]
+async fn connect_agentcore() -> Result<(String, Option<ProviderSession>), String> {
+    agentcore::connect().await
+}
+
+#[cfg(not(feature = "agentcore"))]
+async fn connect_agentcore() -> Result<(String, Option<ProviderSession>), String> {
+    Err("AgentCore provider requires the 'agentcore' feature. Rebuild with: cargo build --features agentcore".to_string())
+}
+
+#[cfg(feature = "agentcore")]
+async fn close_agentcore_session(session_id: &str) -> Result<(), String> {
+    agentcore::close_session(session_id).await
+}
+
+#[cfg(not(feature = "agentcore"))]
+async fn close_agentcore_session(_session_id: &str) -> Result<(), String> {
+    Ok(())
+}
+
+// Stub functions when agentcore feature is disabled
+#[cfg(not(feature = "agentcore"))]
+pub fn get_agentcore_info() -> Option<()> { None }
+
+#[cfg(not(feature = "agentcore"))]
+pub fn take_agentcore_ws_headers() -> Option<Vec<(String, String)>> { None }
+
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_connect_provider_unknown() {
+        let rt = tokio::runtime::Runtime::new().unwrap();
+        let result = rt.block_on(connect_provider("unknown-provider"));
+        assert!(result.is_err());
+        assert!(result.unwrap_err().contains("Unknown provider"));
+    }
+
+    #[test]
+    fn test_connect_provider_agentcore_without_feature() {
+        // Without agentcore feature, should return helpful error
+        #[cfg(not(feature = "agentcore"))]
+        {
+            let rt = tokio::runtime::Runtime::new().unwrap();
+            let result = rt.block_on(connect_provider("agentcore"));
+            assert!(result.is_err());
+            assert!(result.unwrap_err().contains("agentcore"));
+        }
+    }
+
+    #[cfg(feature = "agentcore")]
+    #[test]
+    fn test_agentcore_env_defaults() {
+        // Test that default values are used when env vars not set
+        std::env::remove_var("AGENTCORE_REGION");
+        std::env::remove_var("AGENTCORE_BROWSER_ID");
+        std::env::remove_var("AGENTCORE_SESSION_TIMEOUT");
+
+        // These would be used in connect() - just verify they don't panic
+        let region = std::env::var("AGENTCORE_REGION")
+            .or_else(|_| std::env::var("AWS_REGION"))
+            .unwrap_or_else(|_| "us-east-1".to_string());
+        assert_eq!(region, "us-east-1");
+
+        let browser_id = std::env::var("AGENTCORE_BROWSER_ID")
+            .unwrap_or_else(|_| "aws.browser.v1".to_string());
+        assert_eq!(browser_id, "aws.browser.v1");
+    }
+
+    #[cfg(feature = "agentcore")]
+    #[test]
+    fn test_agentcore_session_info_storage() {
+        let info = agentcore::AgentCoreSessionInfo {
+            session_id: "test-session".to_string(),
+            browser_identifier: "aws.browser.v1".to_string(),
+            region: "us-east-1".to_string(),
+            live_view_url: "https://example.com".to_string(),
+        };
+
+        agentcore::set_agentcore_info(info);
+        let retrieved = get_agentcore_info();
+        assert!(retrieved.is_some());
+        let retrieved = retrieved.unwrap();
+        assert_eq!(retrieved.session_id, "test-session");
+        assert_eq!(retrieved.region, "us-east-1");
+    }
+
+    #[cfg(feature = "agentcore")]
+    #[test]
+    fn test_agentcore_ws_headers_storage() {
+        let headers = vec![
+            ("Authorization".to_string(), "AWS4-HMAC-SHA256...".to_string()),
+            ("X-Amz-Date".to_string(), "20260304T180000Z".to_string()),
+        ];
+
+        agentcore::set_agentcore_ws_headers(headers);
+        let taken = take_agentcore_ws_headers();
+        assert!(taken.is_some());
+        assert_eq!(taken.unwrap().len(), 2);
+
+        // Should be None after take
+        let taken_again = take_agentcore_ws_headers();
+        assert!(taken_again.is_none());
+    }
+}
@@ -2808,7 +2808,7 @@ Options:
                             e.g., --proxy-bypass "localhost,*.internal.com"
  --ignore-https-errors      Ignore HTTPS certificate errors
  --allow-file-access        Allow file:// URLs to access local files (Chromium only)
-  -p, --provider <name>      Browser provider: ios, browserbase, kernel, browseruse, browserless
+  -p, --provider <name>      Browser provider: ios, browserbase, kernel, browseruse, browserless, agentcore
  --device <name>            iOS device name (e.g., "iPhone 15 Pro")
  --json                     JSON output
  --annotate                 Annotated screenshot with numbered labels and legend
@@ -2863,7 +2863,7 @@ Environment:
  AGENT_BROWSER_ANNOTATE         Annotated screenshot with numbered labels and legend
  AGENT_BROWSER_DEBUG            Debug output
  AGENT_BROWSER_IGNORE_HTTPS_ERRORS Ignore HTTPS certificate errors
-  AGENT_BROWSER_PROVIDER         Browser provider (ios, browserbase, kernel, browseruse, browserless)
+  AGENT_BROWSER_PROVIDER         Browser provider (ios, browserbase, kernel, browseruse, browserless, agentcore)
  AGENT_BROWSER_AUTO_CONNECT     Auto-discover and connect to running Chrome
  AGENT_BROWSER_ALLOW_FILE_ACCESS Allow file:// URLs to access local files
  AGENT_BROWSER_COLOR_SCHEME     Color scheme preference (dark, light, no-preference)