chrome-use/cli/src/native/relay.rs

//! Relay between the `ab-connect` browser extension and the daemon's `CdpClient`.
//!
//! The extension speaks a small CDP-over-WebSocket "envelope" protocol (adapted
//! from openclaw-browser-relay) and drives the user's real tabs via per-tab
//! `chrome.debugger`. The daemon's `CdpClient`, however, expects a **browser-
//! level** CDP endpoint (`Target.getTargets` / `Target.attachToTarget` → a
//! `sessionId`, then per-session commands). This relay bridges the two: it
//! tracks the targets the extension reports, answers the browser-level
//! `Target.*` discovery commands LOCALLY, and forwards everything else to the
//! extension as `forwardCDPCommand`. That keeps `CdpClient` and `browser.rs`
//! unchanged.
//!
//! ## Multiple clients (concurrent agents on one shared browser)
//!
//! Several agent-browser daemons (one per `--session`) can connect to the same
//! relay/Chrome at once. The extension is a single peer, so the relay must
//! demultiplex: every forwarded command is re-keyed to a relay-global id mapped
//! back to the originating client, and the extension's reply is routed to **only
//! that client** (with its original id restored). Command ids from different
//! clients therefore never collide, and one client never sees another's command
//! replies. CDP *events* (no id) fan out to all clients, which ignore events for
//! sessions they didn't attach.
//!
//! This module is the pure translation core (no I/O) so the protocol can be
//! unit-tested; the tokio WebSocket server that drives it lives alongside.

use std::collections::HashMap;

use serde_json::{json, Value};

/// Protocol version advertised in the connect handshake (matches the extension).
pub const RELAY_PROTOCOL: i64 = 3;

/// Identifies one connected CDP client (agent-browser daemon) for routing.
pub type ClientId = u64;

/// One target (tab) the extension has attached, as the relay tracks it.
#[derive(Clone)]
struct TargetEntry {
    session_id: String,
    target_info: Value,
}

/// Relay translation state: the targets the extension exposes, plus the
/// in-flight command map used to route extension replies back to the right
/// client.
#[derive(Default)]
pub struct RelayState {
    /// targetId -> entry
    targets: HashMap<String, TargetEntry>,
    /// relay-global command id -> (client that sent it, its original id)
    pending: HashMap<i64, (ClientId, Value)>,
    /// monotonic source of relay-global command ids
    next_global_id: i64,
}

/// What to do with a raw CDP command received from a `CdpClient`.
#[derive(Debug, PartialEq)]
pub enum ClientRoute {
    /// Answer locally; the value is a raw CDP response `{id, result}` to send
    /// back to the originating client only.
    Local(Value),
    /// Forward to the extension; the value is a `forwardCDPCommand` envelope
    /// already re-keyed to a relay-global id.
    Forward(Value),
}

/// An output the relay emits while handling an extension message.
#[derive(Debug, PartialEq)]
pub enum RelayOut {
    /// Send this raw CDP message to clients. `to = Some(id)` targets one client
    /// (a command reply); `to = None` broadcasts (a CDP event).
    ToClient { to: Option<ClientId>, msg: Value },
    /// Send this envelope message back to the extension.
    ToExt(Value),
}

impl RelayState {
    pub fn new() -> Self {
        Self::default()
    }

    /// The challenge the relay sends to the extension as soon as it connects,
    /// kicking off the connect handshake.
    pub fn connect_challenge(nonce: &str) -> Value {
        json!({ "type": "event", "event": "connect.challenge", "payload": { "nonce": nonce } })
    }

    /// A keepalive ping for the extension.
    pub fn ping() -> Value {
        json!({ "method": "ping" })
    }

    /// Forget a disconnected client's in-flight commands so its orphaned
    /// `pending` entries don't leak.
    pub fn drop_client(&mut self, client_id: ClientId) {
        self.pending.retain(|_, (cid, _)| *cid != client_id);
    }

    /// Route a raw CDP command `{id, method, params?, sessionId?}` from a
    /// `CdpClient`: answer browser-level `Target.*` discovery locally, forward
    /// the rest to the extension under a relay-global id keyed to `client_id`.
    pub fn route_client_command(&mut self, client_id: ClientId, raw: &Value) -> ClientRoute {
        let id = raw.get("id").cloned().unwrap_or(Value::Null);
        let method = raw.get("method").and_then(|m| m.as_str()).unwrap_or("");
        let params = raw.get("params").cloned().unwrap_or_else(|| json!({}));
        let session_id = raw.get("sessionId").and_then(|s| s.as_str());

        match method {
            // Browser-level command the daemon uses as its liveness probe
            // (`is_connection_alive` → `Browser.getVersion`). The extension only
            // speaks per-tab `chrome.debugger`, so forwarding it errors → the
            // daemon would deem the connection dead and reconnect+re-discover on
            // EVERY command, resetting the active tab (eval/screenshot drift).
            // Answer it locally so the relay connection reads as alive.
            "Browser.getVersion" => ClientRoute::Local(json!({
                "id": id,
                "result": {
                    "protocolVersion": "1.3",
                    "product": "Chrome/ab-connect-relay",
                    "revision": "",
                    "userAgent": "",
                    "jsVersion": ""
                }
            })),
            // Discovery is best-effort and event-driven in real CDP; abs only
            // reads the getTargets result, so an empty ack is enough here.
            "Target.setDiscoverTargets" | "Target.setAutoAttach" => {
                ClientRoute::Local(json!({ "id": id, "result": {} }))
            }
            "Target.getTargets" => {
                let infos: Vec<Value> = self
                    .targets
                    .values()
                    .map(|t| t.target_info.clone())
                    .collect();
                ClientRoute::Local(json!({ "id": id, "result": { "targetInfos": infos } }))
            }
            "Target.attachToTarget" => {
                let target_id = params
                    .get("targetId")
                    .and_then(|t| t.as_str())
                    .unwrap_or("");
                match self.targets.get(target_id) {
                    Some(entry) => ClientRoute::Local(
                        json!({ "id": id, "result": { "sessionId": entry.session_id } }),
                    ),
                    None => ClientRoute::Local(json!({
                        "id": id,
                        "error": { "code": -32602, "message": format!("No such target {target_id}") }
                    })),
                }
            }
            // Everything else goes to the extension's chrome.debugger. Re-key the
            // id so this client's reply can be routed back unambiguously.
            _ => {
                self.next_global_id += 1;
                let gid = self.next_global_id;
                self.pending.insert(gid, (client_id, id));
                ClientRoute::Forward(json!({
                    "id": gid,
                    "method": "forwardCDPCommand",
                    "params": { "method": method, "params": params, "sessionId": session_id },
                }))
            }
        }
    }

    /// Handle one decoded message from the extension. Updates target state and
    /// returns the messages to emit (routed to a client and/or back to the
    /// extension). `expected_token` is matched against the connect handshake.
    pub fn handle_ext_message(&mut self, msg: &Value, expected_token: &str) -> Vec<RelayOut> {
        // Connect handshake request from the extension.
        if msg.get("type").and_then(|t| t.as_str()) == Some("req")
            && msg.get("method").and_then(|m| m.as_str()) == Some("connect")
        {
            let id = msg.get("id").cloned().unwrap_or(Value::Null);
            let token = msg
                .get("params")
                .and_then(|p| p.get("auth"))
                .and_then(|a| a.get("token"))
                .and_then(|t| t.as_str())
                .unwrap_or("");
            let ok = !expected_token.is_empty() && token == expected_token;
            let mut res = json!({ "type": "res", "id": id, "ok": ok });
            if !ok {
                res["error"] = json!({ "message": "invalid relay token" });
            }
            return vec![RelayOut::ToExt(res)];
        }

        // Keepalive.
        if msg.get("method").and_then(|m| m.as_str()) == Some("pong") {
            return vec![];
        }

        // Response to a forwardCDPCommand we sent → route the raw CDP response
        // back to the client that issued it, with its original id restored.
        if msg.get("id").is_some()
            && (msg.get("result").is_some() || msg.get("error").is_some())
            && msg.get("method").is_none()
        {
            let gid = msg.get("id").and_then(|i| i.as_i64());
            let (to, orig_id) = match gid.and_then(|g| self.pending.remove(&g)) {
                Some((client_id, orig)) => (Some(client_id), orig),
                // No mapping (stale/unknown id) — fall back to broadcasting with
                // whatever id the extension echoed.
                None => (None, msg.get("id").cloned().unwrap_or(Value::Null)),
            };
            let mut out = json!({ "id": orig_id });
            if let Some(r) = msg.get("result") {
                out["result"] = r.clone();
            }
            if let Some(e) = msg.get("error") {
                // CdpClient expects an error object; wrap a bare string.
                out["error"] = match e {
                    Value::String(s) => json!({ "code": -32000, "message": s }),
                    other => other.clone(),
                };
            }
            return vec![RelayOut::ToClient { to, msg: out }];
        }

        // CDP event forwarded from a tab.
        if msg.get("method").and_then(|m| m.as_str()) == Some("forwardCDPEvent") {
            let p = msg.get("params").cloned().unwrap_or_else(|| json!({}));
            let inner_method = p.get("method").and_then(|m| m.as_str()).unwrap_or("");
            let inner_params = p.get("params").cloned().unwrap_or_else(|| json!({}));
            let session_id = p.get("sessionId").and_then(|s| s.as_str());

            // Learn/forget targets from the extension's synthesized Target events.
            // We consume these to maintain state and do NOT forward them: abs
            // discovers targets by pulling getTargets, and forwarding a second
            // attachedToTarget would duplicate the one attachToTarget emits.
            match inner_method {
                "Target.attachedToTarget" => {
                    if let Some(info) = inner_params.get("targetInfo") {
                        if let Some(tid) = info.get("targetId").and_then(|t| t.as_str()) {
                            let sid = inner_params
                                .get("sessionId")
                                .and_then(|s| s.as_str())
                                .unwrap_or("")
                                .to_string();
                            self.targets.insert(
                                tid.to_string(),
                                TargetEntry {
                                    session_id: sid,
                                    target_info: info.clone(),
                                },
                            );
                        }
                    }
                    return vec![];
                }
                "Target.detachedFromTarget" => {
                    let gone = inner_params.get("sessionId").and_then(|s| s.as_str());
                    if let Some(gone) = gone {
                        self.targets.retain(|_, e| e.session_id != gone);
                    }
                    return vec![];
                }
                _ => {}
            }

            // Regular CDP event → fan out to all clients (each filters by the
            // sessions it attached to).
            let mut ev = json!({ "method": inner_method, "params": inner_params });
            if let Some(sid) = session_id {
                ev["sessionId"] = json!(sid);
            }
            return vec![RelayOut::ToClient { to: None, msg: ev }];
        }

        vec![]
    }

    #[cfg(test)]
    fn seed_target(&mut self, target_id: &str, session_id: &str) {
        self.targets.insert(
            target_id.to_string(),
            TargetEntry {
                session_id: session_id.to_string(),
                target_info: json!({
                    "targetId": target_id,
                    "type": "page",
                    "title": "",
                    "url": "about:blank",
                    "attached": true,
                }),
            },
        );
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn attached_event(target_id: &str, session_id: &str) -> Value {
        json!({
            "method": "forwardCDPEvent",
            "params": {
                "sessionId": session_id,
                "method": "Target.attachedToTarget",
                "params": {
                    "sessionId": session_id,
                    "targetInfo": { "targetId": target_id, "type": "page", "url": "https://x", "title": "X" }
                }
            }
        })
    }

    #[test]
    fn learns_target_from_attached_event_and_does_not_forward_it() {
        let mut s = RelayState::new();
        let out = s.handle_ext_message(&attached_event("T1", "cb-tab-1"), "tok");
        assert!(
            out.is_empty(),
            "attachedToTarget should be consumed, not forwarded"
        );
        // Now getTargets must report it.
        let route = s.route_client_command(1, &json!({ "id": 1, "method": "Target.getTargets" }));
        match route {
            ClientRoute::Local(v) => {
                let infos = v["result"]["targetInfos"].as_array().unwrap();
                assert_eq!(infos.len(), 1);
                assert_eq!(infos[0]["targetId"], "T1");
            }
            _ => panic!("getTargets must be local"),
        }
    }

    #[test]
    fn browser_get_version_is_answered_locally() {
        // Liveness probe must NOT be forwarded (the extension can't do
        // browser-level commands) — else the daemon reconnects on every command.
        let mut s = RelayState::new();
        let route = s.route_client_command(1, &json!({ "id": 7, "method": "Browser.getVersion" }));
        match route {
            ClientRoute::Local(v) => {
                assert_eq!(v["id"], 7);
                assert!(v["result"]["protocolVersion"].is_string());
            }
            _ => panic!("Browser.getVersion must be answered locally"),
        }
    }

    #[test]
    fn attach_to_target_returns_known_session() {
        let mut s = RelayState::new();
        s.seed_target("T1", "cb-tab-1");
        let route = s.route_client_command(
            7,
            &json!({ "id": 5, "method": "Target.attachToTarget", "params": { "targetId": "T1", "flatten": true } }),
        );
        assert_eq!(
            route,
            ClientRoute::Local(json!({ "id": 5, "result": { "sessionId": "cb-tab-1" } }))
        );
    }

    #[test]
    fn attach_to_unknown_target_errors_locally() {
        let mut s = RelayState::new();
        let route = s.route_client_command(
            1,
            &json!({ "id": 6, "method": "Target.attachToTarget", "params": { "targetId": "nope" } }),
        );
        match route {
            ClientRoute::Local(v) => assert!(v.get("error").is_some()),
            _ => panic!("should answer locally"),
        }
    }

    #[test]
    fn other_commands_forward_under_global_id() {
        let mut s = RelayState::new();
        let route = s.route_client_command(
            42,
            &json!({ "id": 9, "method": "Page.navigate", "params": { "url": "https://x" }, "sessionId": "cb-tab-1" }),
        );
        match route {
            ClientRoute::Forward(v) => {
                assert_eq!(v["method"], "forwardCDPCommand");
                // id is re-keyed to a relay-global id (not the client's 9).
                assert_eq!(v["id"], 1);
                assert_eq!(v["params"]["method"], "Page.navigate");
                assert_eq!(v["params"]["sessionId"], "cb-tab-1");
                assert_eq!(v["params"]["params"]["url"], "https://x");
            }
            _ => panic!("Page.navigate must forward"),
        }
    }

    #[test]
    fn reply_routes_back_to_the_issuing_client_with_original_id() {
        let mut s = RelayState::new();
        // Two clients each send a command that happens to share original id 1.
        let r1 = s.route_client_command(
            100,
            &json!({ "id": 1, "method": "Page.navigate", "params": {} }),
        );
        let r2 = s.route_client_command(
            200,
            &json!({ "id": 1, "method": "Page.reload", "params": {} }),
        );
        let g1 = match r1 {
            ClientRoute::Forward(v) => v["id"].as_i64().unwrap(),
            _ => panic!(),
        };
        let g2 = match r2 {
            ClientRoute::Forward(v) => v["id"].as_i64().unwrap(),
            _ => panic!(),
        };
        assert_ne!(g1, g2, "global ids must be distinct across clients");

        // Extension replies for g2 → must go to client 200 with original id 1.
        let out = s.handle_ext_message(&json!({ "id": g2, "result": { "ok": true } }), "tok");
        assert_eq!(
            out,
            vec![RelayOut::ToClient {
                to: Some(200),
                msg: json!({ "id": 1, "result": { "ok": true } })
            }]
        );
        // And g1 → client 100.
        let out = s.handle_ext_message(&json!({ "id": g1, "result": { "ok": false } }), "tok");
        assert_eq!(
            out,
            vec![RelayOut::ToClient {
                to: Some(100),
                msg: json!({ "id": 1, "result": { "ok": false } })
            }]
        );
    }

    #[test]
    fn forward_command_error_is_wrapped_and_routed() {
        let mut s = RelayState::new();
        let r = s.route_client_command(
            5,
            &json!({ "id": 3, "method": "Page.navigate", "params": {} }),
        );
        let gid = match r {
            ClientRoute::Forward(v) => v["id"].as_i64().unwrap(),
            _ => panic!(),
        };
        let out = s.handle_ext_message(&json!({ "id": gid, "error": "boom" }), "tok");
        match &out[0] {
            RelayOut::ToClient { to, msg } => {
                assert_eq!(*to, Some(5));
                assert_eq!(msg["id"], 3);
                assert_eq!(msg["error"]["message"], "boom");
            }
            _ => panic!("expected ToClient"),
        }
    }

    #[test]
    fn regular_event_broadcasts_with_session() {
        let mut s = RelayState::new();
        let ev = json!({
            "method": "forwardCDPEvent",
            "params": { "sessionId": "cb-tab-1", "method": "Page.loadEventFired", "params": { "timestamp": 1.0 } }
        });
        let out = s.handle_ext_message(&ev, "tok");
        assert_eq!(
            out,
            vec![RelayOut::ToClient {
                to: None,
                msg: json!({
                    "method": "Page.loadEventFired",
                    "params": { "timestamp": 1.0 },
                    "sessionId": "cb-tab-1"
                })
            }]
        );
    }

    #[test]
    fn drop_client_clears_its_pending() {
        let mut s = RelayState::new();
        let r = s.route_client_command(
            9,
            &json!({ "id": 1, "method": "Page.navigate", "params": {} }),
        );
        let gid = match r {
            ClientRoute::Forward(v) => v["id"].as_i64().unwrap(),
            _ => panic!(),
        };
        s.drop_client(9);
        // Reply now has no mapping → broadcast fallback (to: None), echoed id.
        let out = s.handle_ext_message(&json!({ "id": gid, "result": {} }), "tok");
        match &out[0] {
            RelayOut::ToClient { to, .. } => assert_eq!(*to, None),
            _ => panic!(),
        }
    }

    #[test]
    fn connect_handshake_validates_token() {
        let mut s = RelayState::new();
        let req = json!({ "type": "req", "id": "c1", "method": "connect", "params": { "auth": { "token": "good" } } });
        let ok = s.handle_ext_message(&req, "good");
        assert_eq!(
            ok,
            vec![RelayOut::ToExt(
                json!({ "type": "res", "id": "c1", "ok": true })
            )]
        );

        let bad = s.handle_ext_message(&req, "different");
        match &bad[0] {
            RelayOut::ToExt(v) => {
                assert_eq!(v["ok"], false);
                assert!(v.get("error").is_some());
            }
            _ => panic!("expected ToExt"),
        }
    }
}