diff --git a/.gitignore b/.gitignore index 5a40802..7bd71de 100644 --- a/.gitignore +++ b/.gitignore @@ -61,6 +61,9 @@ docs/package-lock.json # pnpm .pnpm-store/ +# TypeScript +*.tsbuildinfo + # next .next/ out/ diff --git a/README.md b/README.md index 9bd22b5..06b27af 100644 --- a/README.md +++ b/README.md @@ -203,21 +203,24 @@ agent-browser wait "#spinner" --state hidden ### Batch Execution -Execute multiple commands in a single invocation by piping a JSON array of -string arrays to `batch`. This avoids per-command process startup overhead -when running multi-step workflows. +Execute multiple commands in a single invocation. Commands can be passed as +quoted arguments or piped as JSON via stdin. This avoids per-command process +startup overhead when running multi-step workflows. ```bash -# Pipe commands as JSON +# Argument mode: each quoted argument is a full command +agent-browser batch "open https://example.com" "snapshot -i" "screenshot" + +# With --bail to stop on first error +agent-browser batch --bail "open https://example.com" "click @e1" "screenshot" + +# Stdin mode: pipe commands as JSON echo '[ ["open", "https://example.com"], ["snapshot", "-i"], ["click", "@e1"], ["screenshot", "result.png"] ]' | agent-browser batch --json - -# Stop on first error -agent-browser batch --bail < commands.json ``` ### Clipboard @@ -548,6 +551,7 @@ The `snapshot` command supports filtering to reduce output size: ```bash agent-browser snapshot # Full accessibility tree agent-browser snapshot -i # Interactive elements only (buttons, inputs, links) +agent-browser snapshot -i --urls # Interactive elements with link URLs agent-browser snapshot -c # Compact (remove empty structural elements) agent-browser snapshot -d 3 # Limit depth to 3 levels agent-browser snapshot -s "#main" # Scope to CSS selector @@ -557,6 +561,7 @@ agent-browser snapshot -i -c -d 5 # Combine options | Option | Description | | ---------------------- | ----------------------------------------------------------------------- | | `-i, --interactive` | Only show interactive elements (buttons, links, inputs) | +| `-u, --urls` | Include href URLs for link elements | | `-c, --compact` | Remove empty structural elements | | `-d, --depth ` | Limit tree depth | | `-s, --selector ` | Scope to CSS selector | @@ -650,6 +655,19 @@ The dashboard displays: - **Activity feed** -- chronological command/result stream with timing and expandable details - **Console output** -- browser console messages (log, warn, error) - **Session creation** -- create new sessions from the UI with local engines (Chrome, Lightpanda) or cloud providers (AgentCore, Browserbase, Browserless, Browser Use, Kernel) +- **AI Chat** -- chat with an AI assistant directly in the dashboard (requires Vercel AI Gateway configuration) + +### AI Chat + +The dashboard includes an optional AI chat panel powered by the Vercel AI Gateway. Set these environment variables to enable it: + +```bash +export AI_GATEWAY_API_KEY=gw_your_key_here +export AI_GATEWAY_MODEL=anthropic/claude-sonnet-4.6 # optional, this is the default +export AI_GATEWAY_URL=https://ai-gateway.vercel.sh # optional, this is the default +``` + +The Chat tab is always visible in the dashboard. When `AI_GATEWAY_API_KEY` is set, the Rust server proxies requests to the gateway and streams responses back using the Vercel AI SDK's UI Message Stream protocol. Without the key, sending a message shows an error inline. ## Configuration diff --git a/cli/Cargo.lock b/cli/Cargo.lock index eac2dfd..31612c2 100644 --- a/cli/Cargo.lock +++ b/cli/Cargo.lock @@ -605,6 +605,12 @@ version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" +[[package]] +name = "futures-io" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718" + [[package]] name = "futures-macro" version = "0.3.32" @@ -635,9 +641,11 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" dependencies = [ "futures-core", + "futures-io", "futures-macro", "futures-sink", "futures-task", + "memchr", "pin-project-lite", "slab", ] @@ -1678,6 +1686,7 @@ dependencies = [ "base64", "bytes", "futures-core", + "futures-util", "http", "http-body", "http-body-util", @@ -1697,12 +1706,14 @@ dependencies = [ "sync_wrapper", "tokio", "tokio-rustls", + "tokio-util", "tower", "tower-http", "tower-service", "url", "wasm-bindgen", "wasm-bindgen-futures", + "wasm-streams", "web-sys", "webpki-roots 1.0.5", ] @@ -2128,6 +2139,19 @@ dependencies = [ "webpki-roots 0.26.11", ] +[[package]] +name = "tokio-util" +version = "0.7.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ae9cec805b01e8fc3fd2fe289f89149a9b66dd16786abd8b19cfa7b48cb0098" +dependencies = [ + "bytes", + "futures-core", + "futures-sink", + "pin-project-lite", + "tokio", +] + [[package]] name = "tower" version = "0.5.3" @@ -2430,6 +2454,19 @@ dependencies = [ "wasmparser", ] +[[package]] +name = "wasm-streams" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15053d8d85c7eccdbefef60f06769760a563c7f0a9d6902a13d35c7800b0ad65" +dependencies = [ + "futures-util", + "js-sys", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + [[package]] name = "wasmparser" version = "0.244.0" diff --git a/cli/Cargo.toml b/cli/Cargo.toml index f4e6ab8..877e52a 100644 --- a/cli/Cargo.toml +++ b/cli/Cargo.toml @@ -22,7 +22,7 @@ futures-util = "0.3" url = "2" uuid = { version = "1", features = ["v4"] } image = "0.25" -reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls-webpki-roots"] } +reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls-webpki-roots", "stream"] } sha2 = "0.10" aes-gcm = "0.10" async-trait = "0.1" diff --git a/cli/src/color.rs b/cli/src/color.rs index 07eb526..2c10e0b 100644 --- a/cli/src/color.rs +++ b/cli/src/color.rs @@ -1,15 +1,30 @@ -//! Color output utilities respecting NO_COLOR environment variable. +//! Color output utilities. //! -//! When the NO_COLOR environment variable is present (regardless of value), -//! all color formatting is disabled per https://no-color.org/ +//! Colors are off by default (agent-friendly). Enable with +//! `AGENT_BROWSER_COLOR=1`. Setting `NO_COLOR` to any value disables +//! colors per . use std::env; use std::sync::OnceLock; -/// Returns true if color output is enabled (NO_COLOR is NOT set) +fn env_is_truthy(name: &str) -> Option { + env::var(name) + .ok() + .map(|val| !matches!(val.to_lowercase().as_str(), "0" | "false" | "no")) +} + +/// Returns true if color output is enabled. +/// +/// Priority: `NO_COLOR` (presence disables, per spec) > +/// `AGENT_BROWSER_COLOR` (truthy enables) > default (off). pub fn is_enabled() -> bool { static COLORS_ENABLED: OnceLock = OnceLock::new(); - *COLORS_ENABLED.get_or_init(|| env::var("NO_COLOR").is_err()) + *COLORS_ENABLED.get_or_init(|| { + if env::var_os("NO_COLOR").is_some() { + return false; + } + env_is_truthy("AGENT_BROWSER_COLOR").unwrap_or(false) + }) } /// Format text in red (errors) diff --git a/cli/src/commands.rs b/cli/src/commands.rs index a1b31ba..1c6e365 100644 --- a/cli/src/commands.rs +++ b/cli/src/commands.rs @@ -552,9 +552,11 @@ fn parse_command_inner(args: &[String], flags: &Flags) -> Result { - // deprecated, cursor-interactive elements are referred by default now obj.insert("cursor".to_string(), json!(true)); } + "-u" | "--urls" => { + obj.insert("urls".to_string(), json!(true)); + } "-d" | "--depth" => { if let Some(d) = rest.get(i + 1) { if let Ok(n) = d.parse::() { @@ -1409,7 +1411,12 @@ fn parse_command_inner(args: &[String], flags: &Flags) -> Result { let bail = rest.contains(&"--bail"); - Ok(json!({ "id": id, "action": "batch", "bail": bail })) + let commands: Vec<&str> = rest.iter().filter(|a| **a != "--bail").copied().collect(); + let mut cmd = json!({ "id": id, "action": "batch", "bail": bail }); + if !commands.is_empty() { + cmd["commands"] = json!(commands); + } + Ok(cmd) } _ => Err(ParseError::UnknownCommand { @@ -2289,6 +2296,38 @@ fn parse_storage(rest: &[&str], id: &str) -> Result { } } +/// Split a string into arguments respecting shell quoting (double/single quotes, backslash escapes). +pub fn shell_words_split(s: &str) -> Vec { + let mut args = Vec::new(); + let mut current = String::new(); + let mut in_double = false; + let mut in_single = false; + let mut chars = s.chars().peekable(); + + while let Some(c) = chars.next() { + match c { + '\\' if !in_single => { + if let Some(&next) = chars.peek() { + chars.next(); + current.push(next); + } + } + '"' if !in_single => in_double = !in_double, + '\'' if !in_double => in_single = !in_single, + ' ' if !in_double && !in_single => { + if !current.is_empty() { + args.push(std::mem::take(&mut current)); + } + } + _ => current.push(c), + } + } + if !current.is_empty() { + args.push(current); + } + args +} + #[cfg(test)] mod tests { use super::*; @@ -3005,6 +3044,21 @@ mod tests { assert_eq!(cmd["maxDepth"], 3); } + #[test] + fn test_snapshot_urls() { + let cmd = parse_command(&args("snapshot -i --urls"), &default_flags()).unwrap(); + assert_eq!(cmd["action"], "snapshot"); + assert_eq!(cmd["interactive"], true); + assert_eq!(cmd["urls"], true); + } + + #[test] + fn test_snapshot_urls_short() { + let cmd = parse_command(&args("snapshot -i -u"), &default_flags()).unwrap(); + assert_eq!(cmd["action"], "snapshot"); + assert_eq!(cmd["urls"], true); + } + // === Wait === #[test] @@ -4360,4 +4414,41 @@ mod tests { assert_eq!(cmd["action"], "batch"); assert_eq!(cmd["bail"], true); } + + #[test] + fn test_batch_with_args() { + let cmd_args = vec![ + "batch".to_string(), + "open https://example.com".to_string(), + "screenshot".to_string(), + ]; + let cmd = parse_command(&cmd_args, &default_flags()).unwrap(); + assert_eq!(cmd["action"], "batch"); + assert_eq!(cmd["bail"], false); + let commands = cmd["commands"].as_array().unwrap(); + assert_eq!(commands.len(), 2); + assert_eq!(commands[0], "open https://example.com"); + assert_eq!(commands[1], "screenshot"); + } + + #[test] + fn test_batch_with_args_and_bail() { + let cmd_args = vec![ + "batch".to_string(), + "--bail".to_string(), + "open https://example.com".to_string(), + "screenshot".to_string(), + ]; + let cmd = parse_command(&cmd_args, &default_flags()).unwrap(); + assert_eq!(cmd["action"], "batch"); + assert_eq!(cmd["bail"], true); + let commands = cmd["commands"].as_array().unwrap(); + assert_eq!(commands.len(), 2); + } + + #[test] + fn test_batch_no_args_no_commands_field() { + let cmd = parse_command(&args("batch"), &default_flags()).unwrap(); + assert!(cmd.get("commands").is_none()); + } } diff --git a/cli/src/main.rs b/cli/src/main.rs index 5ab72bb..a8fa038 100644 --- a/cli/src/main.rs +++ b/cli/src/main.rs @@ -1257,10 +1257,16 @@ fn main() { } } - // Handle batch command: read commands from stdin, execute sequentially + // Handle batch command: from args or stdin if cmd.get("action").and_then(|v| v.as_str()) == Some("batch") { let bail = cmd.get("bail").and_then(|v| v.as_bool()).unwrap_or(false); - run_batch(&flags, bail); + let arg_commands = cmd.get("commands").and_then(|v| v.as_array()).map(|arr| { + arr.iter() + .filter_map(|v| v.as_str()) + .map(commands::shell_words_split) + .collect::>>() + }); + run_batch(&flags, bail, arg_commands); return; } @@ -1340,36 +1346,40 @@ fn main() { } } -fn run_batch(flags: &Flags, bail: bool) { - use std::io::Read as _; +fn run_batch(flags: &Flags, bail: bool, arg_commands: Option>>) { + let commands: Vec> = if let Some(cmds) = arg_commands { + cmds + } else { + use std::io::Read as _; - let mut input = String::new(); - if let Err(e) = std::io::stdin().read_to_string(&mut input) { - if flags.json { - print_json_error(format!("Failed to read stdin: {}", e)); - } else { - eprintln!("{} Failed to read stdin: {}", color::error_indicator(), e); - } - exit(1); - } - - let commands: Vec> = match serde_json::from_str(&input) { - Ok(c) => c, - Err(e) => { + let mut input = String::new(); + if let Err(e) = std::io::stdin().read_to_string(&mut input) { if flags.json { - print_json_error(format!( - "Invalid JSON input: {}. Expected an array of string arrays, e.g. [[\"open\", \"https://example.com\"], [\"snapshot\"]]", - e - )); + print_json_error(format!("Failed to read stdin: {}", e)); } else { - eprintln!( - "{} Invalid JSON input: {}. Expected an array of string arrays.", - color::error_indicator(), - e - ); + eprintln!("{} Failed to read stdin: {}", color::error_indicator(), e); } exit(1); } + + match serde_json::from_str(&input) { + Ok(c) => c, + Err(e) => { + if flags.json { + print_json_error(format!( + "Invalid JSON input: {}. Expected an array of string arrays, e.g. [[\"open\", \"https://example.com\"], [\"snapshot\"]]", + e + )); + } else { + eprintln!( + "{} Invalid JSON input: {}. Expected an array of string arrays.", + color::error_indicator(), + e + ); + } + exit(1); + } + } }; if commands.is_empty() { diff --git a/cli/src/native/actions.rs b/cli/src/native/actions.rs index 8a77a74..338d3f5 100644 --- a/cli/src/native/actions.rs +++ b/cli/src/native/actions.rs @@ -2304,6 +2304,7 @@ async fn handle_snapshot(cmd: &Value, state: &mut DaemonState) -> Result, + pub urls: bool, } struct TreeNode { @@ -98,7 +99,8 @@ struct TreeNode { has_ref: bool, ref_id: Option, depth: usize, - cursor_info: Option, // cursor-interactive information + cursor_info: Option, + url: Option, } impl TreeNode { @@ -121,10 +123,10 @@ impl TreeNode { ref_id: None, depth: 0, cursor_info: None, + url: None, } } - // Clear node content fn clear(&mut self) { self.role = String::new(); self.name = String::new(); @@ -139,6 +141,7 @@ impl TreeNode { self.children.clear(); self.parent_idx = None; self.has_ref = false; + self.url = None; self.ref_id = None; self.depth = 0; self.cursor_info = None; @@ -383,6 +386,75 @@ pub async fn take_snapshot( ref_map.set_next_ref_num(next_ref); + if options.urls { + let link_nodes: Vec<(usize, i64)> = tree_nodes + .iter() + .enumerate() + .filter(|(_, n)| n.role == "link" && n.has_ref && n.backend_node_id.is_some()) + .filter_map(|(i, n)| n.backend_node_id.map(|bid| (i, bid))) + .collect(); + + if !link_nodes.is_empty() { + // CDP has no batch resolve API, so we parallelize individual calls. + // Phase 1: resolve all backend node IDs to JS object IDs in parallel. + let resolve_futs = link_nodes.iter().map(|&(idx, bid)| async move { + let resolved = client + .send_command( + "DOM.resolveNode", + Some(serde_json::json!({ "backendNodeId": bid })), + Some(session_id), + ) + .await; + let obj_id = resolved.ok().and_then(|r| { + r.get("object") + .and_then(|o| o.get("objectId")) + .and_then(|v| v.as_str()) + .map(|s| s.to_string()) + }); + (idx, obj_id) + }); + let resolved: Vec<(usize, Option)> = + futures_util::future::join_all(resolve_futs).await; + + // Phase 2: fetch hrefs for all resolved objects in parallel. + let href_futs: Vec<_> = resolved + .iter() + .filter_map(|(idx, obj_id)| { + let oid = obj_id.as_ref()?; + Some(async move { + let result = client + .send_command( + "Runtime.callFunctionOn", + Some(serde_json::json!({ + "objectId": oid, + "functionDeclaration": "function() { return this.href || ''; }", + "returnByValue": true, + })), + Some(session_id), + ) + .await; + let href = result.ok().and_then(|r| { + r.get("result") + .and_then(|r| r.get("value")) + .and_then(|v| v.as_str()) + .filter(|s| !s.is_empty()) + .map(|s| s.to_string()) + }); + (*idx, href) + }) + }) + .collect(); + let hrefs: Vec<(usize, Option)> = + futures_util::future::join_all(href_futs).await; + + for (idx, href) in hrefs { + if let Some(url) = href { + tree_nodes[idx].url = Some(url); + } + } + } + } + let mut output = String::new(); for &root_idx in &effective_roots { render_tree(&tree_nodes, root_idx, 0, &mut output, options); @@ -797,6 +869,7 @@ fn build_tree(nodes: &[AXNode]) -> (Vec, Vec) { ref_id: None, depth: 0, cursor_info: None, + url: None, }); id_to_idx.insert(node.node_id.clone(), i); } @@ -993,6 +1066,10 @@ fn render_tree( attrs.push(format!("ref={}", ref_id)); } + if let Some(ref url) = node.url { + attrs.push(format!("url={}", url)); + } + if !attrs.is_empty() { line.push_str(&format!(" [{}]", attrs.join(", "))); } diff --git a/cli/src/native/stream.rs b/cli/src/native/stream.rs deleted file mode 100644 index 795952f..0000000 --- a/cli/src/native/stream.rs +++ /dev/null @@ -1,1785 +0,0 @@ -use serde_json::{json, Value}; -use std::net::SocketAddr; -use std::path::{Path, PathBuf}; -use std::sync::Arc; - -use futures_util::{SinkExt, StreamExt}; -use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt}; -use tokio::net::TcpListener; -use tokio::sync::{broadcast, watch, Mutex, Notify, RwLock}; -use tokio_tungstenite::tungstenite::Message; - -use super::cdp::client::CdpClient; -use super::network; -#[cfg(windows)] -use crate::connection::get_port_for_session; -use crate::connection::get_socket_dir; -#[cfg(windows)] -use crate::connection::resolve_port; -use crate::install::get_dashboard_dir; - -/// Frame metadata from CDP Page.screencastFrame events. -#[derive(Debug, Clone)] -pub struct FrameMetadata { - pub offset_top: f64, - pub page_scale_factor: f64, - pub device_width: u32, - pub device_height: u32, - pub scroll_offset_x: f64, - pub scroll_offset_y: f64, - pub timestamp: u64, -} - -impl Default for FrameMetadata { - fn default() -> Self { - Self { - offset_top: 0.0, - page_scale_factor: 1.0, - device_width: 1280, - device_height: 720, - scroll_offset_x: 0.0, - scroll_offset_y: 0.0, - timestamp: 0, - } - } -} - -pub struct StreamServer { - port: u16, - session_name: String, - frame_tx: broadcast::Sender, - client_count: Arc>, - client_slot: Arc>>>, - /// The active CDP page session ID (from Target.attachToTarget). - cdp_session_id: Arc>>, - client_notify: Arc, - screencasting: Arc>, - viewport_width: Arc>, - viewport_height: Arc>, - dashboard_dir: Option, - last_tabs: Arc>>, - last_engine: Arc>, - last_frame: Arc>>, - recording: Arc>, - shutdown_tx: watch::Sender, - accept_task: Mutex>>, - cdp_task: Mutex>>, -} - -impl StreamServer { - pub async fn start( - preferred_port: u16, - client: Arc, - session_id: String, - ) -> Result { - let client_slot = Arc::new(RwLock::new(Some(client))); - let (server, _) = Self::start_inner(preferred_port, client_slot, session_id, true).await?; - Ok(server) - } - - /// Start the stream server without a CDP client. - /// Returns the server and a shared slot to set the client when the browser launches. - /// Input messages are ignored until the client is set. - /// When `allow_port_fallback` is true, binding to an occupied port falls back to an - /// OS-assigned port (used by daemon startup). When false, the error propagates - /// (used by the runtime `stream_enable` command). - pub async fn start_without_client( - preferred_port: u16, - session_id: String, - allow_port_fallback: bool, - ) -> Result<(Self, Arc>>>), String> { - let client_slot = Arc::new(RwLock::new(None::>)); - Self::start_inner(preferred_port, client_slot, session_id, allow_port_fallback).await - } - - /// Resolve the dashboard directory if it exists. - fn resolve_dashboard_dir() -> Option { - let dir = dirs::home_dir()?.join(".agent-browser").join("dashboard"); - if dir.join("index.html").exists() { - Some(dir) - } else { - None - } - } - - /// Notify the background CDP listener that the client has changed (browser launched/closed). - pub fn notify_client_changed(&self) { - self.client_notify.notify_one(); - } - - /// Update the active CDP page session ID used for screencast commands. - pub async fn set_cdp_session_id(&self, session_id: Option) { - let mut guard = self.cdp_session_id.write().await; - *guard = session_id; - } - - /// Check whether the server currently has active screencast running. - pub async fn is_screencasting(&self) -> bool { - *self.screencasting.lock().await - } - - /// Update the stored viewport dimensions used by status messages and screencast. - /// Also notifies the screencast event loop to restart with the new dimensions. - pub async fn set_viewport(&self, width: u32, height: u32) { - *self.viewport_width.lock().await = width; - *self.viewport_height.lock().await = height; - self.client_notify.notify_one(); - } - - /// Get the current viewport dimensions. - pub async fn viewport(&self) -> (u32, u32) { - let w = *self.viewport_width.lock().await; - let h = *self.viewport_height.lock().await; - (w, h) - } - - /// Override the cached screencast state for explicit CLI start/stop commands. - pub async fn set_screencasting(&self, active: bool) { - let mut guard = self.screencasting.lock().await; - *guard = active; - } - - /// Update and broadcast the recording state. - pub async fn set_recording(&self, active: bool, engine: &str) { - *self.recording.lock().await = active; - let connected = self.client_slot.read().await.is_some(); - let sc = *self.screencasting.lock().await; - let (vw, vh) = self.viewport().await; - self.broadcast_status(connected, sc, vw, vh, engine).await; - } - - /// Shut down the accept loop and background CDP listener, releasing the bound port. - pub async fn shutdown(&self) { - let _ = self.shutdown_tx.send(true); - - if let Some(task) = self.accept_task.lock().await.take() { - let _ = task.await; - } - if let Some(task) = self.cdp_task.lock().await.take() { - let _ = task.await; - } - } - - async fn start_inner( - preferred_port: u16, - client_slot: Arc>>>, - session_id: String, - allow_port_fallback: bool, - ) -> Result<(Self, Arc>>>), String> { - let addr = format!("127.0.0.1:{}", preferred_port); - let listener = match TcpListener::bind(&addr).await { - Ok(l) => l, - Err(_) if allow_port_fallback && preferred_port != 0 => { - TcpListener::bind("127.0.0.1:0") - .await - .map_err(|e| format!("Failed to bind stream server: {}", e))? - } - Err(e) => return Err(format!("Failed to bind stream server: {}", e)), - }; - - let actual_addr = listener - .local_addr() - .map_err(|e| format!("Failed to get stream address: {}", e))?; - let port = actual_addr.port(); - - let dashboard_dir = Self::resolve_dashboard_dir(); - - let (frame_tx, _) = broadcast::channel::(64); - let client_count = Arc::new(Mutex::new(0usize)); - let client_notify = Arc::new(Notify::new()); - let screencasting = Arc::new(Mutex::new(false)); - let cdp_session_id = Arc::new(RwLock::new(None::)); - let viewport_width = Arc::new(Mutex::new(1280u32)); - let viewport_height = Arc::new(Mutex::new(720u32)); - let last_tabs = Arc::new(RwLock::new(Vec::::new())); - let last_engine = Arc::new(RwLock::new("chrome".to_string())); - let last_frame = Arc::new(RwLock::new(None::)); - let recording = Arc::new(Mutex::new(false)); - let (shutdown_tx, shutdown_rx) = watch::channel(false); - - let frame_tx_clone = frame_tx.clone(); - let client_count_clone = client_count.clone(); - let client_slot_clone = client_slot.clone(); - let notify_clone = client_notify.clone(); - let screencasting_clone = screencasting.clone(); - let cdp_session_clone = cdp_session_id.clone(); - - let vw_clone = viewport_width.clone(); - let vh_clone = viewport_height.clone(); - let dashboard_dir_clone = dashboard_dir.clone(); - let last_tabs_clone = last_tabs.clone(); - let last_engine_clone = last_engine.clone(); - let last_frame_clone = last_frame.clone(); - let recording_clone = recording.clone(); - let accept_shutdown_rx = shutdown_rx.clone(); - let session_name_clone = session_id.clone(); - let accept_task = tokio::spawn(async move { - accept_loop( - listener, - frame_tx_clone, - client_count_clone, - client_slot_clone, - notify_clone, - screencasting_clone, - cdp_session_clone, - vw_clone, - vh_clone, - dashboard_dir_clone, - last_tabs_clone, - last_engine_clone, - last_frame_clone, - recording_clone, - accept_shutdown_rx, - session_name_clone, - ) - .await; - }); - - // Background CDP event listener for real-time frame broadcasting - let frame_tx_bg = frame_tx.clone(); - let client_slot_bg = client_slot.clone(); - let client_notify_bg = client_notify.clone(); - let screencasting_bg = screencasting.clone(); - let client_count_bg = client_count.clone(); - let cdp_session_bg = cdp_session_id.clone(); - let vw_bg = viewport_width.clone(); - let vh_bg = viewport_height.clone(); - let last_frame_bg = last_frame.clone(); - let last_tabs_bg = last_tabs.clone(); - let last_engine_bg = last_engine.clone(); - let recording_bg = recording.clone(); - let cdp_task = tokio::spawn(async move { - cdp_event_loop( - frame_tx_bg, - client_slot_bg, - client_notify_bg, - screencasting_bg, - client_count_bg, - cdp_session_bg, - vw_bg, - vh_bg, - last_frame_bg, - last_tabs_bg, - last_engine_bg, - recording_bg, - shutdown_rx, - ) - .await; - }); - - Ok(( - Self { - port, - session_name: session_id, - frame_tx, - client_count, - client_slot: client_slot.clone(), - cdp_session_id, - client_notify, - screencasting, - viewport_width, - viewport_height, - dashboard_dir, - last_tabs, - last_engine, - last_frame, - recording, - shutdown_tx, - accept_task: Mutex::new(Some(accept_task)), - cdp_task: Mutex::new(Some(cdp_task)), - }, - client_slot, - )) - } - - pub fn port(&self) -> u16 { - self.port - } - - /// Broadcast a raw frame string (legacy). - pub fn broadcast_frame(&self, frame_json: &str) { - let s = frame_json.to_string(); - if let Ok(mut lf) = self.last_frame.try_write() { - *lf = Some(s.clone()); - } - let _ = self.frame_tx.send(s); - } - - /// Broadcast a screencast frame with structured metadata. - pub fn broadcast_screencast_frame(&self, base64_data: &str, metadata: &FrameMetadata) { - let msg = json!({ - "type": "frame", - "data": base64_data, - "metadata": { - "offsetTop": metadata.offset_top, - "pageScaleFactor": metadata.page_scale_factor, - "deviceWidth": metadata.device_width, - "deviceHeight": metadata.device_height, - "scrollOffsetX": metadata.scroll_offset_x, - "scrollOffsetY": metadata.scroll_offset_y, - "timestamp": metadata.timestamp, - } - }); - let s = msg.to_string(); - if let Ok(mut lf) = self.last_frame.try_write() { - *lf = Some(s.clone()); - } - let _ = self.frame_tx.send(s); - } - - /// Broadcast a status message to all connected clients. - pub async fn broadcast_status( - &self, - connected: bool, - screencasting: bool, - viewport_width: u32, - viewport_height: u32, - engine: &str, - ) { - { - let mut guard = self.last_engine.write().await; - *guard = engine.to_string(); - } - let rec = *self.recording.lock().await; - let msg = json!({ - "type": "status", - "connected": connected, - "screencasting": screencasting, - "viewportWidth": viewport_width, - "viewportHeight": viewport_height, - "engine": engine, - "recording": rec, - }); - let _ = self.frame_tx.send(msg.to_string()); - } - - /// Broadcast an error message to all connected clients. - pub fn broadcast_error(&self, message: &str) { - let msg = json!({ - "type": "error", - "message": message, - }); - let _ = self.frame_tx.send(msg.to_string()); - } - - /// Broadcast a command event when a command begins executing. - pub fn broadcast_command(&self, action: &str, id: &str, params: &Value) { - let msg = json!({ - "type": "command", - "action": action, - "id": id, - "params": params, - "timestamp": timestamp_ms(), - }); - let _ = self.frame_tx.send(msg.to_string()); - } - - /// Broadcast a result event after a command finishes executing. - pub fn broadcast_result( - &self, - id: &str, - action: &str, - success: bool, - data: &Value, - duration_ms: u64, - ) { - let msg = json!({ - "type": "result", - "id": id, - "action": action, - "success": success, - "data": data, - "duration_ms": duration_ms, - "timestamp": timestamp_ms(), - }); - let _ = self.frame_tx.send(msg.to_string()); - } - - /// Broadcast a console event from the browser. - pub fn broadcast_console(&self, level: &str, text: &str, args: &[Value]) { - let mut msg = json!({ - "type": "console", - "level": level, - "text": text, - "timestamp": timestamp_ms(), - }); - if !args.is_empty() { - msg.as_object_mut() - .unwrap() - .insert("args".to_string(), Value::Array(args.to_vec())); - } - let _ = self.frame_tx.send(msg.to_string()); - } - - /// Broadcast a page error (uncaught exception) from the browser. - pub fn broadcast_page_error(&self, text: &str, line: Option, column: Option) { - let msg = json!({ - "type": "page_error", - "text": text, - "line": line, - "column": column, - "timestamp": timestamp_ms(), - }); - let _ = self.frame_tx.send(msg.to_string()); - } - - /// Broadcast the current tab list so the dashboard can render a tab bar. - /// Also caches the list so newly connected WebSocket clients receive it immediately. - pub async fn broadcast_tabs(&self, tabs: &[Value]) { - { - let mut guard = self.last_tabs.write().await; - *guard = tabs.to_vec(); - } - let msg = json!({ - "type": "tabs", - "tabs": tabs, - "timestamp": timestamp_ms(), - }); - let _ = self.frame_tx.send(msg.to_string()); - } - - /// Whether the dashboard directory is available. - pub fn has_dashboard(&self) -> bool { - self.dashboard_dir.is_some() - } -} - -#[allow(clippy::too_many_arguments)] -async fn accept_loop( - listener: TcpListener, - frame_tx: broadcast::Sender, - client_count: Arc>, - client_slot: Arc>>>, - client_notify: Arc, - screencasting: Arc>, - cdp_session_id: Arc>>, - viewport_width: Arc>, - viewport_height: Arc>, - dashboard_dir: Option, - last_tabs: Arc>>, - last_engine: Arc>, - last_frame: Arc>>, - recording: Arc>, - mut shutdown_rx: watch::Receiver, - session_name: String, -) { - let dashboard_dir = dashboard_dir.map(Arc::from); - let session_name: Arc = Arc::from(session_name); - loop { - tokio::select! { - changed = shutdown_rx.changed() => { - if changed.is_err() || *shutdown_rx.borrow() { - break; - } - } - accept_result = listener.accept() => { - let Ok((stream, addr)) = accept_result else { - break; - }; - let frame_tx = frame_tx.clone(); - let client_count = client_count.clone(); - let client_slot = client_slot.clone(); - let client_notify = client_notify.clone(); - let screencasting = screencasting.clone(); - let cdp_session_id = cdp_session_id.clone(); - let vw = viewport_width.clone(); - let vh = viewport_height.clone(); - let dd = dashboard_dir.clone(); - let lt = last_tabs.clone(); - let le = last_engine.clone(); - let lf = last_frame.clone(); - let rec = recording.clone(); - let shutdown_rx = shutdown_rx.clone(); - let sn = session_name.clone(); - - tokio::spawn(async move { - handle_connection( - stream, - addr, - frame_tx, - client_count, - client_slot, - client_notify, - screencasting, - cdp_session_id, - vw, - vh, - dd, - lt, - le, - lf, - rec, - shutdown_rx, - sn, - ) - .await; - }); - } - } - } -} - -fn is_websocket_upgrade(request: &str) -> bool { - request.lines().any(|line| { - if let Some((name, value)) = line.split_once(':') { - name.trim().eq_ignore_ascii_case("upgrade") - && value.trim().eq_ignore_ascii_case("websocket") - } else { - false - } - }) -} - -/// Peek at the TCP stream to dispatch between WebSocket upgrade and plain HTTP. -#[allow(clippy::too_many_arguments)] -async fn handle_connection( - stream: tokio::net::TcpStream, - addr: SocketAddr, - frame_tx: broadcast::Sender, - client_count: Arc>, - client_slot: Arc>>>, - client_notify: Arc, - screencasting: Arc>, - cdp_session_id: Arc>>, - viewport_width: Arc>, - viewport_height: Arc>, - dashboard_dir: Option>, - last_tabs: Arc>>, - last_engine: Arc>, - last_frame: Arc>>, - recording: Arc>, - shutdown_rx: watch::Receiver, - session_name: Arc, -) { - let mut buf = [0u8; 4096]; - let n = match stream.peek(&mut buf).await { - Ok(n) => n, - Err(_) => return, - }; - let request = String::from_utf8_lossy(&buf[..n]); - - if is_websocket_upgrade(&request) { - let frame_rx = frame_tx.subscribe(); - handle_ws_client( - stream, - addr, - frame_rx, - client_count, - client_slot, - client_notify, - screencasting, - cdp_session_id, - viewport_width, - viewport_height, - last_tabs, - last_engine, - last_frame, - recording, - shutdown_rx, - ) - .await; - } else { - handle_http_request( - stream, - &request, - n, - dashboard_dir.as_deref().map(|p| p.as_path()), - &last_tabs, - &last_engine, - &session_name, - ) - .await; - } -} - -#[allow(clippy::result_large_err, clippy::too_many_arguments)] -async fn handle_ws_client( - stream: tokio::net::TcpStream, - _addr: SocketAddr, - mut frame_rx: broadcast::Receiver, - client_count: Arc>, - client_slot: Arc>>>, - client_notify: Arc, - screencasting: Arc>, - cdp_session_id: Arc>>, - viewport_width: Arc>, - viewport_height: Arc>, - last_tabs: Arc>>, - last_engine: Arc>, - last_frame: Arc>>, - recording: Arc>, - mut shutdown_rx: watch::Receiver, -) { - let callback = - |req: &tokio_tungstenite::tungstenite::handshake::server::Request, - resp: tokio_tungstenite::tungstenite::handshake::server::Response| { - let origin = req - .headers() - .get("origin") - .and_then(|v| v.to_str().ok()) - .map(|s| s.to_string()); - if !is_allowed_origin(origin.as_deref()) { - let mut reject = - tokio_tungstenite::tungstenite::handshake::server::ErrorResponse::new(Some( - "Origin not allowed".to_string(), - )); - *reject.status_mut() = tokio_tungstenite::tungstenite::http::StatusCode::FORBIDDEN; - return Err(reject); - } - Ok(resp) - }; - - let ws_stream = match tokio_tungstenite::accept_hdr_async(stream, callback).await { - Ok(ws) => ws, - Err(_) => return, - }; - - { - let mut count = client_count.lock().await; - *count += 1; - } - - let (mut ws_tx, mut ws_rx) = ws_stream.split(); - - // Send initial status with current viewport dimensions - { - let guard = client_slot.read().await; - let connected = guard.is_some(); - let sc = *screencasting.lock().await; - let vw = *viewport_width.lock().await; - let vh = *viewport_height.lock().await; - let eng = last_engine.read().await.clone(); - let rec = *recording.lock().await; - let status = json!({ - "type": "status", - "connected": connected, - "screencasting": sc, - "viewportWidth": vw, - "viewportHeight": vh, - "engine": eng, - "recording": rec, - }); - let _ = ws_tx.send(Message::Text(status.to_string())).await; - - let tabs = last_tabs.read().await; - if !tabs.is_empty() { - let tabs_msg = json!({ - "type": "tabs", - "tabs": *tabs, - "timestamp": timestamp_ms(), - }); - let _ = ws_tx.send(Message::Text(tabs_msg.to_string())).await; - } - - // Send the most recent screencast frame so new clients see content immediately - if let Some(ref cached) = *last_frame.read().await { - let _ = ws_tx.send(Message::Text(cached.clone())).await; - } - } - - // Notify the CDP event loop that a client connected (may trigger auto-start screencast) - client_notify.notify_one(); - - loop { - tokio::select! { - changed = shutdown_rx.changed() => { - if changed.is_err() || *shutdown_rx.borrow() { - let _ = ws_tx.send(Message::Close(None)).await; - break; - } - } - frame = frame_rx.recv() => { - match frame { - Ok(data) => { - if ws_tx.send(Message::Text(data)).await.is_err() { - break; - } - } - Err(broadcast::error::RecvError::Lagged(_)) => { - // Slow consumer; skip missed frames and continue - continue; - } - Err(broadcast::error::RecvError::Closed) => break, - } - } - msg = ws_rx.next() => { - match msg { - Some(Ok(Message::Text(text))) => { - let guard = client_slot.read().await; - if let Some(ref client) = *guard { - let sid = cdp_session_id.read().await; - handle_client_message(&text, client.as_ref(), sid.as_deref()).await; - } - } - Some(Ok(Message::Close(_))) | None => break, - _ => {} - } - } - } - } - - { - let mut count = client_count.lock().await; - *count = count.saturating_sub(1); - } - - // Notify the CDP event loop that a client disconnected (may trigger auto-stop screencast) - client_notify.notify_one(); -} - -/// Background task that subscribes to CDP events and broadcasts screencast frames in real-time. -/// Also handles auto-start/stop of screencast based on WebSocket client count. -#[allow(clippy::too_many_arguments)] -async fn cdp_event_loop( - frame_tx: broadcast::Sender, - client_slot: Arc>>>, - client_notify: Arc, - screencasting: Arc>, - client_count: Arc>, - cdp_session_id: Arc>>, - viewport_width: Arc>, - viewport_height: Arc>, - last_frame: Arc>>, - last_tabs: Arc>>, - last_engine: Arc>, - recording: Arc>, - mut shutdown_rx: watch::Receiver, -) { - loop { - // Wait until we're notified of a client/connection change - tokio::select! { - changed = shutdown_rx.changed() => { - if changed.is_err() || *shutdown_rx.borrow() { - let session_id = cdp_session_id.read().await.clone(); - if *screencasting.lock().await { - if let Some(ref client) = *client_slot.read().await { - let _ = client - .send_command_no_params("Page.stopScreencast", session_id.as_deref()) - .await; - } - let mut sc = screencasting.lock().await; - *sc = false; - } - return; - } - } - _ = client_notify.notified() => {} - } - - // Check if we have WS clients and a CDP client - let count = *client_count.lock().await; - let guard = client_slot.read().await; - - if count > 0 { - if let Some(ref client) = *guard { - // We have WS clients and a CDP client — start screencast and listen for frames - let mut event_rx = client.subscribe(); - let client_arc = Arc::clone(client); - drop(guard); - - // Get the CDP page session ID for targeted commands - let session_id = cdp_session_id.read().await.clone(); - - // Use the current viewport dimensions for screencast - let vw = *viewport_width.lock().await; - let vh = *viewport_height.lock().await; - - let eng = last_engine.read().await.clone(); - let supports_screencast = eng == "chrome"; - - if supports_screencast { - let _ = client_arc - .send_command( - "Page.startScreencast", - Some(json!({ - "format": "jpeg", - "quality": 80, - "maxWidth": vw, - "maxHeight": vh, - "everyNthFrame": 1, - })), - session_id.as_deref(), - ) - .await; - } - - { - let mut sc = screencasting.lock().await; - *sc = supports_screencast; - } - - // Broadcast connection status with current viewport - let rec = *recording.lock().await; - let status = json!({ - "type": "status", - "connected": true, - "screencasting": supports_screencast, - "viewportWidth": vw, - "viewportHeight": vh, - "engine": eng, - "recording": rec, - }); - let _ = frame_tx.send(status.to_string()); - - // Process CDP events in real-time until client disconnects or CDP closes - loop { - tokio::select! { - changed = shutdown_rx.changed() => { - if changed.is_err() || *shutdown_rx.borrow() { - if supports_screencast { - let session_id = cdp_session_id.read().await.clone(); - let _ = client_arc - .send_command_no_params("Page.stopScreencast", session_id.as_deref()) - .await; - } - let mut sc = screencasting.lock().await; - *sc = false; - return; - } - } - event = event_rx.recv() => { - match event { - Ok(evt) => { - if evt.method == "Page.frameNavigated" { - if let Some(frame) = evt.params.get("frame") { - let is_main = frame - .get("parentId") - .and_then(|v| v.as_str()) - .is_none_or(|s| s.is_empty()); - if is_main { - if let Some(url) = frame.get("url").and_then(|v| v.as_str()) { - // Update the cached tab list so the active tab URL is current - { - let mut tabs = last_tabs.write().await; - for tab in tabs.iter_mut() { - if tab.get("active").and_then(|v| v.as_bool()).unwrap_or(false) { - tab.as_object_mut().map(|o| o.insert("url".to_string(), json!(url))); - } - } - } - let msg = json!({ - "type": "url", - "url": url, - "timestamp": timestamp_ms(), - }); - let _ = frame_tx.send(msg.to_string()); - } - } - } - } else if evt.method == "Page.screencastFrame" { - if let Some(sid) = evt.params.get("sessionId").and_then(|v| v.as_i64()) { - let _ = client_arc.send_command( - "Page.screencastFrameAck", - Some(json!({ "sessionId": sid })), - evt.session_id.as_deref(), - ).await; - } - - if let Some(data) = evt.params.get("data").and_then(|v| v.as_str()) { - let meta = evt.params.get("metadata"); - let msg = json!({ - "type": "frame", - "data": data, - "metadata": { - "offsetTop": meta.and_then(|m| m.get("offsetTop")).and_then(|v| v.as_f64()).unwrap_or(0.0), - "pageScaleFactor": meta.and_then(|m| m.get("pageScaleFactor")).and_then(|v| v.as_f64()).unwrap_or(1.0), - "deviceWidth": meta.and_then(|m| m.get("deviceWidth")).and_then(|v| v.as_u64()).unwrap_or(1280), - "deviceHeight": meta.and_then(|m| m.get("deviceHeight")).and_then(|v| v.as_u64()).unwrap_or(720), - "scrollOffsetX": meta.and_then(|m| m.get("scrollOffsetX")).and_then(|v| v.as_f64()).unwrap_or(0.0), - "scrollOffsetY": meta.and_then(|m| m.get("scrollOffsetY")).and_then(|v| v.as_f64()).unwrap_or(0.0), - "timestamp": meta.and_then(|m| m.get("timestamp")).and_then(|v| v.as_u64()).unwrap_or(0), - } - }); - let msg_str = msg.to_string(); - { - let mut lf = last_frame.write().await; - *lf = Some(msg_str.clone()); - } - let _ = frame_tx.send(msg_str); - } - } else if evt.method == "Runtime.consoleAPICalled" { - let level = evt.params.get("type") - .and_then(|v| v.as_str()) - .unwrap_or("log"); - let raw_args = evt.params.get("args") - .and_then(|v| v.as_array()) - .cloned() - .unwrap_or_default(); - let text = network::format_console_args(&raw_args); - if !text.is_empty() { - let mut msg = json!({ - "type": "console", - "level": level, - "text": text, - "timestamp": timestamp_ms(), - }); - if !raw_args.is_empty() { - msg.as_object_mut().unwrap().insert( - "args".to_string(), - Value::Array(raw_args), - ); - } - let _ = frame_tx.send(msg.to_string()); - } - } else if evt.method == "Runtime.exceptionThrown" { - let text = evt.params.get("exceptionDetails") - .and_then(|d| { - d.get("exception") - .and_then(|e| e.get("description").and_then(|v| v.as_str())) - .or_else(|| d.get("text").and_then(|v| v.as_str())) - }) - .unwrap_or("Unknown error"); - let line = evt.params.get("exceptionDetails") - .and_then(|d| d.get("lineNumber").and_then(|v| v.as_i64())); - let column = evt.params.get("exceptionDetails") - .and_then(|d| d.get("columnNumber").and_then(|v| v.as_i64())); - let msg = json!({ - "type": "page_error", - "text": text, - "line": line, - "column": column, - "timestamp": timestamp_ms(), - }); - let _ = frame_tx.send(msg.to_string()); - } - } - Err(broadcast::error::RecvError::Lagged(_)) => continue, - Err(broadcast::error::RecvError::Closed) => break, - } - } - // Also check for notify (client count change, CDP client change, session switch, or viewport change) - _ = client_notify.notified() => { - let count = *client_count.lock().await; - let new_session_id = cdp_session_id.read().await.clone(); - if count == 0 { - if supports_screencast { - let _ = client_arc - .send_command_no_params("Page.stopScreencast", session_id.as_deref()) - .await; - } - let mut sc = screencasting.lock().await; - *sc = false; - break; - } - let client_changed = { - let guard = client_slot.read().await; - let same = guard - .as_ref() - .is_some_and(|c| Arc::ptr_eq(c, &client_arc)); - !same - }; - let session_changed = new_session_id != session_id; - let new_vw = *viewport_width.lock().await; - let new_vh = *viewport_height.lock().await; - let viewport_changed = new_vw != vw || new_vh != vh; - if client_changed || session_changed || viewport_changed { - if supports_screencast { - let _ = client_arc - .send_command_no_params("Page.stopScreencast", session_id.as_deref()) - .await; - } - let mut sc = screencasting.lock().await; - *sc = false; - client_notify.notify_one(); - break; - } - } - } - } - } else { - drop(guard); - // No CDP client yet — wait for next notification - } - } else { - // No WS clients — if screencasting, stop it - let was_screencasting = *screencasting.lock().await; - if was_screencasting { - if let Some(ref client) = *guard { - let session_id = cdp_session_id.read().await.clone(); - let _ = client - .send_command_no_params("Page.stopScreencast", session_id.as_deref()) - .await; - } - let mut sc = screencasting.lock().await; - *sc = false; - } - drop(guard); - } - } -} - -async fn handle_client_message(msg: &str, client: &CdpClient, session_id: Option<&str>) { - let parsed: Value = match serde_json::from_str(msg) { - Ok(v) => v, - Err(_) => return, - }; - - let msg_type = parsed.get("type").and_then(|v| v.as_str()).unwrap_or(""); - - match msg_type { - "input_mouse" => { - let _ = client - .send_command( - "Input.dispatchMouseEvent", - Some(json!({ - "type": parsed.get("eventType").and_then(|v| v.as_str()).unwrap_or("mouseMoved"), - "x": parsed.get("x").and_then(|v| v.as_f64()).unwrap_or(0.0), - "y": parsed.get("y").and_then(|v| v.as_f64()).unwrap_or(0.0), - "button": parsed.get("button").and_then(|v| v.as_str()).unwrap_or("none"), - "clickCount": parsed.get("clickCount").and_then(|v| v.as_i64()).unwrap_or(0), - "deltaX": parsed.get("deltaX").and_then(|v| v.as_f64()).unwrap_or(0.0), - "deltaY": parsed.get("deltaY").and_then(|v| v.as_f64()).unwrap_or(0.0), - "modifiers": parsed.get("modifiers").and_then(|v| v.as_i64()).unwrap_or(0), - })), - session_id, - ) - .await; - } - "input_keyboard" => { - let _ = client - .send_command( - "Input.dispatchKeyEvent", - Some(json!({ - "type": parsed.get("eventType").and_then(|v| v.as_str()).unwrap_or("keyDown"), - "key": parsed.get("key"), - "code": parsed.get("code"), - "text": parsed.get("text"), - "windowsVirtualKeyCode": parsed.get("windowsVirtualKeyCode").and_then(|v| v.as_i64()).unwrap_or(0), - "modifiers": parsed.get("modifiers").and_then(|v| v.as_i64()).unwrap_or(0), - })), - session_id, - ) - .await; - } - "input_touch" => { - let _ = client - .send_command( - "Input.dispatchTouchEvent", - Some(json!({ - "type": parsed.get("eventType").and_then(|v| v.as_str()).unwrap_or("touchStart"), - "touchPoints": parsed.get("touchPoints").unwrap_or(&json!([])), - "modifiers": parsed.get("modifiers").and_then(|v| v.as_i64()).unwrap_or(0), - })), - session_id, - ) - .await; - } - "status" => { - // Client requesting status -- handled via broadcast_status from the caller - } - _ => {} - } -} - -const CORS_HEADERS: &str = "Access-Control-Allow-Origin: *\r\nAccess-Control-Allow-Methods: GET, POST, OPTIONS\r\nAccess-Control-Allow-Headers: Content-Type\r\n"; - -/// Serve an HTTP request for dashboard static files or the fallback page. -async fn handle_http_request( - mut stream: tokio::net::TcpStream, - request: &str, - peeked_len: usize, - dashboard_dir: Option<&Path>, - last_tabs: &Arc>>, - last_engine: &Arc>, - session_name: &str, -) { - let mut discard = vec![0u8; peeked_len]; - let _ = stream.read_exact(&mut discard).await; - - let first_line = request.lines().next().unwrap_or(""); - let method = first_line.split_whitespace().next().unwrap_or("GET"); - let path = first_line.split_whitespace().nth(1).unwrap_or("/"); - - // Handle CORS preflight - if method == "OPTIONS" { - let response = format!( - "HTTP/1.1 204 No Content\r\n{CORS_HEADERS}Access-Control-Max-Age: 86400\r\nContent-Length: 0\r\nConnection: close\r\n\r\n" - ); - let _ = stream.write_all(response.as_bytes()).await; - return; - } - - // Handle POST /api/sessions (spawn new session) - if method == "POST" && path == "/api/sessions" { - let body_str = extract_http_body(request).unwrap_or(""); - let result = spawn_session(body_str).await; - let (status, resp_body) = match result { - Ok(msg) => ("200 OK", msg), - Err(e) => ( - "400 Bad Request", - format!( - r#"{{"success":false,"error":{}}}"#, - serde_json::to_string(&e).unwrap_or_else(|_| format!("\"{}\"", e)) - ), - ), - }; - let response = format!( - "HTTP/1.1 {status}\r\nContent-Type: application/json; charset=utf-8\r\nContent-Length: {}\r\nConnection: close\r\n{CORS_HEADERS}\r\n", - resp_body.len() - ); - let _ = stream.write_all(response.as_bytes()).await; - let _ = stream.write_all(resp_body.as_bytes()).await; - return; - } - - // Handle POST /api/command - if method == "POST" && path == "/api/command" { - let body = extract_http_body(request).unwrap_or(""); - let result = relay_command_to_daemon(session_name, body).await; - let (status, resp_body) = match result { - Ok(resp) => ("200 OK", resp), - Err(e) => ( - "502 Bad Gateway", - format!( - r#"{{"success":false,"error":{}}}"#, - serde_json::to_string(&e).unwrap_or_else(|_| format!("\"{}\"", e)) - ), - ), - }; - let response = format!( - "HTTP/1.1 {status}\r\nContent-Type: application/json; charset=utf-8\r\nContent-Length: {}\r\nConnection: close\r\n{CORS_HEADERS}\r\n", - resp_body.len() - ); - let _ = stream.write_all(response.as_bytes()).await; - let _ = stream.write_all(resp_body.as_bytes()).await; - return; - } - - let (status, content_type, body): (&str, &str, Vec) = if path == "/api/sessions" { - ( - "200 OK", - "application/json; charset=utf-8", - discover_sessions().into_bytes(), - ) - } else if path == "/api/tabs" { - let tabs = last_tabs.read().await; - ( - "200 OK", - "application/json; charset=utf-8", - serde_json::to_string(&*tabs) - .unwrap_or_else(|_| "[]".to_string()) - .into_bytes(), - ) - } else if path == "/api/status" { - let engine = last_engine.read().await; - ( - "200 OK", - "application/json; charset=utf-8", - format!(r#"{{"engine":"{}"}}"#, *engine).into_bytes(), - ) - } else { - match dashboard_dir { - Some(dir) => serve_static_file(dir, path), - None => ( - "200 OK", - "text/html; charset=utf-8", - DASHBOARD_NOT_INSTALLED_HTML.as_bytes().to_vec(), - ), - } - }; - - let response = format!( - "HTTP/1.1 {}\r\nContent-Type: {}\r\nContent-Length: {}\r\nConnection: close\r\n{CORS_HEADERS}\r\n", - status, - content_type, - body.len() - ); - let _ = stream.write_all(response.as_bytes()).await; - let _ = stream.write_all(&body).await; -} - -/// Extract the HTTP body from a raw request string (headers + body in one buffer). -fn extract_http_body(request: &str) -> Option<&str> { - // Body starts after the first blank line (\r\n\r\n) - request - .find("\r\n\r\n") - .map(|pos| &request[pos + 4..]) - .or_else(|| request.find("\n\n").map(|pos| &request[pos + 2..])) -} - -/// Relay a command JSON body to the daemon and return the response. -async fn relay_command_to_daemon(session_name: &str, body: &str) -> Result { - let mut cmd: Value = serde_json::from_str(body).map_err(|e| format!("Invalid JSON: {}", e))?; - - if cmd.get("id").is_none() { - let id = format!( - "dash-{}", - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap_or_default() - .as_millis() - ); - cmd["id"] = json!(id); - } - - let mut json_str = serde_json::to_string(&cmd).map_err(|e| e.to_string())?; - json_str.push('\n'); - - #[cfg(unix)] - let stream = { - let socket_path = get_socket_dir().join(format!("{}.sock", session_name)); - tokio::net::UnixStream::connect(&socket_path) - .await - .map_err(|e| format!("Failed to connect to daemon: {}", e))? - }; - - #[cfg(windows)] - let stream = { - let port = resolve_port(session_name); - tokio::net::TcpStream::connect(format!("127.0.0.1:{}", port)) - .await - .map_err(|e| format!("Failed to connect to daemon: {}", e))? - }; - - let (reader, mut writer) = tokio::io::split(stream); - - writer - .write_all(json_str.as_bytes()) - .await - .map_err(|e| format!("Failed to send command: {}", e))?; - - let mut buf_reader = tokio::io::BufReader::new(reader); - let mut response_line = String::new(); - buf_reader - .read_line(&mut response_line) - .await - .map_err(|e| format!("Failed to read response: {}", e))?; - - Ok(response_line.trim().to_string()) -} - -fn serve_static_file(dir: &Path, url_path: &str) -> (&'static str, &'static str, Vec) { - let clean = url_path.trim_start_matches('/'); - let file_path = if clean.is_empty() { - dir.join("index.html") - } else { - let joined = dir.join(clean); - if joined.is_file() { - joined - } else { - dir.join("index.html") - } - }; - - match std::fs::read(&file_path) { - Ok(content) => { - let ext = file_path.extension().and_then(|e| e.to_str()).unwrap_or(""); - let ct = match ext { - "html" => "text/html; charset=utf-8", - "js" => "application/javascript; charset=utf-8", - "css" => "text/css; charset=utf-8", - "json" => "application/json; charset=utf-8", - "svg" => "image/svg+xml", - "png" => "image/png", - "ico" => "image/x-icon", - _ => "application/octet-stream", - }; - ("200 OK", ct, content) - } - Err(_) => ( - "404 Not Found", - "text/html; charset=utf-8", - b"

404 Not Found

".to_vec(), - ), - } -} - -const DASHBOARD_NOT_INSTALLED_HTML: &str = r#" - -agent-browser - - - -
-

Dashboard not installed

-

Run agent-browser dashboard install to download the dashboard.

-
- -"#; - -/// Discover all active streaming sessions by reading `*.stream` files. -/// Stale entries (dead process) are removed on the fly. -fn discover_sessions() -> String { - let dir = get_socket_dir(); - let mut sessions = Vec::new(); - - if let Ok(entries) = std::fs::read_dir(&dir) { - for entry in entries.flatten() { - let name = entry.file_name(); - let name_str = name.to_string_lossy(); - if let Some(session) = name_str.strip_suffix(".stream") { - if let Ok(port_str) = std::fs::read_to_string(entry.path()) { - if let Ok(port) = port_str.trim().parse::() { - let pid_path = dir.join(format!("{}.pid", session)); - if is_process_alive(&pid_path) { - let engine_path = dir.join(format!("{}.engine", session)); - let engine = std::fs::read_to_string(&engine_path) - .ok() - .filter(|s| !s.trim().is_empty()) - .unwrap_or_else(|| "chrome".to_string()); - - let provider_path = dir.join(format!("{}.provider", session)); - let provider = std::fs::read_to_string(&provider_path) - .ok() - .filter(|s| !s.trim().is_empty()); - - let extensions = read_extensions_metadata(&dir, session); - - let mut entry = json!({ - "session": session, - "port": port, - "engine": engine.trim(), - }); - if let Some(ref p) = provider { - entry["provider"] = json!(p.trim()); - } - if !extensions.is_empty() { - entry["extensions"] = json!(extensions); - } - sessions.push(entry); - } else { - let _ = std::fs::remove_file(entry.path()); - } - } - } - } - } - } - - serde_json::to_string(&sessions).unwrap_or_else(|_| "[]".to_string()) -} - -fn read_extensions_metadata(dir: &std::path::Path, session: &str) -> Vec { - let ext_path = dir.join(format!("{}.extensions", session)); - let ext_str = match std::fs::read_to_string(&ext_path) { - Ok(s) => s, - Err(_) => return Vec::new(), - }; - - ext_str - .split(',') - .map(|p| p.trim()) - .filter(|p| !p.is_empty()) - .filter_map(|path| { - let manifest_path = std::path::Path::new(path).join("manifest.json"); - let manifest_str = std::fs::read_to_string(&manifest_path).ok()?; - let manifest: Value = serde_json::from_str(&manifest_str).ok()?; - - let name = manifest - .get("name") - .and_then(|v| v.as_str()) - .unwrap_or("Unknown") - .to_string(); - let version = manifest - .get("version") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let description = manifest - .get("description") - .and_then(|v| v.as_str()) - .map(|s| s.to_string()); - - let mut ext = json!({ - "name": name, - "version": version, - "path": path, - }); - if let Some(desc) = description { - ext["description"] = json!(desc); - } - Some(ext) - }) - .collect() -} - -fn is_process_alive(pid_path: &Path) -> bool { - let pid_str = match std::fs::read_to_string(pid_path) { - Ok(s) => s, - Err(_) => return false, - }; - let pid: u32 = match pid_str.trim().parse() { - Ok(p) => p, - Err(_) => return false, - }; - #[cfg(unix)] - { - unsafe { libc::kill(pid as i32, 0) == 0 } - } - #[cfg(not(unix))] - { - let _ = pid; - // On non-Unix, just check if the pid file exists - true - } -} - -fn timestamp_ms() -> u64 { - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .map(|d| d.as_millis() as u64) - .unwrap_or(0) -} - -pub fn is_allowed_origin(origin: Option<&str>) -> bool { - match origin { - None => true, - Some(o) => { - if o.starts_with("file://") { - return true; - } - if let Ok(url) = url::Url::parse(o) { - let host = url.host_str().unwrap_or(""); - host == "localhost" || host == "127.0.0.1" || host == "::1" || host == "[::1]" - } else { - false - } - } - } -} - -pub async fn start_screencast( - client: &CdpClient, - session_id: &str, - format: &str, - quality: i32, - max_width: i32, - max_height: i32, -) -> Result<(), String> { - client - .send_command( - "Page.startScreencast", - Some(json!({ - "format": format, - "quality": quality, - "maxWidth": max_width, - "maxHeight": max_height, - "everyNthFrame": 1, - })), - Some(session_id), - ) - .await?; - Ok(()) -} - -pub async fn stop_screencast(client: &CdpClient, session_id: &str) -> Result<(), String> { - client - .send_command_no_params("Page.stopScreencast", Some(session_id)) - .await?; - Ok(()) -} - -pub async fn ack_screencast_frame( - client: &CdpClient, - session_id: &str, - screencast_session_id: i64, -) -> Result<(), String> { - client - .send_command( - "Page.screencastFrameAck", - Some(json!({ "sessionId": screencast_session_id })), - Some(session_id), - ) - .await?; - Ok(()) -} - -/// Standalone dashboard HTTP server (no browser, no WebSocket streaming). -/// Serves static files and `/api/sessions` for session discovery. -pub async fn run_dashboard_server(port: u16) { - let addr = format!("127.0.0.1:{}", port); - let listener = match TcpListener::bind(&addr).await { - Ok(l) => l, - Err(e) => { - eprintln!("Failed to bind dashboard server on {}: {}", addr, e); - return; - } - }; - - let dashboard_dir: Arc = Arc::from(get_dashboard_dir()); - - loop { - let Ok((stream, _addr)) = listener.accept().await else { - break; - }; - let dash_dir = dashboard_dir.clone(); - tokio::spawn(async move { - handle_dashboard_connection(stream, dash_dir).await; - }); - } -} - -async fn handle_dashboard_connection( - mut stream: tokio::net::TcpStream, - dashboard_dir: Arc, -) { - use tokio::io::AsyncReadExt; - - let mut buf = vec![0u8; 8192]; - let n = match stream.read(&mut buf).await { - Ok(n) if n > 0 => n, - _ => return, - }; - - let first_line = std::str::from_utf8(&buf[..n]) - .unwrap_or("") - .lines() - .next() - .unwrap_or("") - .to_string(); - let method = first_line.split_whitespace().next().unwrap_or("GET"); - let path = first_line.split_whitespace().nth(1).unwrap_or("/"); - - if method == "OPTIONS" { - let response = format!( - "HTTP/1.1 204 No Content\r\n{CORS_HEADERS}Access-Control-Max-Age: 86400\r\nContent-Length: 0\r\nConnection: close\r\n\r\n" - ); - let _ = stream.write_all(response.as_bytes()).await; - return; - } - - if method == "POST" && (path == "/api/sessions" || path == "/api/exec" || path == "/api/kill") { - let body_str = read_post_body(&mut stream, &buf, n).await; - let result = if path == "/api/exec" { - exec_cli(&body_str).await - } else if path == "/api/kill" { - kill_session(&body_str).await - } else { - spawn_session(&body_str).await - }; - let (status, resp_body) = match result { - Ok(msg) => ("200 OK", msg), - Err(e) => ( - "400 Bad Request", - format!( - r#"{{"success":false,"error":{}}}"#, - serde_json::to_string(&e).unwrap_or_else(|_| format!("\"{}\"", e)) - ), - ), - }; - let response = format!( - "HTTP/1.1 {status}\r\nContent-Type: application/json; charset=utf-8\r\nContent-Length: {}\r\nConnection: close\r\n{CORS_HEADERS}\r\n", - resp_body.len() - ); - let _ = stream.write_all(response.as_bytes()).await; - let _ = stream.write_all(resp_body.as_bytes()).await; - return; - } - - let (status, content_type, body): (&str, &str, Vec) = if path == "/api/sessions" { - ( - "200 OK", - "application/json; charset=utf-8", - discover_sessions().into_bytes(), - ) - } else if dashboard_dir.join("index.html").exists() { - serve_static_file(&dashboard_dir, path) - } else { - ( - "200 OK", - "text/html; charset=utf-8", - DASHBOARD_NOT_INSTALLED_HTML.as_bytes().to_vec(), - ) - }; - - let response = format!( - "HTTP/1.1 {}\r\nContent-Type: {}\r\nContent-Length: {}\r\nConnection: close\r\n{CORS_HEADERS}\r\n", - status, - content_type, - body.len() - ); - let _ = stream.write_all(response.as_bytes()).await; - let _ = stream.write_all(&body).await; -} - -/// Read the full POST body from a request. First checks if the body is already -/// present in the initial read buffer; if not, reads remaining bytes based on -/// Content-Length. -async fn read_post_body(stream: &mut tokio::net::TcpStream, initial: &[u8], n: usize) -> String { - use tokio::io::AsyncReadExt; - let header_str = String::from_utf8_lossy(&initial[..n]); - let body = extract_http_body(&header_str).unwrap_or("").to_string(); - - if !body.is_empty() { - return body; - } - - let cl = header_str - .lines() - .find_map(|l| { - let lower = l.to_lowercase(); - lower - .strip_prefix("content-length:") - .map(|v| v.trim().parse::().unwrap_or(0)) - }) - .unwrap_or(0); - - if cl > 0 { - let mut remaining = vec![0u8; cl]; - if stream.read_exact(&mut remaining).await.is_ok() { - return String::from_utf8_lossy(&remaining).to_string(); - } - } - - String::new() -} - -/// Execute an agent-browser CLI command and return JSON with stdout/stderr. -async fn exec_cli(body: &str) -> Result { - let parsed: Value = serde_json::from_str(body).map_err(|e| format!("Invalid JSON: {}", e))?; - let args: Vec = parsed - .get("args") - .and_then(|v| v.as_array()) - .ok_or("Missing \"args\" array")? - .iter() - .filter_map(|v| v.as_str().map(|s| s.to_string())) - .collect(); - - if args.is_empty() { - return Err("Empty args array".to_string()); - } - - let exe = std::env::current_exe().map_err(|e| format!("Cannot resolve executable: {}", e))?; - - let mut cmd = tokio::process::Command::new(&exe); - cmd.args(&args) - .arg("--json") - .env_remove("AGENT_BROWSER_DASHBOARD") - .env_remove("AGENT_BROWSER_DASHBOARD_PORT") - .env_remove("AGENT_BROWSER_STREAM_PORT"); - - let output = cmd - .output() - .await - .map_err(|e| format!("Failed to execute: {}", e))?; - - let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string(); - let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string(); - - Ok(json!({ - "success": output.status.success(), - "exit_code": output.status.code(), - "stdout": stdout, - "stderr": stderr, - }) - .to_string()) -} - -/// Kill a session daemon by sending SIGTERM, then SIGKILL if it survives. -/// Cleans up socket/pid/stream/engine files afterward. -async fn kill_session(body: &str) -> Result { - let parsed: Value = serde_json::from_str(body).map_err(|e| format!("Invalid JSON: {}", e))?; - let session = parsed - .get("session") - .and_then(|v| v.as_str()) - .ok_or("Missing \"session\" field")?; - - if session.is_empty() || session.len() > 64 { - return Err("Session name must be 1-64 characters".to_string()); - } - - let dir = get_socket_dir(); - let pid_path = dir.join(format!("{}.pid", session)); - - let pid_str = std::fs::read_to_string(&pid_path) - .map_err(|_| format!("No PID file for session '{}'", session))?; - let pid: u32 = pid_str - .trim() - .parse() - .map_err(|_| format!("Invalid PID in file: {}", pid_str.trim()))?; - - #[cfg(unix)] - { - unsafe { - libc::kill(pid as i32, libc::SIGTERM); - } - tokio::time::sleep(std::time::Duration::from_millis(500)).await; - if unsafe { libc::kill(pid as i32, 0) } == 0 { - unsafe { - libc::kill(pid as i32, libc::SIGKILL); - } - } - } - - for ext in &["pid", "sock", "stream", "engine", "extensions"] { - let _ = std::fs::remove_file(dir.join(format!("{}.{}", session, ext))); - } - - Ok(json!({ "success": true, "killed_pid": pid }).to_string()) -} - -/// Spawn a new session daemon from a POST /api/sessions request. -async fn spawn_session(body: &str) -> Result { - let parsed: Value = serde_json::from_str(body).map_err(|e| format!("Invalid JSON: {}", e))?; - let session = parsed - .get("session") - .and_then(|v| v.as_str()) - .ok_or("Missing \"session\" field")?; - - if session.is_empty() || session.len() > 64 { - return Err("Session name must be 1-64 characters".to_string()); - } - - let exe = std::env::current_exe().map_err(|e| format!("Cannot resolve executable: {}", e))?; - - let mut cmd = tokio::process::Command::new(&exe); - cmd.arg("open") - .arg("about:blank") - .arg("--session") - .arg(session); - - cmd.stdout(std::process::Stdio::null()); - cmd.stderr(std::process::Stdio::null()); - - let status = cmd - .status() - .await - .map_err(|e| format!("Failed to spawn session: {}", e))?; - - if status.success() { - Ok(format!( - r#"{{"success":true,"session":{}}}"#, - serde_json::to_string(session).unwrap_or_default() - )) - } else { - Err(format!("Session process exited with {}", status)) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_allowed_origin_none() { - assert!(is_allowed_origin(None)); - } - - #[test] - fn test_allowed_origin_file() { - assert!(is_allowed_origin(Some("file:///path/to/file"))); - } - - #[test] - fn test_allowed_origin_localhost() { - assert!(is_allowed_origin(Some("http://localhost:3000"))); - assert!(is_allowed_origin(Some("http://127.0.0.1:8080"))); - } - - #[test] - fn test_disallowed_origin() { - assert!(!is_allowed_origin(Some("http://evil.com"))); - } - - #[test] - fn test_frame_metadata_default() { - let meta = FrameMetadata::default(); - assert_eq!(meta.device_width, 1280); - assert_eq!(meta.device_height, 720); - assert_eq!(meta.page_scale_factor, 1.0); - } -} diff --git a/cli/src/native/stream/cdp_loop.rs b/cli/src/native/stream/cdp_loop.rs new file mode 100644 index 0000000..47cf0e8 --- /dev/null +++ b/cli/src/native/stream/cdp_loop.rs @@ -0,0 +1,325 @@ +use serde_json::{json, Value}; +use std::sync::Arc; + +use tokio::sync::{broadcast, watch, Mutex, RwLock}; + +use crate::native::cdp::client::CdpClient; +use crate::native::network; + +use super::timestamp_ms; + +/// Background task that subscribes to CDP events and broadcasts screencast frames in real-time. +/// Also handles auto-start/stop of screencast based on WebSocket client count. +#[allow(clippy::too_many_arguments)] +pub(super) async fn cdp_event_loop( + frame_tx: broadcast::Sender, + client_slot: Arc>>>, + client_notify: Arc, + screencasting: Arc>, + client_count: Arc>, + cdp_session_id: Arc>>, + viewport_width: Arc>, + viewport_height: Arc>, + last_frame: Arc>>, + last_tabs: Arc>>, + last_engine: Arc>, + recording: Arc>, + mut shutdown_rx: watch::Receiver, +) { + loop { + tokio::select! { + changed = shutdown_rx.changed() => { + if changed.is_err() || *shutdown_rx.borrow() { + let session_id = cdp_session_id.read().await.clone(); + if *screencasting.lock().await { + if let Some(ref client) = *client_slot.read().await { + let _ = client + .send_command_no_params("Page.stopScreencast", session_id.as_deref()) + .await; + } + let mut sc = screencasting.lock().await; + *sc = false; + } + return; + } + } + _ = client_notify.notified() => {} + } + + let count = *client_count.lock().await; + let guard = client_slot.read().await; + + if count > 0 { + if let Some(ref client) = *guard { + let mut event_rx = client.subscribe(); + let client_arc = Arc::clone(client); + drop(guard); + + let session_id = cdp_session_id.read().await.clone(); + + let vw = *viewport_width.lock().await; + let vh = *viewport_height.lock().await; + + let eng = last_engine.read().await.clone(); + let supports_screencast = eng == "chrome"; + + if supports_screencast { + let _ = client_arc + .send_command( + "Page.startScreencast", + Some(json!({ + "format": "jpeg", + "quality": 80, + "maxWidth": vw, + "maxHeight": vh, + "everyNthFrame": 1, + })), + session_id.as_deref(), + ) + .await; + } + + { + let mut sc = screencasting.lock().await; + *sc = supports_screencast; + } + + let rec = *recording.lock().await; + let status = json!({ + "type": "status", + "connected": true, + "screencasting": supports_screencast, + "viewportWidth": vw, + "viewportHeight": vh, + "engine": eng, + "recording": rec, + }); + let _ = frame_tx.send(status.to_string()); + + loop { + tokio::select! { + changed = shutdown_rx.changed() => { + if changed.is_err() || *shutdown_rx.borrow() { + if supports_screencast { + let session_id = cdp_session_id.read().await.clone(); + let _ = client_arc + .send_command_no_params("Page.stopScreencast", session_id.as_deref()) + .await; + } + let mut sc = screencasting.lock().await; + *sc = false; + return; + } + } + event = event_rx.recv() => { + match event { + Ok(evt) => { + if evt.method == "Page.frameNavigated" { + if let Some(frame) = evt.params.get("frame") { + let is_main = frame + .get("parentId") + .and_then(|v| v.as_str()) + .is_none_or(|s| s.is_empty()); + if is_main { + if let Some(url) = frame.get("url").and_then(|v| v.as_str()) { + { + let mut tabs = last_tabs.write().await; + for tab in tabs.iter_mut() { + if tab.get("active").and_then(|v| v.as_bool()).unwrap_or(false) { + tab.as_object_mut().map(|o| o.insert("url".to_string(), json!(url))); + } + } + } + let msg = json!({ + "type": "url", + "url": url, + "timestamp": timestamp_ms(), + }); + let _ = frame_tx.send(msg.to_string()); + } + } + } + } else if evt.method == "Page.screencastFrame" { + if let Some(sid) = evt.params.get("sessionId").and_then(|v| v.as_i64()) { + let _ = client_arc.send_command( + "Page.screencastFrameAck", + Some(json!({ "sessionId": sid })), + evt.session_id.as_deref(), + ).await; + } + + if let Some(data) = evt.params.get("data").and_then(|v| v.as_str()) { + let meta = evt.params.get("metadata"); + let msg = json!({ + "type": "frame", + "data": data, + "metadata": { + "offsetTop": meta.and_then(|m| m.get("offsetTop")).and_then(|v| v.as_f64()).unwrap_or(0.0), + "pageScaleFactor": meta.and_then(|m| m.get("pageScaleFactor")).and_then(|v| v.as_f64()).unwrap_or(1.0), + "deviceWidth": meta.and_then(|m| m.get("deviceWidth")).and_then(|v| v.as_u64()).unwrap_or(1280), + "deviceHeight": meta.and_then(|m| m.get("deviceHeight")).and_then(|v| v.as_u64()).unwrap_or(720), + "scrollOffsetX": meta.and_then(|m| m.get("scrollOffsetX")).and_then(|v| v.as_f64()).unwrap_or(0.0), + "scrollOffsetY": meta.and_then(|m| m.get("scrollOffsetY")).and_then(|v| v.as_f64()).unwrap_or(0.0), + "timestamp": meta.and_then(|m| m.get("timestamp")).and_then(|v| v.as_u64()).unwrap_or(0), + } + }); + let msg_str = msg.to_string(); + { + let mut lf = last_frame.write().await; + *lf = Some(msg_str.clone()); + } + let _ = frame_tx.send(msg_str); + } + } else if evt.method == "Runtime.consoleAPICalled" { + let level = evt.params.get("type") + .and_then(|v| v.as_str()) + .unwrap_or("log"); + let raw_args = evt.params.get("args") + .and_then(|v| v.as_array()) + .cloned() + .unwrap_or_default(); + let text = network::format_console_args(&raw_args); + if !text.is_empty() { + let mut msg = json!({ + "type": "console", + "level": level, + "text": text, + "timestamp": timestamp_ms(), + }); + if !raw_args.is_empty() { + msg.as_object_mut().unwrap().insert( + "args".to_string(), + Value::Array(raw_args), + ); + } + let _ = frame_tx.send(msg.to_string()); + } + } else if evt.method == "Runtime.exceptionThrown" { + let text = evt.params.get("exceptionDetails") + .and_then(|d| { + d.get("exception") + .and_then(|e| e.get("description").and_then(|v| v.as_str())) + .or_else(|| d.get("text").and_then(|v| v.as_str())) + }) + .unwrap_or("Unknown error"); + let line = evt.params.get("exceptionDetails") + .and_then(|d| d.get("lineNumber").and_then(|v| v.as_i64())); + let column = evt.params.get("exceptionDetails") + .and_then(|d| d.get("columnNumber").and_then(|v| v.as_i64())); + let msg = json!({ + "type": "page_error", + "text": text, + "line": line, + "column": column, + "timestamp": timestamp_ms(), + }); + let _ = frame_tx.send(msg.to_string()); + } + } + Err(broadcast::error::RecvError::Lagged(_)) => continue, + Err(broadcast::error::RecvError::Closed) => break, + } + } + _ = client_notify.notified() => { + let count = *client_count.lock().await; + let new_session_id = cdp_session_id.read().await.clone(); + if count == 0 { + if supports_screencast { + let _ = client_arc + .send_command_no_params("Page.stopScreencast", session_id.as_deref()) + .await; + } + let mut sc = screencasting.lock().await; + *sc = false; + break; + } + let client_changed = { + let guard = client_slot.read().await; + let same = guard + .as_ref() + .is_some_and(|c| Arc::ptr_eq(c, &client_arc)); + !same + }; + let session_changed = new_session_id != session_id; + let new_vw = *viewport_width.lock().await; + let new_vh = *viewport_height.lock().await; + let viewport_changed = new_vw != vw || new_vh != vh; + if client_changed || session_changed || viewport_changed { + if supports_screencast { + let _ = client_arc + .send_command_no_params("Page.stopScreencast", session_id.as_deref()) + .await; + } + let mut sc = screencasting.lock().await; + *sc = false; + client_notify.notify_one(); + break; + } + } + } + } + } else { + drop(guard); + } + } else { + let was_screencasting = *screencasting.lock().await; + if was_screencasting { + if let Some(ref client) = *guard { + let session_id = cdp_session_id.read().await.clone(); + let _ = client + .send_command_no_params("Page.stopScreencast", session_id.as_deref()) + .await; + } + let mut sc = screencasting.lock().await; + *sc = false; + } + drop(guard); + } + } +} + +pub async fn start_screencast( + client: &CdpClient, + session_id: &str, + format: &str, + quality: i32, + max_width: i32, + max_height: i32, +) -> Result<(), String> { + client + .send_command( + "Page.startScreencast", + Some(json!({ + "format": format, + "quality": quality, + "maxWidth": max_width, + "maxHeight": max_height, + "everyNthFrame": 1, + })), + Some(session_id), + ) + .await?; + Ok(()) +} + +pub async fn stop_screencast(client: &CdpClient, session_id: &str) -> Result<(), String> { + client + .send_command_no_params("Page.stopScreencast", Some(session_id)) + .await?; + Ok(()) +} + +pub async fn ack_screencast_frame( + client: &CdpClient, + session_id: &str, + screencast_session_id: i64, +) -> Result<(), String> { + client + .send_command( + "Page.screencastFrameAck", + Some(json!({ "sessionId": screencast_session_id })), + Some(session_id), + ) + .await?; + Ok(()) +} diff --git a/cli/src/native/stream/chat.rs b/cli/src/native/stream/chat.rs new file mode 100644 index 0000000..e271735 --- /dev/null +++ b/cli/src/native/stream/chat.rs @@ -0,0 +1,970 @@ +use std::sync::OnceLock; + +use serde_json::{json, Value}; + +use tokio::io::AsyncWriteExt; + +use super::http::cors_headers_for_origin; + +const DEFAULT_AI_GATEWAY_URL: &str = "https://ai-gateway.vercel.sh"; + +static HTTP_CLIENT: OnceLock = OnceLock::new(); + +fn http_client() -> &'static reqwest::Client { + HTTP_CLIENT.get_or_init(reqwest::Client::new) +} + +fn is_chat_enabled() -> bool { + std::env::var("AI_GATEWAY_API_KEY").is_ok() +} + +pub(super) fn chat_status_json() -> String { + let enabled = is_chat_enabled(); + let mut obj = json!({ "enabled": enabled }); + if enabled { + if let Ok(model) = std::env::var("AI_GATEWAY_MODEL") { + obj["model"] = Value::String(model); + } + } + obj.to_string() +} + +pub(super) async fn handle_models_request( + stream: &mut tokio::net::TcpStream, + origin: Option<&str>, +) { + let cors = cors_headers_for_origin(origin); + let gateway_url = std::env::var("AI_GATEWAY_URL") + .unwrap_or_else(|_| DEFAULT_AI_GATEWAY_URL.to_string()) + .trim_end_matches('/') + .to_string(); + let api_key = match std::env::var("AI_GATEWAY_API_KEY") { + Ok(k) => k, + Err(_) => { + let body = r#"{"data":[]}"#; + let resp = format!( + "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n{cors}\r\n", + body.len() + ); + let _ = stream.write_all(resp.as_bytes()).await; + let _ = stream.write_all(body.as_bytes()).await; + return; + } + }; + + let url = format!("{}/v1/models", gateway_url); + let client = http_client(); + let result = client + .get(&url) + .header("Authorization", format!("Bearer {}", api_key)) + .send() + .await; + + let body = match result { + Ok(r) if r.status().is_success() => r + .text() + .await + .unwrap_or_else(|_| r#"{"data":[]}"#.to_string()), + _ => r#"{"data":[]}"#.to_string(), + }; + + let resp = format!( + "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n{cors}\r\n", + body.len() + ); + let _ = stream.write_all(resp.as_bytes()).await; + let _ = stream.write_all(body.as_bytes()).await; +} + +const SKILL_NAMES: &[&str] = &["agent-browser", "slack", "electron", "dogfood", "agentcore"]; + +/// Locate the `skills/` directory by walking up from the executable. +/// Works for npm installs (binary in `bin/`, skills at `../skills/`) and +/// dev builds (binary deep in `cli/target/`, skills at repo root). +fn find_skills_dir() -> Option { + let exe = std::env::current_exe().ok()?; + let real = exe.canonicalize().unwrap_or(exe); + let mut dir = real.parent(); + while let Some(d) = dir { + let candidate = d.join("skills"); + if candidate.join("agent-browser").join("SKILL.md").exists() { + return Some(candidate); + } + dir = d.parent(); + } + None +} + +fn load_skills() -> Vec<(String, String)> { + let Some(skills_dir) = find_skills_dir() else { + return Vec::new(); + }; + SKILL_NAMES + .iter() + .filter_map(|name| { + let path = skills_dir.join(name).join("SKILL.md"); + let content = std::fs::read_to_string(&path).ok()?; + Some((name.to_string(), content)) + }) + .collect() +} + +fn strip_frontmatter(s: &str) -> &str { + if !s.starts_with("---") { + return s; + } + if let Some(end) = s[3..].find("---") { + let after = &s[3 + end + 3..]; + after.trim_start_matches(['\n', '\r']) + } else { + s + } +} + +fn get_system_prompt() -> &'static str { + static PROMPT: OnceLock = OnceLock::new(); + PROMPT.get_or_init(|| { + let skills = load_skills(); + + let mut sections = String::new(); + for (name, content) in &skills { + let body = strip_frontmatter(content); + sections.push_str(&format!("\n\n\n{}\n", name, body.trim())); + } + + format!( + r#"You are an AI assistant that controls a browser through agent-browser. You have an active browser session, but you can also create new sessions. + +RULES: +- You MUST use the agent_browser tool for every browser action. NEVER claim you performed an action without calling the tool. +- If the user asks you to do something, call the tool first, then describe the result. +- If a request is outside your capabilities (e.g. system operations), say so honestly. Do not improvise or pretend. +- One tool call per command. Do not chain with `&&` or `;`. +- Do not add `--json`. +- Do not run non-agent-browser programs. +- Keep responses concise. +- For screenshots, omit the path argument so they save to the default location (which will be displayed inline). Screenshots from tool calls are ALREADY shown to the user. Do NOT re-display them with markdown image syntax in your text response. Never use `![...]()` to reference screenshots. +- To create a new session: add `--session ` to any command (e.g. `agent-browser --session my-session open https://example.com`). If the session does not exist, it will be created automatically. +- To use a different browser engine: add `--engine ` (e.g. `agent-browser --session lp-session --engine lightpanda open https://example.com`). Supported engines: chrome (default), lightpanda. + +The following skill references describe agent-browser capabilities in detail. Use them when deciding which commands to run and how to approach tasks. +{sections}"#, + ) + }) +} + +const CHAT_TOOLS: &str = r#"[{"type":"function","function":{"name":"agent_browser","description":"Execute an agent-browser command. Runs against the active session by default. Add --session to target or create a different session, and --engine to choose a browser engine.","parameters":{"type":"object","properties":{"command":{"type":"string","description":"The command to execute, e.g. 'agent-browser open https://google.com' or 'agent-browser --session new-session open https://example.com' or 'agent-browser snapshot -i' or 'agent-browser click @e3'"}},"required":["command"]}}}]"#; + +const COMPACT_THRESHOLD_CHARS: usize = 200_000; +const KEEP_RECENT_MESSAGES: usize = 6; + +fn estimate_chars(messages: &[Value]) -> usize { + messages + .iter() + .map(|m| { + let content_len = m + .get("content") + .map(|c| { + if let Some(s) = c.as_str() { + s.len() + } else { + c.to_string().len() + } + }) + .unwrap_or(0); + let tc_len = m + .get("tool_calls") + .map(|t| t.to_string().len()) + .unwrap_or(0); + content_len + tc_len + }) + .sum() +} + +fn find_safe_split(messages: &[Value], keep_recent: usize) -> usize { + if messages.len() <= keep_recent + 1 { + return 1; + } + let desired = messages.len() - keep_recent; + for i in (1..=desired).rev() { + if messages[i].get("role").and_then(|r| r.as_str()) == Some("user") { + return i; + } + } + desired.max(1) +} + +fn build_summary_text(messages: &[Value]) -> String { + let mut text = String::new(); + for msg in messages { + let role = msg + .get("role") + .and_then(|r| r.as_str()) + .unwrap_or("unknown"); + if let Some(content) = msg.get("content").and_then(|c| c.as_str()) { + if !content.is_empty() { + let truncated = if content.len() > 2000 { + format!("{}...[truncated]", &content[..2000]) + } else { + content.to_string() + }; + text.push_str(&format!("[{}] {}\n\n", role, truncated)); + } + } + if let Some(tcs) = msg.get("tool_calls").and_then(|t| t.as_array()) { + for tc in tcs { + let name = tc + .get("function") + .and_then(|f| f.get("name")) + .and_then(|n| n.as_str()) + .unwrap_or(""); + let args = tc + .get("function") + .and_then(|f| f.get("arguments")) + .and_then(|a| a.as_str()) + .unwrap_or(""); + text.push_str(&format!("[assistant tool:{}] {}\n", name, args)); + } + } + } + text +} + +async fn summarize_for_compaction( + client: &reqwest::Client, + url: &str, + api_key: &str, + model: &str, + messages: &[Value], +) -> Option { + let conversation = build_summary_text(messages); + if conversation.is_empty() { + return None; + } + + let body = json!({ + "model": model, + "messages": [ + { + "role": "system", + "content": "Summarize this browser automation conversation concisely. Preserve: URLs visited, actions performed, current page state, errors encountered, and user goals. Output only the summary." + }, + { + "role": "user", + "content": conversation + } + ], + "max_tokens": 1024, + "stream": false, + }); + + let resp = client + .post(url) + .header("Authorization", format!("Bearer {}", api_key)) + .header("Content-Type", "application/json") + .body(body.to_string()) + .send() + .await + .ok()?; + + if !resp.status().is_success() { + return None; + } + + let result: Value = resp.json().await.ok()?; + result + .get("choices") + .and_then(|c| c.get(0)) + .and_then(|c| c.get("message")) + .and_then(|m| m.get("content")) + .and_then(|c| c.as_str()) + .map(|s| s.to_string()) +} + +const SCREENSHOT_MAX_WIDTH: u32 = 1024; +const SCREENSHOT_JPEG_QUALITY: u8 = 40; + +fn compress_image_to_jpeg(raw_bytes: &[u8]) -> Option> { + let img = image::load_from_memory(raw_bytes).ok()?; + let img = if img.width() > SCREENSHOT_MAX_WIDTH { + img.resize( + SCREENSHOT_MAX_WIDTH, + u32::MAX, + image::imageops::FilterType::Triangle, + ) + } else { + img + }; + let mut buf = std::io::Cursor::new(Vec::new()); + let encoder = + image::codecs::jpeg::JpegEncoder::new_with_quality(&mut buf, SCREENSHOT_JPEG_QUALITY); + img.write_with_encoder(encoder).ok()?; + Some(buf.into_inner()) +} + +fn has_image_extension(s: &str) -> bool { + let lower = s.to_lowercase(); + lower.ends_with(".png") || lower.ends_with(".jpg") || lower.ends_with(".jpeg") +} + +fn extract_image_path(text: &str) -> Option { + for line in text.lines() { + let trimmed = line.trim(); + // Whole line is a path (handles paths with spaces) + if has_image_extension(trimmed) && std::path::Path::new(trimmed).exists() { + return Some(trimmed.to_string()); + } + for suffix in [".png", ".jpg", ".jpeg"] { + if let Some(pos) = trimmed.to_lowercase().rfind(suffix) { + let end = pos + suffix.len(); + let candidate = &trimmed[..end]; + let start = candidate + .rfind(|c: char| c.is_whitespace()) + .map(|i| i + 1) + .unwrap_or(0); + let path = &candidate[start..]; + if !path.is_empty() && std::path::Path::new(path).exists() { + return Some(path.to_string()); + } + } + } + } + None +} + +fn enrich_tool_output(result: &str) -> String { + let Some(path) = extract_image_path(result) else { + return result.to_string(); + }; + + let Ok(raw_bytes) = std::fs::read(&path) else { + return result.to_string(); + }; + + let (jpeg_bytes, mime) = match compress_image_to_jpeg(&raw_bytes) { + Some(compressed) => (compressed, "image/jpeg"), + None => { + let lower = path.to_lowercase(); + ( + raw_bytes, + if lower.ends_with(".png") { + "image/png" + } else { + "image/jpeg" + }, + ) + } + }; + + let b64 = base64::Engine::encode(&base64::engine::general_purpose::STANDARD, &jpeg_bytes); + let data_url = format!("data:{};base64,{}", mime, b64); + + json!({ + "text": result, + "image": data_url + }) + .to_string() +} + +const ALLOWED_COMMANDS: &[&str] = &[ + "open", + "goto", + "navigate", + "back", + "forward", + "reload", + "click", + "dblclick", + "fill", + "type", + "hover", + "focus", + "check", + "uncheck", + "select", + "drag", + "upload", + "download", + "press", + "key", + "keydown", + "keyup", + "keyboard", + "scroll", + "scrollintoview", + "scrollinto", + "wait", + "screenshot", + "pdf", + "snapshot", + "eval", + "close", + "quit", + "exit", + "inspect", + "auth", + "confirm", + "deny", + "connect", + "cookies", + "storage", + "window", + "frame", + "dialog", + "trace", + "profiler", + "record", + "har", + "network", + "title", + "url", + "console", + "errors", + "highlight", + "state", + "emulate", + "video", + "tap", + "swipe", + "device", + "batch", + "diff", + "find", + "role", + "text", + "label", + "placeholder", + "alt", + "testid", + "first", + "last", + "nth", + "mouse", + "touchscreen", + "attribute", + "property", + "set", + "get", + "is", + "stream", + "tab", + "clipboard", + "session", +]; + +const ALLOWED_GLOBAL_FLAGS: &[&str] = &["--session", "--engine"]; + +async fn execute_chat_tool(session: &str, command: &str) -> String { + let exe = match std::env::current_exe() { + Ok(p) => p, + Err(e) => return format!("Failed to resolve executable: {}", e), + }; + + let single = command.split("&&").next().unwrap_or(command); + let single = single.split(';').next().unwrap_or(single).trim(); + let stripped = single.strip_prefix("agent-browser ").unwrap_or(single); + let words = crate::commands::shell_words_split(stripped); + + let mut global_flags: Vec = Vec::new(); + let mut cmd_words: Vec = Vec::new(); + let mut has_session_flag = false; + let mut i = 0; + while i < words.len() { + if ALLOWED_GLOBAL_FLAGS.contains(&words[i].as_str()) { + if words[i] == "--session" { + has_session_flag = true; + } + global_flags.push(words[i].clone()); + if i + 1 < words.len() { + global_flags.push(words[i + 1].clone()); + i += 2; + } else { + i += 1; + } + } else { + cmd_words.push(words[i].clone()); + i += 1; + } + } + + let first_cmd = cmd_words.first().map(|s| s.as_str()).unwrap_or(""); + if !ALLOWED_COMMANDS.contains(&first_cmd) { + return format!( + "Blocked: '{}' is not a valid agent-browser command.", + first_cmd + ); + } + + let mut args: Vec = Vec::new(); + if !has_session_flag { + args.push("--session".into()); + args.push(session.into()); + } + args.extend(global_flags); + args.extend(cmd_words); + + let mut cmd = tokio::process::Command::new(&exe); + cmd.args(&args) + .env_remove("AGENT_BROWSER_DASHBOARD") + .env_remove("AGENT_BROWSER_DASHBOARD_PORT") + .env_remove("AGENT_BROWSER_STREAM_PORT"); + + match cmd.output().await { + Ok(output) => { + let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string(); + let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string(); + if stdout.is_empty() && !stderr.is_empty() { + stderr + } else if stdout.is_empty() { + "Command completed with no output.".to_string() + } else { + stdout + } + } + Err(e) => format!("Failed to execute command: {}", e), + } +} + +async fn stream_gateway_response( + stream: &mut tokio::net::TcpStream, + gw_response: reqwest::Response, +) -> Vec<(String, String, String)> { + use futures_util::StreamExt as _; + + let mut text_part_id = uuid::Uuid::new_v4().to_string(); + let mut text_started = false; + let mut tool_calls: Vec<(String, String, String)> = Vec::new(); + let mut tool_call_args: std::collections::HashMap = + std::collections::HashMap::new(); + let mut byte_stream = gw_response.bytes_stream(); + let mut buffer = String::new(); + + while let Some(chunk_result) = byte_stream.next().await { + let chunk = match chunk_result { + Ok(c) => c, + Err(_) => break, + }; + + buffer.push_str(&String::from_utf8_lossy(&chunk)); + + while let Some(newline_pos) = buffer.find('\n') { + let line = buffer[..newline_pos].trim_end_matches('\r').to_string(); + buffer = buffer[newline_pos + 1..].to_string(); + + if line.is_empty() { + continue; + } + let Some(data) = line.strip_prefix("data: ") else { + continue; + }; + if data == "[DONE]" { + if text_started { + let ev = format!("data: {}\n\n", json!({"type":"text-end","id":text_part_id})); + let _ = stream.write_all(ev.as_bytes()).await; + } + let mut indices: Vec = tool_call_args.keys().copied().collect(); + indices.sort(); + for idx in indices { + if let Some(tc) = tool_call_args.remove(&idx) { + tool_calls.push(tc); + } + } + return tool_calls; + } + let Ok(sse_json) = serde_json::from_str::(data) else { + continue; + }; + let delta = sse_json + .get("choices") + .and_then(|c| c.get(0)) + .and_then(|c| c.get("delta")); + let Some(delta) = delta else { continue }; + + if let Some(text) = delta.get("content").and_then(|c| c.as_str()) { + if !text.is_empty() { + if !text_started { + let ev = format!( + "data: {}\n\n", + json!({"type":"text-start","id":text_part_id}) + ); + if stream.write_all(ev.as_bytes()).await.is_err() { + return tool_calls; + } + text_started = true; + } + let ev = format!( + "data: {}\n\n", + json!({"type":"text-delta","id":text_part_id,"delta":text}) + ); + if stream.write_all(ev.as_bytes()).await.is_err() { + return tool_calls; + } + } + } + + if let Some(tcs) = delta.get("tool_calls").and_then(|t| t.as_array()) { + if text_started { + let ev = format!("data: {}\n\n", json!({"type":"text-end","id":text_part_id})); + let _ = stream.write_all(ev.as_bytes()).await; + text_started = false; + text_part_id = uuid::Uuid::new_v4().to_string(); + } + + for tc in tcs { + let idx = tc.get("index").and_then(|i| i.as_u64()).unwrap_or(0) as usize; + if let std::collections::hash_map::Entry::Vacant(e) = tool_call_args.entry(idx) + { + let id = tc + .get("id") + .and_then(|i| i.as_str()) + .unwrap_or("") + .to_string(); + let name = tc + .get("function") + .and_then(|f| f.get("name")) + .and_then(|n| n.as_str()) + .unwrap_or("") + .to_string(); + let ev = format!( + "data: {}\n\n", + json!({"type":"tool-input-start","toolCallId":id,"toolName":name}) + ); + let _ = stream.write_all(ev.as_bytes()).await; + e.insert((id, name, String::new())); + } + if let Some(arg_delta) = tc + .get("function") + .and_then(|f| f.get("arguments")) + .and_then(|a| a.as_str()) + { + let entry = tool_call_args.get_mut(&idx).unwrap(); + entry.2.push_str(arg_delta); + let ev = format!( + "data: {}\n\n", + json!({"type":"tool-input-delta","toolCallId":entry.0,"inputTextDelta":arg_delta}) + ); + let _ = stream.write_all(ev.as_bytes()).await; + } + } + } + } + } + + if text_started { + let ev = format!("data: {}\n\n", json!({"type":"text-end","id":text_part_id})); + let _ = stream.write_all(ev.as_bytes()).await; + } + let mut indices: Vec = tool_call_args.keys().copied().collect(); + indices.sort(); + for idx in indices { + if let Some(tc) = tool_call_args.remove(&idx) { + tool_calls.push(tc); + } + } + tool_calls +} + +pub(super) async fn handle_chat_request( + stream: &mut tokio::net::TcpStream, + body: &str, + origin: Option<&str>, +) { + let cors = cors_headers_for_origin(origin); + let gateway_url = std::env::var("AI_GATEWAY_URL") + .unwrap_or_else(|_| DEFAULT_AI_GATEWAY_URL.to_string()) + .trim_end_matches('/') + .to_string(); + let api_key = match std::env::var("AI_GATEWAY_API_KEY") { + Ok(k) => k, + Err(_) => { + let err = r#"{"error":"AI_GATEWAY_API_KEY not set. Set the AI_GATEWAY_API_KEY environment variable to enable AI chat."}"#; + let resp = format!( + "HTTP/1.1 500 Internal Server Error\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n{cors}\r\n", + err.len() + ); + let _ = stream.write_all(resp.as_bytes()).await; + let _ = stream.write_all(err.as_bytes()).await; + return; + } + }; + + let default_model = std::env::var("AI_GATEWAY_MODEL") + .unwrap_or_else(|_| "anthropic/claude-sonnet-4.6".to_string()); + + let parsed: Value = match serde_json::from_str(body) { + Ok(v) => v, + Err(e) => { + let err = format!(r#"{{"error":"Invalid JSON: {}"}}"#, e); + let resp = format!( + "HTTP/1.1 400 Bad Request\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n{cors}\r\n", + err.len() + ); + let _ = stream.write_all(resp.as_bytes()).await; + let _ = stream.write_all(err.as_bytes()).await; + return; + } + }; + + let messages = parsed.get("messages").cloned().unwrap_or(json!([])); + let model = parsed + .get("model") + .and_then(|v| v.as_str()) + .unwrap_or(&default_model) + .to_string(); + let session = parsed + .get("session") + .and_then(|v| v.as_str()) + .unwrap_or("default") + .to_string(); + + let mut openai_messages: Vec = + vec![json!({"role": "system", "content": get_system_prompt()})]; + let mut frontend_boundaries: Vec = Vec::new(); + let frontend_arr = messages.as_array(); + let frontend_count = frontend_arr.map(|a| a.len()).unwrap_or(0); + if let Some(arr) = frontend_arr { + for msg in arr { + frontend_boundaries.push(openai_messages.len()); + let Some(role) = msg.get("role").and_then(|r| r.as_str()) else { + continue; + }; + if let Some(parts) = msg.get("parts").and_then(|p| p.as_array()) { + let mut content_parts: Vec = Vec::new(); + for part in parts { + match part.get("type").and_then(|t| t.as_str()) { + Some("text") => { + if let Some(text) = part.get("text").and_then(|t| t.as_str()) { + if !text.is_empty() { + content_parts.push(json!({"type": "text", "text": text})); + } + } + } + Some("file") => { + if let (Some(url), Some(media_type)) = ( + part.get("url").and_then(|u| u.as_str()), + part.get("mediaType").and_then(|m| m.as_str()), + ) { + if media_type.starts_with("image/") { + content_parts.push(json!({ + "type": "image_url", + "image_url": { "url": url } + })); + } + } + } + _ => {} + } + } + if !content_parts.is_empty() { + let content = if content_parts.len() == 1 + && content_parts[0].get("type").and_then(|t| t.as_str()) == Some("text") + { + content_parts[0]["text"].clone() + } else { + json!(content_parts) + }; + openai_messages.push(json!({"role": role, "content": content})); + } + } else if let Some(content) = msg.get("content").and_then(|c| c.as_str()) { + openai_messages.push(json!({"role": role, "content": content})); + } + } + } + + let tools: Value = serde_json::from_str(CHAT_TOOLS).unwrap(); + let url = format!("{}/v1/chat/completions", gateway_url); + let client = http_client(); + + let total_chars = estimate_chars(&openai_messages); + let mut compaction_summary: Option = None; + let mut compaction_failed = false; + let mut keep_last_n: usize = frontend_count; + + if total_chars > COMPACT_THRESHOLD_CHARS && openai_messages.len() > KEEP_RECENT_MESSAGES + 2 { + let split = find_safe_split(&openai_messages, KEEP_RECENT_MESSAGES); + let to_summarize = &openai_messages[1..split]; + + if let Some(summary) = + summarize_for_compaction(client, &url, &api_key, &model, to_summarize).await + { + let summary_msg = json!({ + "role": "system", + "content": format!("[Conversation summary]\n{}", summary) + }); + let recent = openai_messages[split..].to_vec(); + openai_messages = vec![openai_messages[0].clone(), summary_msg]; + openai_messages.extend(recent); + + let kept_frontend = frontend_boundaries + .iter() + .filter(|&&boundary| boundary >= split) + .count(); + keep_last_n = kept_frontend; + compaction_summary = Some(summary); + } else { + compaction_failed = true; + } + } + + let headers = format!( + "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nCache-Control: no-cache\r\nConnection: keep-alive\r\nx-vercel-ai-ui-message-stream: v1\r\n{cors}\r\n" + ); + if stream.write_all(headers.as_bytes()).await.is_err() { + return; + } + + let message_id = uuid::Uuid::new_v4().to_string(); + let start_ev = format!( + "data: {}\n\n", + json!({"type":"start","messageId":message_id}) + ); + if stream.write_all(start_ev.as_bytes()).await.is_err() { + return; + } + + if let Some(ref summary) = compaction_summary { + let ev = format!( + "data: {}\n\n", + json!({ + "type": "message-metadata", + "messageMetadata": { + "compacted": true, + "summary": summary, + "keepLastN": keep_last_n + } + }) + ); + let _ = stream.write_all(ev.as_bytes()).await; + } else if compaction_failed { + let ev = format!( + "data: {}\n\n", + json!({ + "type": "message-metadata", + "messageMetadata": { + "compacted": false, + "warning": "Conversation is large but compaction failed. Responses may be degraded." + } + }) + ); + let _ = stream.write_all(ev.as_bytes()).await; + } + + let total_deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(300); + const TOOL_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60); + + for _step in 0..50 { + if tokio::time::Instant::now() >= total_deadline { + let ev = format!( + "data: {}\n\n", + json!({"type":"error","errorText":"Chat session timed out (5 minute limit)."}) + ); + let _ = stream.write_all(ev.as_bytes()).await; + break; + } + + let step_ev = "data: {\"type\":\"start-step\"}\n\n"; + if stream.write_all(step_ev.as_bytes()).await.is_err() { + return; + } + + let gateway_body = json!({ + "model": model, + "messages": openai_messages, + "tools": tools, + "stream": true, + }); + + let gw_response = match client + .post(&url) + .header("Authorization", format!("Bearer {}", api_key)) + .header("Content-Type", "application/json") + .body(gateway_body.to_string()) + .send() + .await + { + Ok(r) => r, + Err(e) => { + let ev = format!( + "data: {}\n\n", + json!({"type":"error","errorText":format!("Gateway request failed: {}", e)}) + ); + let _ = stream.write_all(ev.as_bytes()).await; + break; + } + }; + + if !gw_response.status().is_success() { + let body_text = gw_response.text().await.unwrap_or_default(); + let ev = format!( + "data: {}\n\n", + json!({"type":"error","errorText":body_text}) + ); + let _ = stream.write_all(ev.as_bytes()).await; + break; + } + + let tool_calls = stream_gateway_response(stream, gw_response).await; + + if tool_calls.is_empty() { + let finish_step_ev = "data: {\"type\":\"finish-step\"}\n\n"; + let _ = stream.write_all(finish_step_ev.as_bytes()).await; + break; + } + + let tc_values: Vec = tool_calls.iter().map(|(id, name, args)| { + json!({"id": id, "type": "function", "function": {"name": name, "arguments": args}}) + }).collect(); + openai_messages.push(json!({"role": "assistant", "tool_calls": tc_values})); + + for (tc_id, tc_name, tc_args) in &tool_calls { + let input: Value = serde_json::from_str(tc_args).unwrap_or(json!({})); + let command = input.get("command").and_then(|c| c.as_str()).unwrap_or(""); + + let ev = format!( + "data: {}\n\n", + json!({ + "type": "tool-input-available", + "toolCallId": tc_id, + "toolName": tc_name, + "input": input + }) + ); + let _ = stream.write_all(ev.as_bytes()).await; + + let result = match tokio::time::timeout( + TOOL_TIMEOUT, + execute_chat_tool(&session, command), + ) + .await + { + Ok(r) => r, + Err(_) => "Tool execution timed out after 60 seconds.".to_string(), + }; + + let frontend_output = enrich_tool_output(&result); + let ev = format!( + "data: {}\n\n", + json!({ + "type": "tool-output-available", + "toolCallId": tc_id, + "output": frontend_output + }) + ); + let _ = stream.write_all(ev.as_bytes()).await; + + openai_messages.push(json!({ + "role": "tool", + "tool_call_id": tc_id, + "content": result + })); + } + + let finish_step_ev = "data: {\"type\":\"finish-step\"}\n\n"; + let _ = stream.write_all(finish_step_ev.as_bytes()).await; + } + + let finish_ev = "data: {\"type\":\"finish\"}\n\n"; + let _ = stream.write_all(finish_ev.as_bytes()).await; + let done_ev = "data: [DONE]\n\n"; + let _ = stream.write_all(done_ev.as_bytes()).await; +} diff --git a/cli/src/native/stream/dashboard.rs b/cli/src/native/stream/dashboard.rs new file mode 100644 index 0000000..0bde1e9 --- /dev/null +++ b/cli/src/native/stream/dashboard.rs @@ -0,0 +1,310 @@ +use serde_json::{json, Value}; +use std::path::PathBuf; +use std::sync::Arc; + +use tokio::io::AsyncWriteExt; +use tokio::net::TcpListener; + +use crate::connection::get_socket_dir; +use crate::install::get_dashboard_dir; + +use super::chat::{chat_status_json, handle_chat_request, handle_models_request}; +use super::discovery::discover_sessions; +use super::http::{serve_static_file, CORS_HEADERS, DASHBOARD_NOT_INSTALLED_HTML}; + +pub async fn run_dashboard_server(port: u16) { + let addr = format!("127.0.0.1:{}", port); + let listener = match TcpListener::bind(&addr).await { + Ok(l) => l, + Err(e) => { + eprintln!("Failed to bind dashboard server on {}: {}", addr, e); + return; + } + }; + + let dashboard_dir: Arc = Arc::from(get_dashboard_dir()); + + loop { + let Ok((stream, _addr)) = listener.accept().await else { + break; + }; + let dash_dir = dashboard_dir.clone(); + tokio::spawn(async move { + handle_dashboard_connection(stream, dash_dir).await; + }); + } +} + +async fn handle_dashboard_connection( + mut stream: tokio::net::TcpStream, + dashboard_dir: Arc, +) { + use tokio::io::AsyncReadExt; + + let mut buf = vec![0u8; 8192]; + let n = match stream.read(&mut buf).await { + Ok(n) if n > 0 => n, + _ => return, + }; + + let header_str = std::str::from_utf8(&buf[..n]).unwrap_or(""); + let first_line = header_str.lines().next().unwrap_or("").to_string(); + let method = first_line.split_whitespace().next().unwrap_or("GET"); + let path = first_line.split_whitespace().nth(1).unwrap_or("/"); + let origin = header_str.lines().find_map(|line| { + if line.len() > 8 && line[..8].eq_ignore_ascii_case("origin: ") { + Some(line[8..].trim().to_string()) + } else { + None + } + }); + + if method == "OPTIONS" { + let response = format!( + "HTTP/1.1 204 No Content\r\n{CORS_HEADERS}Access-Control-Max-Age: 86400\r\nContent-Length: 0\r\nConnection: close\r\n\r\n" + ); + let _ = stream.write_all(response.as_bytes()).await; + return; + } + + if method == "POST" && path == "/api/chat" { + let body_str = read_post_body(&mut stream, &buf, n).await; + handle_chat_request(&mut stream, &body_str, origin.as_deref()).await; + return; + } + + if method == "GET" && path == "/api/models" { + handle_models_request(&mut stream, origin.as_deref()).await; + return; + } + + if method == "POST" && (path == "/api/sessions" || path == "/api/exec" || path == "/api/kill") { + let body_str = read_post_body(&mut stream, &buf, n).await; + let result = if path == "/api/exec" { + exec_cli(&body_str).await + } else if path == "/api/kill" { + kill_session(&body_str).await + } else { + spawn_session(&body_str).await + }; + let (status, resp_body) = match result { + Ok(msg) => ("200 OK", msg), + Err(e) => ( + "400 Bad Request", + format!( + r#"{{"success":false,"error":{}}}"#, + serde_json::to_string(&e).unwrap_or_else(|_| format!("\"{}\"", e)) + ), + ), + }; + let response = format!( + "HTTP/1.1 {status}\r\nContent-Type: application/json; charset=utf-8\r\nContent-Length: {}\r\nConnection: close\r\n{CORS_HEADERS}\r\n", + resp_body.len() + ); + let _ = stream.write_all(response.as_bytes()).await; + let _ = stream.write_all(resp_body.as_bytes()).await; + return; + } + + let (status, content_type, body): (&str, &str, Vec) = if path == "/api/sessions" { + ( + "200 OK", + "application/json; charset=utf-8", + discover_sessions().into_bytes(), + ) + } else if path == "/api/chat/status" { + ( + "200 OK", + "application/json; charset=utf-8", + chat_status_json().into_bytes(), + ) + } else if dashboard_dir.join("index.html").exists() { + serve_static_file(&dashboard_dir, path) + } else { + ( + "200 OK", + "text/html; charset=utf-8", + DASHBOARD_NOT_INSTALLED_HTML.as_bytes().to_vec(), + ) + }; + + let response = format!( + "HTTP/1.1 {}\r\nContent-Type: {}\r\nContent-Length: {}\r\nConnection: close\r\n{CORS_HEADERS}\r\n", + status, + content_type, + body.len() + ); + let _ = stream.write_all(response.as_bytes()).await; + let _ = stream.write_all(&body).await; +} + +async fn read_post_body(stream: &mut tokio::net::TcpStream, initial: &[u8], n: usize) -> String { + use tokio::io::AsyncReadExt; + + let header_end = initial[..n] + .windows(4) + .position(|w| w == b"\r\n\r\n") + .map(|p| p + 4) + .or_else(|| { + initial[..n] + .windows(2) + .position(|w| w == b"\n\n") + .map(|p| p + 2) + }); + let Some(header_end) = header_end else { + return String::new(); + }; + + let header_str = String::from_utf8_lossy(&initial[..header_end]); + let content_length: usize = header_str + .lines() + .find_map(|l| { + if l.len() > 16 && l[..16].eq_ignore_ascii_case("content-length: ") { + l[16..].trim().parse().ok() + } else { + let lower = l.to_lowercase(); + lower + .strip_prefix("content-length:") + .and_then(|v| v.trim().parse().ok()) + } + }) + .unwrap_or(0); + + if content_length == 0 { + return String::new(); + } + + let read_body = &initial[header_end..n]; + let already_read = read_body.len().min(content_length); + + let mut body = Vec::with_capacity(content_length); + body.extend_from_slice(&read_body[..already_read]); + + let remaining = content_length - already_read; + if remaining > 0 { + let mut rest = vec![0u8; remaining]; + if stream.read_exact(&mut rest).await.is_ok() { + body.extend_from_slice(&rest); + } + } + + String::from_utf8(body).unwrap_or_default() +} + +async fn exec_cli(body: &str) -> Result { + let parsed: Value = serde_json::from_str(body).map_err(|e| format!("Invalid JSON: {}", e))?; + let args: Vec = parsed + .get("args") + .and_then(|v| v.as_array()) + .ok_or("Missing \"args\" array")? + .iter() + .filter_map(|v| v.as_str().map(|s| s.to_string())) + .collect(); + + if args.is_empty() { + return Err("Empty args array".to_string()); + } + + let exe = std::env::current_exe().map_err(|e| format!("Cannot resolve executable: {}", e))?; + + let mut cmd = tokio::process::Command::new(&exe); + cmd.args(&args) + .arg("--json") + .env_remove("AGENT_BROWSER_DASHBOARD") + .env_remove("AGENT_BROWSER_DASHBOARD_PORT") + .env_remove("AGENT_BROWSER_STREAM_PORT"); + + let output = cmd + .output() + .await + .map_err(|e| format!("Failed to execute: {}", e))?; + + let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string(); + let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string(); + + Ok(json!({ + "success": output.status.success(), + "exit_code": output.status.code(), + "stdout": stdout, + "stderr": stderr, + }) + .to_string()) +} + +async fn kill_session(body: &str) -> Result { + let parsed: Value = serde_json::from_str(body).map_err(|e| format!("Invalid JSON: {}", e))?; + let session = parsed + .get("session") + .and_then(|v| v.as_str()) + .ok_or("Missing \"session\" field")?; + + if session.is_empty() || session.len() > 64 { + return Err("Session name must be 1-64 characters".to_string()); + } + + let dir = get_socket_dir(); + let pid_path = dir.join(format!("{}.pid", session)); + + let pid_str = std::fs::read_to_string(&pid_path) + .map_err(|_| format!("No PID file for session '{}'", session))?; + let pid: u32 = pid_str + .trim() + .parse() + .map_err(|_| format!("Invalid PID in file: {}", pid_str.trim()))?; + + #[cfg(unix)] + { + unsafe { + libc::kill(pid as i32, libc::SIGTERM); + } + tokio::time::sleep(std::time::Duration::from_millis(500)).await; + if unsafe { libc::kill(pid as i32, 0) } == 0 { + unsafe { + libc::kill(pid as i32, libc::SIGKILL); + } + } + } + + for ext in &["pid", "sock", "stream", "engine", "extensions"] { + let _ = std::fs::remove_file(dir.join(format!("{}.{}", session, ext))); + } + + Ok(json!({ "success": true, "killed_pid": pid }).to_string()) +} + +pub(super) async fn spawn_session(body: &str) -> Result { + let parsed: Value = serde_json::from_str(body).map_err(|e| format!("Invalid JSON: {}", e))?; + let session = parsed + .get("session") + .and_then(|v| v.as_str()) + .ok_or("Missing \"session\" field")?; + + if session.is_empty() || session.len() > 64 { + return Err("Session name must be 1-64 characters".to_string()); + } + + let exe = std::env::current_exe().map_err(|e| format!("Cannot resolve executable: {}", e))?; + + let mut cmd = tokio::process::Command::new(&exe); + cmd.arg("open") + .arg("about:blank") + .arg("--session") + .arg(session); + + cmd.stdout(std::process::Stdio::null()); + cmd.stderr(std::process::Stdio::null()); + + let status = cmd + .status() + .await + .map_err(|e| format!("Failed to spawn session: {}", e))?; + + if status.success() { + Ok(format!( + r#"{{"success":true,"session":{}}}"#, + serde_json::to_string(session).unwrap_or_default() + )) + } else { + Err(format!("Session process exited with {}", status)) + } +} diff --git a/cli/src/native/stream/discovery.rs b/cli/src/native/stream/discovery.rs new file mode 100644 index 0000000..6f2359b --- /dev/null +++ b/cli/src/native/stream/discovery.rs @@ -0,0 +1,118 @@ +use serde_json::{json, Value}; +use std::path::Path; + +use crate::connection::get_socket_dir; + +pub(super) fn discover_sessions() -> String { + let dir = get_socket_dir(); + let mut sessions = Vec::new(); + + if let Ok(entries) = std::fs::read_dir(&dir) { + for entry in entries.flatten() { + let name = entry.file_name(); + let name_str = name.to_string_lossy(); + if let Some(session) = name_str.strip_suffix(".stream") { + if let Ok(port_str) = std::fs::read_to_string(entry.path()) { + if let Ok(port) = port_str.trim().parse::() { + let pid_path = dir.join(format!("{}.pid", session)); + if is_process_alive(&pid_path) { + let engine_path = dir.join(format!("{}.engine", session)); + let engine = std::fs::read_to_string(&engine_path) + .ok() + .filter(|s| !s.trim().is_empty()) + .unwrap_or_else(|| "chrome".to_string()); + + let provider_path = dir.join(format!("{}.provider", session)); + let provider = std::fs::read_to_string(&provider_path) + .ok() + .filter(|s| !s.trim().is_empty()); + + let extensions = read_extensions_metadata(&dir, session); + + let mut entry = json!({ + "session": session, + "port": port, + "engine": engine.trim(), + }); + if let Some(ref p) = provider { + entry["provider"] = json!(p.trim()); + } + if !extensions.is_empty() { + entry["extensions"] = json!(extensions); + } + sessions.push(entry); + } else { + let _ = std::fs::remove_file(entry.path()); + } + } + } + } + } + } + + serde_json::to_string(&sessions).unwrap_or_else(|_| "[]".to_string()) +} + +fn read_extensions_metadata(dir: &std::path::Path, session: &str) -> Vec { + let ext_path = dir.join(format!("{}.extensions", session)); + let ext_str = match std::fs::read_to_string(&ext_path) { + Ok(s) => s, + Err(_) => return Vec::new(), + }; + + ext_str + .split(',') + .map(|p| p.trim()) + .filter(|p| !p.is_empty()) + .filter_map(|path| { + let manifest_path = std::path::Path::new(path).join("manifest.json"); + let manifest_str = std::fs::read_to_string(&manifest_path).ok()?; + let manifest: Value = serde_json::from_str(&manifest_str).ok()?; + + let name = manifest + .get("name") + .and_then(|v| v.as_str()) + .unwrap_or("Unknown") + .to_string(); + let version = manifest + .get("version") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let description = manifest + .get("description") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + + let mut ext = json!({ + "name": name, + "version": version, + "path": path, + }); + if let Some(desc) = description { + ext["description"] = json!(desc); + } + Some(ext) + }) + .collect() +} + +fn is_process_alive(pid_path: &Path) -> bool { + let pid_str = match std::fs::read_to_string(pid_path) { + Ok(s) => s, + Err(_) => return false, + }; + let pid: u32 = match pid_str.trim().parse() { + Ok(p) => p, + Err(_) => return false, + }; + #[cfg(unix)] + { + unsafe { libc::kill(pid as i32, 0) == 0 } + } + #[cfg(not(unix))] + { + let _ = pid; + true + } +} diff --git a/cli/src/native/stream/http.rs b/cli/src/native/stream/http.rs new file mode 100644 index 0000000..07e9947 --- /dev/null +++ b/cli/src/native/stream/http.rs @@ -0,0 +1,339 @@ +use serde_json::{json, Value}; +use std::path::Path; +use std::sync::Arc; + +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::sync::RwLock; + +use crate::connection::get_socket_dir; +#[cfg(windows)] +use crate::connection::resolve_port; + +use super::chat::{chat_status_json, handle_chat_request, handle_models_request}; +use super::dashboard::spawn_session; +use super::discovery::discover_sessions; + +pub(super) const CORS_HEADERS: &str = "Access-Control-Allow-Origin: *\r\nAccess-Control-Allow-Methods: GET, POST, OPTIONS\r\nAccess-Control-Allow-Headers: Content-Type\r\n"; + +/// Build CORS headers that reflect the request origin only when it passes +/// `is_allowed_origin`. Used for sensitive endpoints (chat, models) so the +/// API key is not accessible from arbitrary web pages. +pub(super) fn cors_headers_for_origin(origin: Option<&str>) -> String { + let allowed_origin = match origin { + Some(o) if super::is_allowed_origin(Some(o)) => o, + _ => "http://localhost", + }; + format!( + "Access-Control-Allow-Origin: {}\r\nAccess-Control-Allow-Methods: GET, POST, OPTIONS\r\nAccess-Control-Allow-Headers: Content-Type\r\n", + allowed_origin + ) +} + +fn parse_origin(peeked: &[u8]) -> Option { + let header_str = std::str::from_utf8(peeked).ok()?; + for line in header_str.lines() { + if line.len() > 8 && line[..8].eq_ignore_ascii_case("origin: ") { + return Some(line[8..].trim().to_string()); + } + } + None +} + +pub(super) async fn handle_http_request( + mut stream: tokio::net::TcpStream, + peeked: &[u8], + dashboard_dir: Option<&Path>, + last_tabs: &Arc>>, + last_engine: &Arc>, + session_name: &str, +) { + let peeked_len = peeked.len(); + let mut discard = vec![0u8; peeked_len]; + let _ = stream.read_exact(&mut discard).await; + + let request = String::from_utf8_lossy(peeked); + let first_line = request.lines().next().unwrap_or(""); + let method = first_line.split_whitespace().next().unwrap_or("GET"); + let path = first_line.split_whitespace().nth(1).unwrap_or("/"); + let origin = parse_origin(peeked); + + if method == "OPTIONS" { + let response = format!( + "HTTP/1.1 204 No Content\r\n{CORS_HEADERS}Access-Control-Max-Age: 86400\r\nContent-Length: 0\r\nConnection: close\r\n\r\n" + ); + let _ = stream.write_all(response.as_bytes()).await; + return; + } + + if method == "POST" { + let full_body = read_full_body(&mut stream, peeked).await; + if full_body.is_none() + && (path == "/api/chat" || path == "/api/sessions" || path == "/api/command") + { + let body = r#"{"error":"Request body too large"}"#; + let response = format!( + "HTTP/1.1 413 Payload Too Large\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n{CORS_HEADERS}\r\n", + body.len() + ); + let _ = stream.write_all(response.as_bytes()).await; + let _ = stream.write_all(body.as_bytes()).await; + return; + } + let body_str = full_body.as_deref().unwrap_or(""); + + if path == "/api/sessions" { + let result = spawn_session(body_str).await; + let (status, resp_body) = match result { + Ok(msg) => ("200 OK", msg), + Err(e) => ( + "400 Bad Request", + format!( + r#"{{"success":false,"error":{}}}"#, + serde_json::to_string(&e).unwrap_or_else(|_| format!("\"{}\"", e)) + ), + ), + }; + let response = format!( + "HTTP/1.1 {status}\r\nContent-Type: application/json; charset=utf-8\r\nContent-Length: {}\r\nConnection: close\r\n{CORS_HEADERS}\r\n", + resp_body.len() + ); + let _ = stream.write_all(response.as_bytes()).await; + let _ = stream.write_all(resp_body.as_bytes()).await; + return; + } + + if path == "/api/command" { + let result = relay_command_to_daemon(session_name, body_str).await; + let (status, resp_body) = match result { + Ok(resp) => ("200 OK", resp), + Err(e) => ( + "502 Bad Gateway", + format!( + r#"{{"success":false,"error":{}}}"#, + serde_json::to_string(&e).unwrap_or_else(|_| format!("\"{}\"", e)) + ), + ), + }; + let response = format!( + "HTTP/1.1 {status}\r\nContent-Type: application/json; charset=utf-8\r\nContent-Length: {}\r\nConnection: close\r\n{CORS_HEADERS}\r\n", + resp_body.len() + ); + let _ = stream.write_all(response.as_bytes()).await; + let _ = stream.write_all(resp_body.as_bytes()).await; + return; + } + + if path == "/api/chat" { + handle_chat_request(&mut stream, body_str, origin.as_deref()).await; + return; + } + } + + if method == "GET" && path == "/api/models" { + handle_models_request(&mut stream, origin.as_deref()).await; + return; + } + + let (status, content_type, body): (&str, &str, Vec) = if path == "/api/sessions" { + ( + "200 OK", + "application/json; charset=utf-8", + discover_sessions().into_bytes(), + ) + } else if path == "/api/tabs" { + let tabs = last_tabs.read().await; + ( + "200 OK", + "application/json; charset=utf-8", + serde_json::to_string(&*tabs) + .unwrap_or_else(|_| "[]".to_string()) + .into_bytes(), + ) + } else if path == "/api/status" { + let engine = last_engine.read().await; + ( + "200 OK", + "application/json; charset=utf-8", + format!(r#"{{"engine":"{}"}}"#, *engine).into_bytes(), + ) + } else if path == "/api/chat/status" { + ( + "200 OK", + "application/json; charset=utf-8", + chat_status_json().into_bytes(), + ) + } else { + match dashboard_dir { + Some(dir) => serve_static_file(dir, path), + None => ( + "200 OK", + "text/html; charset=utf-8", + DASHBOARD_NOT_INSTALLED_HTML.as_bytes().to_vec(), + ), + } + }; + + let response = format!( + "HTTP/1.1 {}\r\nContent-Type: {}\r\nContent-Length: {}\r\nConnection: close\r\n{CORS_HEADERS}\r\n", + status, + content_type, + body.len() + ); + let _ = stream.write_all(response.as_bytes()).await; + let _ = stream.write_all(&body).await; +} + +fn find_header_end(buf: &[u8]) -> Option { + buf.windows(4) + .position(|w| w == b"\r\n\r\n") + .map(|p| p + 4) + .or_else(|| buf.windows(2).position(|w| w == b"\n\n").map(|p| p + 2)) +} + +fn parse_content_length_bytes(headers: &[u8]) -> Option { + let header_str = std::str::from_utf8(headers).ok()?; + for line in header_str.lines() { + if line.len() > 16 && line[..16].eq_ignore_ascii_case("content-length: ") { + return line[16..].trim().parse().ok(); + } + } + None +} + +const MAX_BODY_SIZE: usize = 10 * 1024 * 1024; + +async fn read_full_body(stream: &mut tokio::net::TcpStream, peeked: &[u8]) -> Option { + let body_offset = find_header_end(peeked)?; + let content_length = parse_content_length_bytes(&peeked[..body_offset])?; + if content_length == 0 { + return Some(String::new()); + } + if content_length > MAX_BODY_SIZE { + return None; + } + + let peeked_body = &peeked[body_offset..]; + let peeked_body_len = peeked_body.len().min(content_length); + + let mut body = Vec::with_capacity(content_length); + body.extend_from_slice(&peeked_body[..peeked_body_len]); + + let remaining = content_length - peeked_body_len; + if remaining > 0 { + let mut rest = vec![0u8; remaining]; + if stream.read_exact(&mut rest).await.is_err() { + return String::from_utf8(body).ok(); + } + body.extend_from_slice(&rest); + } + + String::from_utf8(body).ok() +} + +pub(super) async fn relay_command_to_daemon( + session_name: &str, + body: &str, +) -> Result { + let mut cmd: Value = serde_json::from_str(body).map_err(|e| format!("Invalid JSON: {}", e))?; + + if cmd.get("id").is_none() { + let id = format!( + "dash-{}", + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_millis() + ); + cmd["id"] = json!(id); + } + + let mut json_str = serde_json::to_string(&cmd).map_err(|e| e.to_string())?; + json_str.push('\n'); + + #[cfg(unix)] + let stream = { + let socket_path = get_socket_dir().join(format!("{}.sock", session_name)); + tokio::net::UnixStream::connect(&socket_path) + .await + .map_err(|e| format!("Failed to connect to daemon: {}", e))? + }; + + #[cfg(windows)] + let stream = { + let port = resolve_port(session_name); + tokio::net::TcpStream::connect(format!("127.0.0.1:{}", port)) + .await + .map_err(|e| format!("Failed to connect to daemon: {}", e))? + }; + + let (reader, mut writer) = tokio::io::split(stream); + + writer + .write_all(json_str.as_bytes()) + .await + .map_err(|e| format!("Failed to send command: {}", e))?; + + let mut buf_reader = tokio::io::BufReader::new(reader); + let mut response_line = String::new(); + tokio::io::AsyncBufReadExt::read_line(&mut buf_reader, &mut response_line) + .await + .map_err(|e| format!("Failed to read response: {}", e))?; + + Ok(response_line.trim().to_string()) +} + +pub(super) fn serve_static_file( + dir: &Path, + url_path: &str, +) -> (&'static str, &'static str, Vec) { + let clean = url_path.trim_start_matches('/'); + let file_path = if clean.is_empty() { + dir.join("index.html") + } else { + let joined = dir.join(clean); + if joined.is_file() { + joined + } else { + dir.join("index.html") + } + }; + + match std::fs::read(&file_path) { + Ok(content) => { + let ext = file_path.extension().and_then(|e| e.to_str()).unwrap_or(""); + let ct = match ext { + "html" => "text/html; charset=utf-8", + "js" => "application/javascript; charset=utf-8", + "css" => "text/css; charset=utf-8", + "json" => "application/json; charset=utf-8", + "svg" => "image/svg+xml", + "png" => "image/png", + "ico" => "image/x-icon", + _ => "application/octet-stream", + }; + ("200 OK", ct, content) + } + Err(_) => ( + "404 Not Found", + "text/html; charset=utf-8", + b"

404 Not Found

".to_vec(), + ), + } +} + +pub(super) const DASHBOARD_NOT_INSTALLED_HTML: &str = r#" + +agent-browser + + + +
+

Dashboard not installed

+

Run agent-browser dashboard install to download the dashboard.

+
+ +"#; diff --git a/cli/src/native/stream/mod.rs b/cli/src/native/stream/mod.rs new file mode 100644 index 0000000..9fd433a --- /dev/null +++ b/cli/src/native/stream/mod.rs @@ -0,0 +1,501 @@ +mod cdp_loop; +mod chat; +mod dashboard; +mod discovery; +mod http; +mod websocket; + +pub use cdp_loop::{ack_screencast_frame, start_screencast, stop_screencast}; +pub use dashboard::run_dashboard_server; + +use serde_json::{json, Value}; +use std::path::PathBuf; +use std::sync::Arc; + +use tokio::net::TcpListener; +use tokio::sync::{broadcast, watch, Mutex, Notify, RwLock}; + +use super::cdp::client::CdpClient; + +/// Frame metadata from CDP Page.screencastFrame events. +#[derive(Debug, Clone)] +pub struct FrameMetadata { + pub offset_top: f64, + pub page_scale_factor: f64, + pub device_width: u32, + pub device_height: u32, + pub scroll_offset_x: f64, + pub scroll_offset_y: f64, + pub timestamp: u64, +} + +impl Default for FrameMetadata { + fn default() -> Self { + Self { + offset_top: 0.0, + page_scale_factor: 1.0, + device_width: 1280, + device_height: 720, + scroll_offset_x: 0.0, + scroll_offset_y: 0.0, + timestamp: 0, + } + } +} + +pub struct StreamServer { + port: u16, + session_name: String, + frame_tx: broadcast::Sender, + client_count: Arc>, + client_slot: Arc>>>, + /// The active CDP page session ID (from Target.attachToTarget). + cdp_session_id: Arc>>, + client_notify: Arc, + screencasting: Arc>, + viewport_width: Arc>, + viewport_height: Arc>, + dashboard_dir: Option, + last_tabs: Arc>>, + last_engine: Arc>, + last_frame: Arc>>, + recording: Arc>, + shutdown_tx: watch::Sender, + accept_task: Mutex>>, + cdp_task: Mutex>>, +} + +impl StreamServer { + pub async fn start( + preferred_port: u16, + client: Arc, + session_id: String, + ) -> Result { + let client_slot = Arc::new(RwLock::new(Some(client))); + let (server, _) = Self::start_inner(preferred_port, client_slot, session_id, true).await?; + Ok(server) + } + + /// Start the stream server without a CDP client. + /// Returns the server and a shared slot to set the client when the browser launches. + /// Input messages are ignored until the client is set. + /// When `allow_port_fallback` is true, binding to an occupied port falls back to an + /// OS-assigned port (used by daemon startup). When false, the error propagates + /// (used by the runtime `stream_enable` command). + pub async fn start_without_client( + preferred_port: u16, + session_id: String, + allow_port_fallback: bool, + ) -> Result<(Self, Arc>>>), String> { + let client_slot = Arc::new(RwLock::new(None::>)); + Self::start_inner(preferred_port, client_slot, session_id, allow_port_fallback).await + } + + /// Resolve the dashboard directory if it exists. + fn resolve_dashboard_dir() -> Option { + let dir = dirs::home_dir()?.join(".agent-browser").join("dashboard"); + if dir.join("index.html").exists() { + Some(dir) + } else { + None + } + } + + /// Notify the background CDP listener that the client has changed (browser launched/closed). + pub fn notify_client_changed(&self) { + self.client_notify.notify_one(); + } + + /// Update the active CDP page session ID used for screencast commands. + pub async fn set_cdp_session_id(&self, session_id: Option) { + let mut guard = self.cdp_session_id.write().await; + *guard = session_id; + } + + /// Check whether the server currently has active screencast running. + pub async fn is_screencasting(&self) -> bool { + *self.screencasting.lock().await + } + + /// Update the stored viewport dimensions used by status messages and screencast. + /// Also notifies the screencast event loop to restart with the new dimensions. + pub async fn set_viewport(&self, width: u32, height: u32) { + *self.viewport_width.lock().await = width; + *self.viewport_height.lock().await = height; + self.client_notify.notify_one(); + } + + /// Get the current viewport dimensions. + pub async fn viewport(&self) -> (u32, u32) { + let w = *self.viewport_width.lock().await; + let h = *self.viewport_height.lock().await; + (w, h) + } + + /// Override the cached screencast state for explicit CLI start/stop commands. + pub async fn set_screencasting(&self, active: bool) { + let mut guard = self.screencasting.lock().await; + *guard = active; + } + + /// Update and broadcast the recording state. + pub async fn set_recording(&self, active: bool, engine: &str) { + *self.recording.lock().await = active; + let connected = self.client_slot.read().await.is_some(); + let sc = *self.screencasting.lock().await; + let (vw, vh) = self.viewport().await; + self.broadcast_status(connected, sc, vw, vh, engine).await; + } + + /// Shut down the accept loop and background CDP listener, releasing the bound port. + pub async fn shutdown(&self) { + let _ = self.shutdown_tx.send(true); + + if let Some(task) = self.accept_task.lock().await.take() { + let _ = task.await; + } + if let Some(task) = self.cdp_task.lock().await.take() { + let _ = task.await; + } + } + + async fn start_inner( + preferred_port: u16, + client_slot: Arc>>>, + session_id: String, + allow_port_fallback: bool, + ) -> Result<(Self, Arc>>>), String> { + let addr = format!("127.0.0.1:{}", preferred_port); + let listener = match TcpListener::bind(&addr).await { + Ok(l) => l, + Err(_) if allow_port_fallback && preferred_port != 0 => { + TcpListener::bind("127.0.0.1:0") + .await + .map_err(|e| format!("Failed to bind stream server: {}", e))? + } + Err(e) => return Err(format!("Failed to bind stream server: {}", e)), + }; + + let actual_addr = listener + .local_addr() + .map_err(|e| format!("Failed to get stream address: {}", e))?; + let port = actual_addr.port(); + + let dashboard_dir = Self::resolve_dashboard_dir(); + + let (frame_tx, _) = broadcast::channel::(64); + let client_count = Arc::new(Mutex::new(0usize)); + let client_notify = Arc::new(Notify::new()); + let screencasting = Arc::new(Mutex::new(false)); + let cdp_session_id = Arc::new(RwLock::new(None::)); + let viewport_width = Arc::new(Mutex::new(1280u32)); + let viewport_height = Arc::new(Mutex::new(720u32)); + let last_tabs = Arc::new(RwLock::new(Vec::::new())); + let last_engine = Arc::new(RwLock::new("chrome".to_string())); + let last_frame = Arc::new(RwLock::new(None::)); + let recording = Arc::new(Mutex::new(false)); + let (shutdown_tx, shutdown_rx) = watch::channel(false); + + let frame_tx_clone = frame_tx.clone(); + let client_count_clone = client_count.clone(); + let client_slot_clone = client_slot.clone(); + let notify_clone = client_notify.clone(); + let screencasting_clone = screencasting.clone(); + let cdp_session_clone = cdp_session_id.clone(); + + let vw_clone = viewport_width.clone(); + let vh_clone = viewport_height.clone(); + let dashboard_dir_clone = dashboard_dir.clone(); + let last_tabs_clone = last_tabs.clone(); + let last_engine_clone = last_engine.clone(); + let last_frame_clone = last_frame.clone(); + let recording_clone = recording.clone(); + let accept_shutdown_rx = shutdown_rx.clone(); + let session_name_clone = session_id.clone(); + let accept_task = tokio::spawn(async move { + websocket::accept_loop( + listener, + frame_tx_clone, + client_count_clone, + client_slot_clone, + notify_clone, + screencasting_clone, + cdp_session_clone, + vw_clone, + vh_clone, + dashboard_dir_clone, + last_tabs_clone, + last_engine_clone, + last_frame_clone, + recording_clone, + accept_shutdown_rx, + session_name_clone, + ) + .await; + }); + + let frame_tx_bg = frame_tx.clone(); + let client_slot_bg = client_slot.clone(); + let client_notify_bg = client_notify.clone(); + let screencasting_bg = screencasting.clone(); + let client_count_bg = client_count.clone(); + let cdp_session_bg = cdp_session_id.clone(); + let vw_bg = viewport_width.clone(); + let vh_bg = viewport_height.clone(); + let last_frame_bg = last_frame.clone(); + let last_tabs_bg = last_tabs.clone(); + let last_engine_bg = last_engine.clone(); + let recording_bg = recording.clone(); + let cdp_task = tokio::spawn(async move { + cdp_loop::cdp_event_loop( + frame_tx_bg, + client_slot_bg, + client_notify_bg, + screencasting_bg, + client_count_bg, + cdp_session_bg, + vw_bg, + vh_bg, + last_frame_bg, + last_tabs_bg, + last_engine_bg, + recording_bg, + shutdown_rx, + ) + .await; + }); + + Ok(( + Self { + port, + session_name: session_id, + frame_tx, + client_count, + client_slot: client_slot.clone(), + cdp_session_id, + client_notify, + screencasting, + viewport_width, + viewport_height, + dashboard_dir, + last_tabs, + last_engine, + last_frame, + recording, + shutdown_tx, + accept_task: Mutex::new(Some(accept_task)), + cdp_task: Mutex::new(Some(cdp_task)), + }, + client_slot, + )) + } + + pub fn port(&self) -> u16 { + self.port + } + + /// Broadcast a raw frame string (legacy). + pub fn broadcast_frame(&self, frame_json: &str) { + let s = frame_json.to_string(); + if let Ok(mut lf) = self.last_frame.try_write() { + *lf = Some(s.clone()); + } + let _ = self.frame_tx.send(s); + } + + /// Broadcast a screencast frame with structured metadata. + pub fn broadcast_screencast_frame(&self, base64_data: &str, metadata: &FrameMetadata) { + let msg = json!({ + "type": "frame", + "data": base64_data, + "metadata": { + "offsetTop": metadata.offset_top, + "pageScaleFactor": metadata.page_scale_factor, + "deviceWidth": metadata.device_width, + "deviceHeight": metadata.device_height, + "scrollOffsetX": metadata.scroll_offset_x, + "scrollOffsetY": metadata.scroll_offset_y, + "timestamp": metadata.timestamp, + } + }); + let s = msg.to_string(); + if let Ok(mut lf) = self.last_frame.try_write() { + *lf = Some(s.clone()); + } + let _ = self.frame_tx.send(s); + } + + /// Broadcast a status message to all connected clients. + pub async fn broadcast_status( + &self, + connected: bool, + screencasting: bool, + viewport_width: u32, + viewport_height: u32, + engine: &str, + ) { + { + let mut guard = self.last_engine.write().await; + *guard = engine.to_string(); + } + let rec = *self.recording.lock().await; + let msg = json!({ + "type": "status", + "connected": connected, + "screencasting": screencasting, + "viewportWidth": viewport_width, + "viewportHeight": viewport_height, + "engine": engine, + "recording": rec, + }); + let _ = self.frame_tx.send(msg.to_string()); + } + + /// Broadcast an error message to all connected clients. + pub fn broadcast_error(&self, message: &str) { + let msg = json!({ + "type": "error", + "message": message, + }); + let _ = self.frame_tx.send(msg.to_string()); + } + + /// Broadcast a command event when a command begins executing. + pub fn broadcast_command(&self, action: &str, id: &str, params: &Value) { + let msg = json!({ + "type": "command", + "action": action, + "id": id, + "params": params, + "timestamp": timestamp_ms(), + }); + let _ = self.frame_tx.send(msg.to_string()); + } + + /// Broadcast a result event after a command finishes executing. + pub fn broadcast_result( + &self, + id: &str, + action: &str, + success: bool, + data: &Value, + duration_ms: u64, + ) { + let msg = json!({ + "type": "result", + "id": id, + "action": action, + "success": success, + "data": data, + "duration_ms": duration_ms, + "timestamp": timestamp_ms(), + }); + let _ = self.frame_tx.send(msg.to_string()); + } + + /// Broadcast a console event from the browser. + pub fn broadcast_console(&self, level: &str, text: &str, args: &[Value]) { + let mut msg = json!({ + "type": "console", + "level": level, + "text": text, + "timestamp": timestamp_ms(), + }); + if !args.is_empty() { + msg.as_object_mut() + .unwrap() + .insert("args".to_string(), Value::Array(args.to_vec())); + } + let _ = self.frame_tx.send(msg.to_string()); + } + + /// Broadcast a page error (uncaught exception) from the browser. + pub fn broadcast_page_error(&self, text: &str, line: Option, column: Option) { + let msg = json!({ + "type": "page_error", + "text": text, + "line": line, + "column": column, + "timestamp": timestamp_ms(), + }); + let _ = self.frame_tx.send(msg.to_string()); + } + + /// Broadcast the current tab list so the dashboard can render a tab bar. + /// Also caches the list so newly connected WebSocket clients receive it immediately. + pub async fn broadcast_tabs(&self, tabs: &[Value]) { + { + let mut guard = self.last_tabs.write().await; + *guard = tabs.to_vec(); + } + let msg = json!({ + "type": "tabs", + "tabs": tabs, + "timestamp": timestamp_ms(), + }); + let _ = self.frame_tx.send(msg.to_string()); + } + + /// Whether the dashboard directory is available. + pub fn has_dashboard(&self) -> bool { + self.dashboard_dir.is_some() + } +} + +pub(crate) fn timestamp_ms() -> u64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_millis() as u64) + .unwrap_or(0) +} + +pub fn is_allowed_origin(origin: Option<&str>) -> bool { + match origin { + None => true, + Some(o) => { + if o.starts_with("file://") { + return true; + } + if let Ok(url) = url::Url::parse(o) { + let host = url.host_str().unwrap_or(""); + host == "localhost" || host == "127.0.0.1" || host == "::1" || host == "[::1]" + } else { + false + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_allowed_origin_none() { + assert!(is_allowed_origin(None)); + } + + #[test] + fn test_allowed_origin_file() { + assert!(is_allowed_origin(Some("file:///path/to/file"))); + } + + #[test] + fn test_allowed_origin_localhost() { + assert!(is_allowed_origin(Some("http://localhost:3000"))); + assert!(is_allowed_origin(Some("http://127.0.0.1:8080"))); + } + + #[test] + fn test_disallowed_origin() { + assert!(!is_allowed_origin(Some("http://evil.com"))); + } + + #[test] + fn test_frame_metadata_default() { + let meta = FrameMetadata::default(); + assert_eq!(meta.device_width, 1280); + assert_eq!(meta.device_height, 720); + assert_eq!(meta.page_scale_factor, 1.0); + } +} diff --git a/cli/src/native/stream/websocket.rs b/cli/src/native/stream/websocket.rs new file mode 100644 index 0000000..920681e --- /dev/null +++ b/cli/src/native/stream/websocket.rs @@ -0,0 +1,352 @@ +use serde_json::{json, Value}; +use std::net::SocketAddr; +use std::path::PathBuf; +use std::sync::Arc; + +use futures_util::{SinkExt, StreamExt}; +use tokio::net::TcpListener; +use tokio::sync::{broadcast, watch, Mutex, Notify, RwLock}; +use tokio_tungstenite::tungstenite::Message; + +use crate::native::cdp::client::CdpClient; + +use super::http::handle_http_request; +use super::{is_allowed_origin, timestamp_ms}; + +#[allow(clippy::too_many_arguments)] +pub(super) async fn accept_loop( + listener: TcpListener, + frame_tx: broadcast::Sender, + client_count: Arc>, + client_slot: Arc>>>, + client_notify: Arc, + screencasting: Arc>, + cdp_session_id: Arc>>, + viewport_width: Arc>, + viewport_height: Arc>, + dashboard_dir: Option, + last_tabs: Arc>>, + last_engine: Arc>, + last_frame: Arc>>, + recording: Arc>, + mut shutdown_rx: watch::Receiver, + session_name: String, +) { + let dashboard_dir = dashboard_dir.map(Arc::from); + let session_name: Arc = Arc::from(session_name); + loop { + tokio::select! { + changed = shutdown_rx.changed() => { + if changed.is_err() || *shutdown_rx.borrow() { + break; + } + } + accept_result = listener.accept() => { + let Ok((stream, addr)) = accept_result else { + break; + }; + let frame_tx = frame_tx.clone(); + let client_count = client_count.clone(); + let client_slot = client_slot.clone(); + let client_notify = client_notify.clone(); + let screencasting = screencasting.clone(); + let cdp_session_id = cdp_session_id.clone(); + let vw = viewport_width.clone(); + let vh = viewport_height.clone(); + let dd = dashboard_dir.clone(); + let lt = last_tabs.clone(); + let le = last_engine.clone(); + let lf = last_frame.clone(); + let rec = recording.clone(); + let shutdown_rx = shutdown_rx.clone(); + let sn = session_name.clone(); + + tokio::spawn(async move { + handle_connection( + stream, + addr, + frame_tx, + client_count, + client_slot, + client_notify, + screencasting, + cdp_session_id, + vw, + vh, + dd, + lt, + le, + lf, + rec, + shutdown_rx, + sn, + ) + .await; + }); + } + } + } +} + +fn is_websocket_upgrade(request: &str) -> bool { + request.lines().any(|line| { + if let Some((name, value)) = line.split_once(':') { + name.trim().eq_ignore_ascii_case("upgrade") + && value.trim().eq_ignore_ascii_case("websocket") + } else { + false + } + }) +} + +/// Peek at the TCP stream to dispatch between WebSocket upgrade and plain HTTP. +#[allow(clippy::too_many_arguments)] +async fn handle_connection( + stream: tokio::net::TcpStream, + addr: SocketAddr, + frame_tx: broadcast::Sender, + client_count: Arc>, + client_slot: Arc>>>, + client_notify: Arc, + screencasting: Arc>, + cdp_session_id: Arc>>, + viewport_width: Arc>, + viewport_height: Arc>, + dashboard_dir: Option>, + last_tabs: Arc>>, + last_engine: Arc>, + last_frame: Arc>>, + recording: Arc>, + shutdown_rx: watch::Receiver, + session_name: Arc, +) { + let mut buf = [0u8; 4096]; + let n = match stream.peek(&mut buf).await { + Ok(n) => n, + Err(_) => return, + }; + let request = String::from_utf8_lossy(&buf[..n]); + + if is_websocket_upgrade(&request) { + let frame_rx = frame_tx.subscribe(); + handle_ws_client( + stream, + addr, + frame_rx, + client_count, + client_slot, + client_notify, + screencasting, + cdp_session_id, + viewport_width, + viewport_height, + last_tabs, + last_engine, + last_frame, + recording, + shutdown_rx, + ) + .await; + } else { + handle_http_request( + stream, + &buf[..n], + dashboard_dir.as_deref().map(|p| p.as_path()), + &last_tabs, + &last_engine, + &session_name, + ) + .await; + } +} + +#[allow(clippy::result_large_err, clippy::too_many_arguments)] +async fn handle_ws_client( + stream: tokio::net::TcpStream, + _addr: SocketAddr, + mut frame_rx: broadcast::Receiver, + client_count: Arc>, + client_slot: Arc>>>, + client_notify: Arc, + screencasting: Arc>, + cdp_session_id: Arc>>, + viewport_width: Arc>, + viewport_height: Arc>, + last_tabs: Arc>>, + last_engine: Arc>, + last_frame: Arc>>, + recording: Arc>, + mut shutdown_rx: watch::Receiver, +) { + let callback = + |req: &tokio_tungstenite::tungstenite::handshake::server::Request, + resp: tokio_tungstenite::tungstenite::handshake::server::Response| { + let origin = req + .headers() + .get("origin") + .and_then(|v| v.to_str().ok()) + .map(|s| s.to_string()); + if !is_allowed_origin(origin.as_deref()) { + let mut reject = + tokio_tungstenite::tungstenite::handshake::server::ErrorResponse::new(Some( + "Origin not allowed".to_string(), + )); + *reject.status_mut() = tokio_tungstenite::tungstenite::http::StatusCode::FORBIDDEN; + return Err(reject); + } + Ok(resp) + }; + + let ws_stream = match tokio_tungstenite::accept_hdr_async(stream, callback).await { + Ok(ws) => ws, + Err(_) => return, + }; + + { + let mut count = client_count.lock().await; + *count += 1; + } + + let (mut ws_tx, mut ws_rx) = ws_stream.split(); + + { + let guard = client_slot.read().await; + let connected = guard.is_some(); + let sc = *screencasting.lock().await; + let vw = *viewport_width.lock().await; + let vh = *viewport_height.lock().await; + let eng = last_engine.read().await.clone(); + let rec = *recording.lock().await; + let status = json!({ + "type": "status", + "connected": connected, + "screencasting": sc, + "viewportWidth": vw, + "viewportHeight": vh, + "engine": eng, + "recording": rec, + }); + let _ = ws_tx.send(Message::Text(status.to_string())).await; + + let tabs = last_tabs.read().await; + if !tabs.is_empty() { + let tabs_msg = json!({ + "type": "tabs", + "tabs": *tabs, + "timestamp": timestamp_ms(), + }); + let _ = ws_tx.send(Message::Text(tabs_msg.to_string())).await; + } + + if let Some(ref cached) = *last_frame.read().await { + let _ = ws_tx.send(Message::Text(cached.clone())).await; + } + } + + client_notify.notify_one(); + + loop { + tokio::select! { + changed = shutdown_rx.changed() => { + if changed.is_err() || *shutdown_rx.borrow() { + let _ = ws_tx.send(Message::Close(None)).await; + break; + } + } + frame = frame_rx.recv() => { + match frame { + Ok(data) => { + if ws_tx.send(Message::Text(data)).await.is_err() { + break; + } + } + Err(broadcast::error::RecvError::Lagged(_)) => { + continue; + } + Err(broadcast::error::RecvError::Closed) => break, + } + } + msg = ws_rx.next() => { + match msg { + Some(Ok(Message::Text(text))) => { + let guard = client_slot.read().await; + if let Some(ref client) = *guard { + let sid = cdp_session_id.read().await; + handle_client_message(&text, client.as_ref(), sid.as_deref()).await; + } + } + Some(Ok(Message::Close(_))) | None => break, + _ => {} + } + } + } + } + + { + let mut count = client_count.lock().await; + *count = count.saturating_sub(1); + } + + client_notify.notify_one(); +} + +async fn handle_client_message(msg: &str, client: &CdpClient, session_id: Option<&str>) { + let parsed: Value = match serde_json::from_str(msg) { + Ok(v) => v, + Err(_) => return, + }; + + let msg_type = parsed.get("type").and_then(|v| v.as_str()).unwrap_or(""); + + match msg_type { + "input_mouse" => { + let _ = client + .send_command( + "Input.dispatchMouseEvent", + Some(json!({ + "type": parsed.get("eventType").and_then(|v| v.as_str()).unwrap_or("mouseMoved"), + "x": parsed.get("x").and_then(|v| v.as_f64()).unwrap_or(0.0), + "y": parsed.get("y").and_then(|v| v.as_f64()).unwrap_or(0.0), + "button": parsed.get("button").and_then(|v| v.as_str()).unwrap_or("none"), + "clickCount": parsed.get("clickCount").and_then(|v| v.as_i64()).unwrap_or(0), + "deltaX": parsed.get("deltaX").and_then(|v| v.as_f64()).unwrap_or(0.0), + "deltaY": parsed.get("deltaY").and_then(|v| v.as_f64()).unwrap_or(0.0), + "modifiers": parsed.get("modifiers").and_then(|v| v.as_i64()).unwrap_or(0), + })), + session_id, + ) + .await; + } + "input_keyboard" => { + let _ = client + .send_command( + "Input.dispatchKeyEvent", + Some(json!({ + "type": parsed.get("eventType").and_then(|v| v.as_str()).unwrap_or("keyDown"), + "key": parsed.get("key"), + "code": parsed.get("code"), + "text": parsed.get("text"), + "windowsVirtualKeyCode": parsed.get("windowsVirtualKeyCode").and_then(|v| v.as_i64()).unwrap_or(0), + "modifiers": parsed.get("modifiers").and_then(|v| v.as_i64()).unwrap_or(0), + })), + session_id, + ) + .await; + } + "input_touch" => { + let _ = client + .send_command( + "Input.dispatchTouchEvent", + Some(json!({ + "type": parsed.get("eventType").and_then(|v| v.as_str()).unwrap_or("touchStart"), + "touchPoints": parsed.get("touchPoints").unwrap_or(&json!([])), + "modifiers": parsed.get("modifiers").and_then(|v| v.as_i64()).unwrap_or(0), + })), + session_id, + ) + .await; + } + "status" => {} + _ => {} + } +} diff --git a/cli/src/output.rs b/cli/src/output.rs index 770fad7..4f17720 100644 --- a/cli/src/output.rs +++ b/cli/src/output.rs @@ -1544,6 +1544,7 @@ Designed for AI agents to understand page structure. Options: -i, --interactive Only include interactive elements + -u, --urls Include href URLs for link elements -c, --compact Remove empty structural elements -d, --depth Limit tree depth -s, --selector Scope snapshot to CSS selector @@ -1555,6 +1556,7 @@ Global Options: Examples: agent-browser snapshot agent-browser snapshot -i + agent-browser snapshot -i --urls agent-browser snapshot --compact --depth 5 agent-browser snapshot -s "#main-content" "## @@ -2622,20 +2624,24 @@ Examples: "batch" => { r##" -agent-browser batch - Execute multiple commands from stdin +agent-browser batch - Execute multiple commands sequentially -Usage: echo '' | agent-browser batch [options] +Usage: agent-browser batch [options] "" "" ... + echo '' | agent-browser batch [options] -Reads a JSON array of commands from stdin and executes them sequentially. -Each command is an array of strings matching normal CLI arguments. -Results are printed in order, separated by blank lines (or as a JSON array -with --json). +Runs multiple commands in sequence. Commands can be passed as quoted +arguments or piped as JSON via stdin. Results are printed in order, +separated by blank lines (or as a JSON array with --json). Options: --bail Stop on first error (default: continue all commands) --json Output results as a JSON array -Input Format: +Argument Mode: + Each quoted argument is a full command string: + agent-browser batch "open https://example.com" "snapshot -i" "screenshot" + +Stdin Mode (JSON): A JSON array of string arrays. Each inner array is one command: [ ["open", "https://example.com"], @@ -2646,8 +2652,9 @@ Input Format: ] Examples: + agent-browser batch "open https://example.com" "screenshot" + agent-browser batch --bail "open https://example.com" "click @e1" "screenshot" echo '[["open", "https://example.com"], ["snapshot"]]' | agent-browser batch - echo '[["open", "https://example.com"], ["get", "title"]]' | agent-browser batch --json agent-browser batch --bail < commands.json "## } @@ -2769,8 +2776,8 @@ Streaming: stream status Show streaming status and active port Batch: - batch [--bail] Execute commands from stdin (JSON array of string arrays) - --bail stops on first error (default: continue all) + batch [--bail] ["cmd" ...] Execute multiple commands sequentially (args or stdin) + --bail stops on first error (default: continue all) Auth Vault: auth save [opts] Save auth profile (--url, --username, --password/--password-stdin) @@ -2912,6 +2919,9 @@ Environment: AGENT_BROWSER_SCREENSHOT_DIR Default screenshot output directory AGENT_BROWSER_SCREENSHOT_QUALITY JPEG quality 0-100 AGENT_BROWSER_SCREENSHOT_FORMAT Screenshot format: png, jpeg + AI_GATEWAY_URL Vercel AI Gateway base URL (default: https://ai-gateway.vercel.sh) + AI_GATEWAY_API_KEY API key for the AI Gateway (enables dashboard AI chat) + AI_GATEWAY_MODEL Default AI model (default: anthropic/claude-sonnet-4.6) Install: npm install -g agent-browser # npm @@ -2928,7 +2938,7 @@ Examples: agent-browser get text @e1 agent-browser screenshot --full agent-browser screenshot --annotate # Labeled screenshot for vision models - agent-browser wait --load networkidle # Wait for slow pages to load + agent-browser wait 2000 # Wait for slow pages to settle agent-browser --cdp 9222 snapshot # Connect via CDP port agent-browser --auto-connect snapshot # Auto-discover running Chrome agent-browser stream enable # Start runtime streaming on an auto-selected port @@ -2942,9 +2952,9 @@ Examples: Command Chaining: Chain commands with && in a single shell call (browser persists via daemon): - agent-browser open example.com && agent-browser wait --load networkidle && agent-browser snapshot -i + agent-browser open example.com && agent-browser snapshot -i agent-browser fill @e1 "user@example.com" && agent-browser fill @e2 "pass" && agent-browser click @e3 - agent-browser open example.com && agent-browser wait --load networkidle && agent-browser screenshot page.png + agent-browser open example.com && agent-browser screenshot iOS Simulator (requires Xcode and Appium): agent-browser -p ios open example.com # Use default iPhone diff --git a/docs/src/app/api/docs-chat/route.ts b/docs/src/app/api/docs-chat/route.ts index 4f7f062..c95b533 100644 --- a/docs/src/app/api/docs-chat/route.ts +++ b/docs/src/app/api/docs-chat/route.ts @@ -10,7 +10,7 @@ import { minuteRateLimit, dailyRateLimit } from "@/lib/rate-limit"; export const maxDuration = 60; -const DEFAULT_MODEL = "anthropic/claude-haiku-4.5"; +const DEFAULT_MODEL = "anthropic/claude-sonnet-4.6"; const SYSTEM_PROMPT = `You are a helpful documentation assistant for agent-browser, a browser automation CLI designed for AI agents. diff --git a/docs/src/app/commands/page.mdx b/docs/src/app/commands/page.mdx index 25b7868..54ef73f 100644 --- a/docs/src/app/commands/page.mdx +++ b/docs/src/app/commands/page.mdx @@ -394,18 +394,22 @@ agent-browser reload # Reload page ## Batch execution -Execute multiple commands in a single invocation by piping a JSON array of string arrays to `batch`: +Execute multiple commands in a single invocation. Commands can be passed as quoted arguments or piped as JSON via stdin. ```bash +# Argument mode: each quoted argument is a full command +agent-browser batch "open https://example.com" "snapshot -i" "screenshot" + +# With --bail to stop on first error +agent-browser batch --bail "open https://example.com" "click @e1" "screenshot" + +# Stdin mode: pipe commands as JSON echo '[ ["open", "https://example.com"], ["snapshot", "-i"], ["click", "@e1"], ["screenshot", "result.png"] ]' | agent-browser batch --json - -# Stop on first error -agent-browser batch --bail < commands.json ``` diff --git a/docs/src/app/configuration/page.mdx b/docs/src/app/configuration/page.mdx index 334f402..1b5485b 100644 --- a/docs/src/app/configuration/page.mdx +++ b/docs/src/app/configuration/page.mdx @@ -187,6 +187,9 @@ These environment variables configure additional daemon and runtime behavior: + + +
AGENT_BROWSER_CONFIRM_INTERACTIVEEnable interactive confirmation prompts (auto-denies if stdin is not a TTY).(disabled)
AGENT_BROWSER_ENGINEBrowser engine to use: chrome (default), lightpanda.chrome
AGENT_BROWSER_NO_AUTO_DIALOGDisable automatic dismissal of alert/beforeunload dialogs.(disabled)
AI_GATEWAY_URLVercel AI Gateway base URL.https://ai-gateway.vercel.sh
AI_GATEWAY_API_KEYAPI key for the Vercel AI Gateway. Required to enable AI chat.(none)
AI_GATEWAY_MODELDefault AI model for dashboard chat.anthropic/claude-sonnet-4.6
diff --git a/docs/src/app/dashboard/page.mdx b/docs/src/app/dashboard/page.mdx index 7f60a1c..0026383 100644 --- a/docs/src/app/dashboard/page.mdx +++ b/docs/src/app/dashboard/page.mdx @@ -135,3 +135,38 @@ pnpm build:dashboard The built files are served by the daemon's stream server on the same port used for WebSocket connections. Plain HTTP requests serve the dashboard, while WebSocket upgrade requests are handled as before. When the dashboard is not installed, visiting the HTTP endpoint shows instructions to run `agent-browser dashboard install`. + +## AI Chat + +The dashboard includes an optional AI chat panel powered by the [Vercel AI Gateway](https://vercel.com/docs/ai-gateway). When enabled, a **Chat** tab appears in the right pane alongside Activity, Console, Network, Storage, and Extensions. + +### Setup + +The Chat tab is always visible. Set the API key to enable responses: + +```bash +export AI_GATEWAY_API_KEY=gw_your_key_here +agent-browser dashboard start +``` + +Optionally override the gateway URL or model: + +```bash +export AI_GATEWAY_URL=https://ai-gateway.vercel.sh # this is the default +export AI_GATEWAY_MODEL=openai/gpt-4o-mini # default: anthropic/claude-sonnet-4.6 +``` + +### How it works + +The Rust server proxies chat requests from the dashboard to the Vercel AI Gateway and streams responses back using the Vercel AI SDK's UI Message Stream protocol. The dashboard frontend uses `useChat` from `@ai-sdk/react` with `DefaultChatTransport`. + + + + + + + + + + +
VariableDescriptionDefault
AI_GATEWAY_URLVercel AI Gateway base URL.https://ai-gateway.vercel.sh
AI_GATEWAY_API_KEYAPI key for the AI Gateway. Required to enable AI chat responses.(none)
AI_GATEWAY_MODELDefault AI model for chat requests.anthropic/claude-sonnet-4.6
diff --git a/docs/src/app/snapshots/page.mdx b/docs/src/app/snapshots/page.mdx index b5581e7..6c7ff69 100644 --- a/docs/src/app/snapshots/page.mdx +++ b/docs/src/app/snapshots/page.mdx @@ -21,6 +21,7 @@ agent-browser snapshot -i -c -d 5 # Combine options -i, --interactiveOnly interactive elements (buttons, links, inputs) + -u, --urlsInclude href URLs for link elements -c, --compactRemove empty structural elements -d, --depthLimit tree depth -s, --selectorScope to CSS selector diff --git a/docs/src/components/header.tsx b/docs/src/components/header.tsx index a6ff4b3..7c9240e 100644 --- a/docs/src/components/header.tsx +++ b/docs/src/components/header.tsx @@ -68,7 +68,7 @@ export function Header() { > - 25k + 27k + - - {children} - + + + {children} + + ); diff --git a/packages/dashboard/src/app/page.tsx b/packages/dashboard/src/app/page.tsx index 84bb6d6..25a240b 100644 --- a/packages/dashboard/src/app/page.tsx +++ b/packages/dashboard/src/app/page.tsx @@ -1,14 +1,16 @@ "use client"; -import { useAtomValue } from "jotai/react"; -import { activePortAtom } from "@/store/sessions"; +import { useAtomValue, useSetAtom } from "jotai/react"; +import { activePortAtom, sessionsAtom, newSessionDialogAtom } from "@/store/sessions"; import { useSessionsSync } from "@/store/sessions"; import { useStreamSync, hasConsoleErrorsAtom, consoleLogsAtom } from "@/store/stream"; import { useActivitySync } from "@/store/activity"; import { activeExtensionsAtom } from "@/store/sessions"; +import { useChatStatusSync } from "@/store/chat"; import { useMediaQuery } from "@/hooks/use-media-query"; import { Viewport } from "@/components/viewport"; import { ActivityFeed } from "@/components/activity-feed"; +import { ChatPanel } from "@/components/chat-panel"; import { ConsolePanel } from "@/components/console-panel"; import { StoragePanel } from "@/components/storage-panel"; import { ExtensionsPanel } from "@/components/extensions-panel"; @@ -20,21 +22,28 @@ import { ResizableHandle, } from "@/components/ui/resizable"; import { Tabs, TabsList, TabsTrigger, TabsContent } from "@/components/ui/tabs"; +import { Button } from "@/components/ui/button"; +import { Plus } from "lucide-react"; export default function DashboardPage() { const activePort = useAtomValue(activePortAtom); useStreamSync(activePort); useSessionsSync(); useActivitySync(); + useChatStatusSync(); + const sessions = useAtomValue(sessionsAtom); + const hasSessions = sessions.length > 0; + const setNewSessionDialog = useSetAtom(newSessionDialogAtom); const isDesktop = useMediaQuery("(min-width: 768px)"); const hasConsoleErrors = useAtomValue(hasConsoleErrorsAtom); const activeExtensions = useAtomValue(activeExtensionsAtom); const sidePanel = ( - +
+ Chat Activity Console @@ -67,10 +76,46 @@ export default function DashboardPage() { + + + ); if (isDesktop) { + if (!hasSessions) { + return ( +
+ + + + + + +
+
+
+

No active sessions

+

Create a session to get started

+
+ +
+
+
+
+
+ ); + } + return (
& ExtraProps; +type MdHeadingProps = React.HTMLAttributes & ExtraProps; +type MdAnchorProps = React.AnchorHTMLAttributes & ExtraProps; +type MdPreProps = React.HTMLAttributes & ExtraProps; +type MdCodeProps = React.HTMLAttributes & ExtraProps; + +const chatComponents = { + img: ({ node: _node, src, alt, ...props }: MdImgProps) => { + if (typeof src === "string" && src.startsWith("data:image/")) { + return {alt}; + } + return null; + }, + h1: ({ node: _node, ...props }: MdHeadingProps) =>

, + h2: ({ node: _node, ...props }: MdHeadingProps) =>

, + h3: ({ node: _node, ...props }: MdHeadingProps) =>

, + h4: ({ node: _node, ...props }: MdHeadingProps) =>

, + h5: ({ node: _node, ...props }: MdHeadingProps) =>

, + h6: ({ node: _node, ...props }: MdHeadingProps) =>

, + a: ({ node: _node, href, children, ...props }: MdAnchorProps) => ( + + {children} + + ), + pre: ({ node: _node, ...props }: MdPreProps) => ( +

+  ),
+  code: ({ className, children, node: _node, ...props }: MdCodeProps) => {
+    if (className?.includes("language-")) {
+      return {children};
+    }
+    return (
+      
+        {children}
+      
+    );
+  },
+};
+
+const STORAGE_PREFIX = "dashboard-chat-";
+const IMAGE_DATA_URL_RE = /data:image\/[^;]+;base64,[A-Za-z0-9+/=]+/g;
+
+function stripImagesForStorage(messages: unknown[]): unknown[] {
+  const json = JSON.stringify(messages);
+  return JSON.parse(json.replace(IMAGE_DATA_URL_RE, "[image stripped]"));
+}
+
+const SUGGESTIONS = [
+  "Go to vercel.com",
+  "Take a screenshot",
+  "What's on the page?",
+  "Click the first link",
+];
+
+interface ToolInvocationPart {
+  type: string;
+  toolCallId: string;
+  state: string;
+  input?: Record;
+  output?: unknown;
+}
+
+function isToolPart(part: { type: string }): part is ToolInvocationPart {
+  return part.type.startsWith("tool-");
+}
+
+function truncateOutput(text: string, maxLines = 30): string {
+  const lines = text.split("\n");
+  if (lines.length <= maxLines) return text;
+  return lines.slice(0, maxLines).join("\n") + `\n... (${lines.length - maxLines} more lines)`;
+}
+
+function parseOutputObject(raw: unknown): Record | null {
+  if (typeof raw === "string") {
+    try {
+      const parsed = JSON.parse(raw);
+      if (typeof parsed === "object" && parsed !== null) return parsed;
+    } catch { /* not JSON */ }
+    return null;
+  }
+  if (typeof raw === "object" && raw !== null) return raw as Record;
+  return null;
+}
+
+function formatOutput(raw: unknown): string | null {
+  if (typeof raw === "string") {
+    if (!raw.trim()) return null;
+    const obj = parseOutputObject(raw);
+    if (obj) {
+      if (typeof obj.text === "string" && obj.image) return obj.text as string;
+      const { image: _, ...rest } = obj;
+      return JSON.stringify(rest, null, 2);
+    }
+    return raw;
+  }
+  if (typeof raw === "object" && raw !== null) {
+    const r = raw as Record;
+    if (typeof r.text === "string" && r.image) return r.text as string;
+    const { image: _, ...rest } = r;
+    return JSON.stringify(rest, null, 2);
+  }
+  return null;
+}
+
+function extractImageUrl(raw: unknown): string | null {
+  const obj = parseOutputObject(raw);
+  if (!obj) return null;
+  const img = obj.image;
+  if (typeof img === "string" && img.startsWith("data:image/")) return img;
+  return null;
+}
+
+function ToolCallBlock({ part, onImageLoad }: { part: ToolInvocationPart; onImageLoad?: () => void }) {
+  const [expanded, setExpanded] = useState(false);
+  const toolName = part.type.split("-").slice(1).join("-");
+  const command = (part.input as { command?: string })?.command ?? toolName;
+  const isDone = part.state === "output-available";
+  const isRunning = !isDone;
+  const output = isDone ? formatOutput(part.output) : null;
+  const hasOutput = !!output;
+  const imageUrl = isDone ? extractImageUrl(part.output) : null;
+  const canExpand = hasOutput && !isRunning;
+
+  return (
+    
+
canExpand && setExpanded(!expanded)} + > +
+ {isRunning ? ( + + ) : ( + + )} + {command} +
+ {expanded && hasOutput && ( +
+
+              {truncateOutput(output)}
+            
+
+ )} +
+ {imageUrl && ( + Screenshot + )} +
+ ); +} + +const DEFAULT_CONTEXT_WINDOW = 128000; + +function estimateTokens(text: string): number { + return Math.ceil(text.length / 4); +} + +function formatTokenCount(n: number): string { + if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M`; + if (n >= 1_000) return `${(n / 1_000).toFixed(0)}K`; + return `${n}`; +} + +function ContextMeter({ used, total }: { used: number; total: number }) { + const ratio = Math.min(used / total, 1); + const size = 16; + const strokeWidth = 2; + const r = (size - strokeWidth) / 2; + const circumference = 2 * Math.PI * r; + const offset = circumference * (1 - ratio); + const color = + ratio > 0.9 ? "text-destructive" : ratio > 0.7 ? "text-yellow-500" : "text-muted-foreground/50"; + + return ( +
+ + + + +
+ ); +} + +const DEFAULT_MODEL = "anthropic/claude-sonnet-4.6"; + +function useTimeAgo(ts: number | undefined) { + const [, setTick] = useState(0); + useEffect(() => { + if (!ts) return; + const id = setInterval(() => setTick((t) => t + 1), 30_000); + return () => clearInterval(id); + }, [ts]); + if (!ts) return ""; + const diff = Math.floor((Date.now() - ts) / 1000); + if (diff < 5) return "just now"; + if (diff < 60) return `${diff}s ago`; + const mins = Math.floor(diff / 60); + if (mins < 60) return `${mins}m ago`; + const hrs = Math.floor(mins / 60); + return `${hrs}h ago`; +} + +function MessageFooter({ model, timestamp, text }: { model: string; timestamp?: number; text: string }) { + const [copied, setCopied] = useState(false); + const timeAgo = useTimeAgo(timestamp); + const shortModel = model.includes("/") ? model.split("/").pop()! : model; + + const handleCopy = useCallback(() => { + navigator.clipboard.writeText(text).then(() => { + setCopied(true); + setTimeout(() => setCopied(false), 2000); + }); + }, [text]); + + return ( +
+ {shortModel} + {timeAgo && ( + <> + · + {timeAgo} + + )} + +
+ ); +} + +interface PendingImage { + file: File; + preview: string; +} + +export function ChatPanel() { + const [input, setInput] = useState(""); + const [errorDismissed, setErrorDismissed] = useState(false); + const [pendingImages, setPendingImages] = useState([]); + const fileInputRef = useRef(null); + const defaultModel = useAtomValue(chatModelAtom); + const [selectedModel, setSelectedModel] = useState(defaultModel || DEFAULT_MODEL); + const messagesEndRef = useRef(null); + const inputRef = useRef(null); + const sessionName = useAtomValue(activeSessionNameAtom); + const chatId = sessionName || "default"; + const storageKey = `${STORAGE_PREFIX}${chatId}`; + const sessionRef = useRef(chatId); + sessionRef.current = chatId; + const modelRef = useRef(selectedModel); + modelRef.current = selectedModel; + const messageTimestamps = useRef>({}); + + useEffect(() => { + if (defaultModel) setSelectedModel(defaultModel); + }, [defaultModel]); + + const transport = useRef( + new DefaultChatTransport({ + api: getChatApiUrl(), + body: () => ({ + session: sessionRef.current, + model: modelRef.current, + }), + }), + ).current; + + const { messages, sendMessage, stop, status, setMessages, error } = useChat({ + chatId, + transport, + onError: () => setErrorDismissed(false), + }); + + const visibleError = error && !errorDismissed ? error : undefined; + const isLoading = status === "streaming" || status === "submitted"; + const hasMessages = messages.length > 0 || !!visibleError; + + useEffect(() => { + for (const msg of messages) { + if (msg.role === "assistant" && !messageTimestamps.current[msg.id]) { + messageTimestamps.current[msg.id] = Date.now(); + } + } + }, [messages]); + + const models = useAtomValue(availableModelsAtom); + const estimatedTokens = useMemo(() => { + let total = 0; + for (const msg of messages) { + for (const part of msg.parts) { + if (part.type === "text") total += estimateTokens(part.text); + else if (isToolPart(part)) { + if (part.input) total += estimateTokens(JSON.stringify(part.input)); + if (part.output) { + const raw = typeof part.output === "string" ? part.output : JSON.stringify(part.output); + const stripped = raw.replace(/"image"\s*:\s*"data:[^"]*"/g, '"image":"[omitted]"'); + total += estimateTokens(stripped); + } + } + } + } + return total; + }, [messages]); + const contextWindow = useMemo(() => { + const match = models.find((m) => m.id === selectedModel); + return match?.context_window ?? DEFAULT_CONTEXT_WINDOW; + }, [models, selectedModel]); + + useEffect(() => { + inputRef.current?.focus(); + }, []); + + const scrollToBottom = useCallback(() => { + messagesEndRef.current?.scrollIntoView({ behavior: "smooth" }); + }, []); + + useEffect(() => { + scrollToBottom(); + }, [messages, visibleError, scrollToBottom]); + + // Restore messages from localStorage when chatId changes + useEffect(() => { + try { + const stored = localStorage.getItem(storageKey); + if (stored) { + const parsed = JSON.parse(stored); + if (Array.isArray(parsed) && parsed.length > 0) { + setMessages(parsed); + return; + } + } + } catch { + // ignore + } + setMessages([]); + }, [chatId, storageKey, setMessages]); + + // Persist messages to localStorage (strip base64 images to save space) + useEffect(() => { + if (isLoading) return; + if (messages.length === 0) { + localStorage.removeItem(storageKey); + return; + } + try { + localStorage.setItem(storageKey, JSON.stringify(stripImagesForStorage(messages))); + } catch { + // ignore quota + } + }, [messages, isLoading, storageKey]); + + const addImages = useCallback((files: FileList | null) => { + if (!files) return; + const images = Array.from(files).filter((f) => f.type.startsWith("image/")); + setPendingImages((prev) => [ + ...prev, + ...images.map((file) => ({ file, preview: URL.createObjectURL(file) })), + ]); + }, []); + + const removeImage = useCallback((index: number) => { + setPendingImages((prev) => { + const next = [...prev]; + URL.revokeObjectURL(next[index].preview); + next.splice(index, 1); + return next; + }); + }, []); + + const handleSubmit = useCallback( + (e: React.FormEvent) => { + e.preventDefault(); + if ((!input.trim() && pendingImages.length === 0) || isLoading) return; + const dt = new DataTransfer(); + for (const img of pendingImages) dt.items.add(img.file); + const files = dt.files.length > 0 ? dt.files : undefined; + sendMessage({ text: input, files }); + setInput(""); + setPendingImages((prev) => { + for (const p of prev) URL.revokeObjectURL(p.preview); + return []; + }); + }, + [input, isLoading, sendMessage, pendingImages], + ); + + const lastCompactedId = useRef(null); + useEffect(() => { + if (isLoading || messages.length === 0) return; + const lastAssistant = [...messages].reverse().find((m) => m.role === "assistant"); + if (!lastAssistant) return; + if (lastAssistant.id === lastCompactedId.current) return; + const meta = (lastAssistant as any).metadata as + | { compacted?: boolean; summary?: string; keepLastN?: number } + | undefined; + if (!meta?.compacted || typeof meta.keepLastN !== "number") return; + + lastCompactedId.current = lastAssistant.id; + const keep = meta.keepLastN; + if (keep >= messages.length) return; + + const summaryMsg = { + id: `compaction-${Date.now()}`, + role: "assistant" as const, + parts: [ + { + type: "text" as const, + text: `*Earlier messages were summarized to stay within the context window.*`, + }, + ], + }; + + const kept = messages.slice(messages.length - keep); + setMessages([summaryMsg as any, ...kept]); + }, [isLoading, messages, setMessages]); + + const handleClear = useCallback(() => { + setMessages([]); + setErrorDismissed(true); + localStorage.removeItem(storageKey); + requestAnimationFrame(() => inputRef.current?.focus()); + }, [setMessages, storageKey]); + + const handleDownload = useCallback(() => { + const data = messages.map((msg) => ({ + id: msg.id, + role: msg.role, + parts: msg.parts.map((p) => { + if (p.type === "text") return { type: "text", text: p.text }; + if (p.type === "file") return { type: "file", filename: (p as any).filename }; + if (isToolPart(p)) { + const out = typeof p.output === "string" ? p.output : JSON.stringify(p.output); + const stripped = out?.replace(/"image":"data:[^"]*"/g, '"image":"[stripped]"'); + return { + type: p.type, + toolName: (p as any).toolName, + state: (p as any).state, + input: (p as any).input, + output: stripped, + }; + } + return { type: p.type }; + }), + })); + const json = JSON.stringify({ session: chatId, model: selectedModel, messages: data }, null, 2); + const blob = new Blob([json], { type: "application/json" }); + const url = URL.createObjectURL(blob); + const a = document.createElement("a"); + a.href = url; + a.download = `chat-${chatId}-${Date.now()}.json`; + a.click(); + URL.revokeObjectURL(url); + }, [messages, chatId, selectedModel]); + + const hasVisibleContent = (parts: (typeof messages)[number]["parts"]): boolean => { + return parts.some( + (p) => (p.type === "text" && p.text.length > 0) || p.type === "file" || isToolPart(p), + ); + }; + + return ( +
+ {hasMessages && ( +
+ + +
+ )} + + +
+ {!hasMessages && !isLoading && ( +
+

+ Control the browser with natural language: +

+
+ {SUGGESTIONS.map((s) => ( + + ))} +
+
+ )} + + {messages.map((message) => { + if (message.id.startsWith("compaction-")) { + return ( +
+
+ Earlier messages summarized +
+
+ ); + } + if (!hasVisibleContent(message.parts)) return null; + return ( +
+ {message.role === "user" ? ( +
+ {message.parts.some((p) => p.type === "file") && ( +
+ {message.parts + .filter((p): p is Extract => p.type === "file") + .map((p, i) => ( + {p.filename + ))} +
+ )} +
+ {message.parts + .filter((p): p is Extract => p.type === "text") + .map((p) => p.text) + .join("")} +
+
+ ) : ( +
+ {(() => { + type Group = { type: "tools" | "text"; items: (typeof message.parts)[number][] }; + const groups: Group[] = []; + for (const part of message.parts) { + const groupType = isToolPart(part) ? "tools" : "text"; + const last = groups[groups.length - 1]; + if (last && last.type === groupType) { + last.items.push(part); + } else { + groups.push({ type: groupType, items: [part] }); + } + } + + return groups.map((group, gi) => { + if (group.type === "tools") { + return ( +
+ {group.items.map((part) => { + if (!isToolPart(part)) return null; + return ; + })} +
+ ); + } + const combinedText = group.items + .filter((p): p is Extract => p.type === "text" && !!p.text) + .map((p) => p.text) + .join(""); + if (!combinedText) return null; + return ( +
+ + {combinedText} + +
+ ); + }); + })()} + {(() => { + const isLast = message === messages[messages.length - 1]; + const isComplete = !isLast || !isLoading; + if (!isComplete) return null; + const fullText = message.parts + .filter((p): p is Extract => p.type === "text" && !!p.text) + .map((p) => p.text) + .join(""); + return ( + + ); + })()} +
+ )} +
+ ); + })} + + {isLoading && messages.length > 0 && (() => { + const lastMsg = messages[messages.length - 1]; + const lastPart = lastMsg?.parts[lastMsg.parts.length - 1]; + const noVisibleContent = !lastMsg || !hasVisibleContent(lastMsg.parts); + const lastIsCompletedTool = lastPart && isToolPart(lastPart) && lastPart.state === "output-available"; + if (noVisibleContent || lastIsCompletedTool) { + return ( + + Working... + + ); + } + return null; + })()} + + {visibleError && ( +
+ {(() => { + try { + const parsed = JSON.parse(visibleError.message); + return parsed.message || parsed.error || visibleError.message; + } catch { + return visibleError.message || "Something went wrong."; + } + })()} +
+ )} + +
+
+ + +
+
+ {pendingImages.length > 0 && ( +
+ {pendingImages.map((img, i) => ( +
+ {img.file.name} + +
+ ))} +
+ )} +
+