1use anyhow::{Context, Result};
8use serde_json::{json, Value};
9use std::collections::HashMap;
10use std::sync::Arc;
11use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
12use tokio::sync::{Mutex, RwLock};
13
14use crate::cli::env_resolver::ResolvedBrowser;
15use crate::session::backend::TabBackend;
16
17pub type BidiCache = Arc<Mutex<Option<Arc<crate::bidi::BidiClient>>>>;
22
23#[derive(Clone)]
35pub struct ServerState {
36 pub browser: Arc<RwLock<ResolvedBrowser>>,
37 pub bidi: BidiCache,
38 pub bidi_lock: Arc<Mutex<BidiLockState>>,
45 pub backend: Arc<Mutex<Option<TabBackend>>>,
50 pub active_target_id: Arc<Mutex<Option<String>>>,
55 pub origin_target_ids: Arc<Mutex<HashMap<String, String>>>,
61 pub sidecar: Arc<Mutex<Option<crate::sidecar::Sidecar>>>,
68 pub sidecar_config: crate::sidecar::SidecarConfig,
71 pub op_barrier: Arc<RwLock<()>>,
81}
82
83#[derive(Default)]
87pub enum BidiLockState {
88 #[default]
89 Pending,
90 Acquired(crate::registry::BidiLockGuard),
91 NotApplicable,
92}
93
94impl ServerState {
95 pub fn new(browser: ResolvedBrowser) -> Self {
96 Self::with_sidecar_config(browser, crate::sidecar::SidecarConfig::default())
97 }
98
99 pub fn with_sidecar_config(
102 browser: ResolvedBrowser,
103 sidecar_config: crate::sidecar::SidecarConfig,
104 ) -> Self {
105 Self {
106 browser: Arc::new(RwLock::new(browser)),
107 bidi: Arc::new(Mutex::new(None)),
108 bidi_lock: Arc::new(Mutex::new(BidiLockState::Pending)),
109 backend: Arc::new(Mutex::new(None)),
110 active_target_id: Arc::new(Mutex::new(None)),
111 origin_target_ids: Arc::new(Mutex::new(HashMap::new())),
112 sidecar: Arc::new(Mutex::new(None)),
113 sidecar_config,
114 op_barrier: Arc::new(RwLock::new(())),
115 }
116 }
117
118 pub async fn ensure_sidecar(&self, tool_name: &str) -> Result<crate::sidecar::Sidecar> {
125 self.ensure_sidecar_supported(tool_name).await?;
126 let resolved = self.browser_snapshot().await;
127 let mut guard = self.sidecar.lock().await;
128 if let Some(sc) = guard.as_ref() {
129 return Ok(sc.clone());
130 }
131 let sc = crate::sidecar::Sidecar::start(self.sidecar_config.clone()).await?;
132 sc.connect(&resolved.endpoint).await?;
133 *guard = Some(sc.clone());
134 Ok(sc)
135 }
136
137 pub async fn ensure_sidecar_supported(&self, tool_name: &str) -> Result<()> {
140 self.ensure_active_browser_alive().await?;
141 let resolved = self.browser_snapshot().await;
142 if resolved.engine != crate::detect::Engine::Cdp {
143 return Err(crate::errors::SessionError::EngineUnsupported {
144 tool: tool_name.to_string(),
145 required_engine: "Chromium (CDP)".into(),
146 current_engine: format!("{:?}", resolved.engine),
147 hint: "use engine-agnostic tools such as browser_get_html, browser_fetch, browser_take_screenshot, or switch to a Chromium browser via browser_select",
148 }
149 .into());
150 }
151 Ok(())
152 }
153
154 pub async fn reset_sidecar(&self) {
158 let mut sidecar = self.sidecar.lock().await;
159 let _ = sidecar.take();
160 }
161
162 pub async fn browser_snapshot(&self) -> ResolvedBrowser {
164 self.browser.read().await.clone()
165 }
166
167 pub async fn ensure_bidi_lock(&self) -> Result<()> {
171 use crate::cli::env_resolver::Source;
172 use crate::cli::mcp::acquire_bidi_lock_if_needed;
173 use crate::detect::Engine;
174 use crate::registry::Registry;
175 let mut guard = self.bidi_lock.lock().await;
176 if matches!(*guard, BidiLockState::Pending) {
177 let resolved = self.browser_snapshot().await;
178 if resolved.engine != Engine::Bidi || matches!(resolved.source, Source::External) {
179 *guard = BidiLockState::NotApplicable;
180 return Ok(());
181 }
182 let acquired = tokio::task::spawn_blocking(move || {
188 let registry = Registry::open()?;
189 acquire_bidi_lock_if_needed(®istry, &resolved)
190 })
191 .await??;
192 *guard = match acquired {
193 Some(lock) => BidiLockState::Acquired(lock),
194 None => BidiLockState::NotApplicable,
195 };
196 }
197 Ok(())
198 }
199
200 pub async fn ensure_backend(&self) -> Result<TabBackend> {
205 self.ensure_active_browser_alive().await?;
206 self.ensure_bidi_lock().await?;
207 let mut guard = self.backend.lock().await;
208 if let Some(b) = guard.as_ref() {
209 return Ok(b.clone());
210 }
211 let resolved = self.browser_snapshot().await;
212 let b = crate::session::backend::open_backend(&resolved.endpoint, resolved.engine).await?;
213 *guard = Some(b.clone());
214 Ok(b)
215 }
216
217 pub async fn ensure_active_browser_alive(&self) -> Result<()> {
225 use crate::cli::env_resolver::Source;
226 use crate::registry::BrowserLiveness;
227
228 let resolved = self.browser_snapshot().await;
229 let Source::Registered { name } = resolved.source else {
230 return Ok(());
231 };
232 sync_registry_op(move |registry| {
233 let Some(row) = registry
234 .get_by_name(&name)
235 .with_context(|| format!("checking active browser `{name}`"))?
236 else {
237 let hint = "call `browser_start` to start a browser, or `browser_select` to switch to another live browser";
238 let kind_hint = generated_kind_hint(&name);
239 anyhow::bail!(
240 "active browser `{name}` is no longer registered{kind_hint}; {hint}"
241 );
242 };
243 match crate::registry::liveness(&row) {
244 BrowserLiveness::Alive => Ok(()),
245 BrowserLiveness::DeadPid => {
246 registry
247 .delete(&row.name)
248 .with_context(|| format!("pruning terminated browser {}", row.name))?;
249 anyhow::bail!(
250 "active browser `{}` has exited (pid {}); call `browser_start` with `{}` to launch/reuse a browser, or `browser_select` another live browser",
251 row.name,
252 row.pid,
253 row.kind.as_str()
254 );
255 }
256 BrowserLiveness::EndpointUnreachable => {
257 anyhow::bail!(
258 "active browser `{}` is not reachable at {} (pid {} still exists); retry, call `browser_start` with `{}` to launch/reuse a browser, or `browser_select` another live browser",
259 row.name,
260 row.endpoint,
261 row.pid,
262 row.kind.as_str()
263 );
264 }
265 }
266 })
267 .await
268 }
269
270 pub async fn current_tab(&self) -> Result<(TabBackend, String)> {
289 let backend = self.ensure_backend().await?;
290 let mut pointer = self.active_target_id.lock().await;
291 if let Some(tid) = pointer.as_ref() {
292 let live = backend.live_target_ids().await?;
293 if live.contains(tid) {
294 return Ok((backend, tid.clone()));
295 }
296 }
298 let new_tid = backend.create_tab("about:blank").await?;
299 *pointer = Some(new_tid.clone());
300 Ok((backend, new_tid))
301 }
302
303 pub async fn resolve_or_create_for_origin(&self, url: &str) -> Result<(TabBackend, String)> {
312 let want =
313 url::Url::parse(url).map_err(|e| anyhow::anyhow!("invalid fetch URL `{url}`: {e}"))?;
314 let origin_root = crate::session::attach::origin_root_url(&want);
315 let backend = self.ensure_backend().await?;
316 let mut origin_targets = self.origin_target_ids.lock().await;
317 let live_targets = backend.live_targets().await?;
318 let live_ids: std::collections::HashSet<&str> =
319 live_targets.iter().map(|t| t.id.as_str()).collect();
320
321 if let Some(cached) = origin_targets.get(&origin_root) {
322 if live_ids.contains(cached.as_str()) {
323 return Ok((backend, cached.clone()));
324 }
325 origin_targets.remove(&origin_root);
326 }
327
328 if let Some(existing) = live_targets.iter().find(|t| {
329 url::Url::parse(&t.url)
330 .map(|parsed| crate::session::attach::same_origin(&parsed, &want))
331 .unwrap_or(false)
332 }) {
333 origin_targets.insert(origin_root, existing.id.clone());
334 return Ok((backend, existing.id.clone()));
335 }
336
337 let new_tid = backend.create_tab(&origin_root).await?;
338 origin_targets.insert(origin_root, new_tid.clone());
339 Ok((backend, new_tid))
340 }
341
342 pub async fn resolve_target_for_args(&self, args: &Value) -> Result<(TabBackend, String)> {
357 let tab = args.get("tab").and_then(|v| v.as_str()).map(String::from);
358 let target = args
359 .get("target")
360 .and_then(|v| v.as_str())
361 .map(String::from);
362 match (tab, target) {
363 (Some(_), Some(_)) => Err(anyhow::anyhow!("`tab` and `target` are mutually exclusive")),
364 (Some(name), None) => {
365 let backend = self.ensure_backend().await?;
366 let browser_name = self.registered_browser_name().await?;
367 let bn = browser_name.clone();
371 let n = name.clone();
372 let row = sync_registry_op(move |reg| reg.tab_get(&bn, &n))
373 .await?
374 .ok_or_else(|| crate::errors::SessionError::TabNotFound {
375 browser: browser_name.clone(),
376 name: name.clone(),
377 })?;
378 let live = backend.live_target_ids().await?;
379 if !live.contains(&row.target_id) {
380 let bn = browser_name.clone();
382 let n = name.clone();
383 sync_registry_op(move |reg| reg.tab_delete(&bn, &n)).await?;
384 return Err(crate::errors::SessionError::TabNotFound {
385 browser: browser_name,
386 name,
387 }
388 .into());
389 }
390 let bn = browser_name.clone();
391 let n = name.clone();
392 sync_registry_op(move |reg| reg.tab_touch(&bn, &n)).await?;
393 Ok((backend, row.target_id))
394 }
395 (None, Some(regex)) => {
396 let backend = self.ensure_backend().await?;
397 let target_id = resolve_target_by_regex(&backend, ®ex).await?;
398 Ok((backend, target_id))
399 }
400 (None, None) => self.current_tab().await,
401 }
402 }
403
404 pub async fn registered_browser_name(&self) -> Result<String> {
407 use crate::cli::env_resolver::Source;
408 let resolved = self.browser_snapshot().await;
409 match resolved.source {
410 Source::Registered { name } => Ok(name),
411 Source::External => Err(anyhow::anyhow!(
412 "operation requires a registered browser; external URL endpoints \
413 don't have a stable identity"
414 )),
415 }
416 }
417
418 pub async fn switch_browser(&self, new_browser: ResolvedBrowser) -> Result<()> {
430 {
432 let mut bidi = self.bidi.lock().await;
433 if let Some(client) = bidi.take() {
434 let _ = client.session_end().await;
435 }
436 }
437 {
440 let mut backend = self.backend.lock().await;
441 *backend = None;
442 }
443 {
445 let mut lock = self.bidi_lock.lock().await;
446 *lock = BidiLockState::Pending;
447 }
448 {
450 let mut pointer = self.active_target_id.lock().await;
451 *pointer = None;
452 }
453 {
455 let mut origins = self.origin_target_ids.lock().await;
456 origins.clear();
457 }
458 {
462 let mut sidecar = self.sidecar.lock().await;
463 if let Some(sc) = sidecar.take() {
464 let _ = sc.call("dispose", serde_json::json!({})).await;
465 drop(sc);
467 }
468 }
469 {
471 let mut br = self.browser.write().await;
472 *br = new_browser;
473 }
474 self.ensure_bidi_lock().await?;
477 Ok(())
478 }
479}
480
481pub(crate) async fn sync_registry_op<T, F>(f: F) -> Result<T>
492where
493 F: FnOnce(&crate::registry::Registry) -> Result<T> + Send + 'static,
494 T: Send + 'static,
495{
496 tokio::task::spawn_blocking(move || {
497 let reg = crate::registry::Registry::open()?;
498 f(®)
499 })
500 .await?
501}
502
503fn generated_kind_hint(name: &str) -> String {
504 let Some((prefix, _)) = name.split_once('-') else {
505 return String::new();
506 };
507 if crate::detect::Kind::parse(prefix).is_some() {
508 format!("; it looks like a generated `{prefix}` browser name")
509 } else {
510 String::new()
511 }
512}
513
514pub(crate) async fn resolve_browser_send(
520 selector: crate::cli::env_resolver::BrowserSelector,
521) -> Result<ResolvedBrowser> {
522 use crate::cli::env_resolver::{BrowserSelector, DefaultResolver, Resolver};
523 let startable_kind = crate::cli::mcp::startable_kind_from_selector(&selector);
524 let resolved = match selector {
525 BrowserSelector::Url(u) => match u.scheme() {
526 "ws" | "wss" => Ok(ResolvedBrowser {
527 engine: if u.path().contains("/session") {
528 crate::detect::Engine::Bidi
529 } else {
530 crate::detect::Engine::Cdp
531 },
532 endpoint: u.to_string(),
533 source: crate::cli::env_resolver::Source::External,
534 }),
535 "http" | "https" => {
536 let base = u.as_str().trim_end_matches('/').to_string();
537 let ws = DefaultResolver.fetch_version(&base).await?;
538 let ws_url = url::Url::parse(&ws)?;
539 Ok(ResolvedBrowser {
540 engine: if ws_url.path().contains("/session") {
541 crate::detect::Engine::Bidi
542 } else {
543 crate::detect::Engine::Cdp
544 },
545 endpoint: ws,
546 source: crate::cli::env_resolver::Source::External,
547 })
548 }
549 other => anyhow::bail!("unsupported URL scheme: {other}"),
550 },
551 other => {
552 tokio::task::spawn_blocking(move || {
553 let reg = crate::registry::Registry::open()?;
554 let rt = tokio::runtime::Builder::new_current_thread().build()?;
559 rt.block_on(crate::cli::env_resolver::resolve_with(
560 other,
561 ®,
562 &DefaultResolver,
563 ))
564 })
565 .await?
566 }
567 };
568 match resolved {
569 Ok(browser) => Ok(browser),
570 Err(resolve_err) => {
571 if let Some(kind) = startable_kind {
572 crate::cli::mcp::start_and_resolve(Some(kind.as_str().to_string()), false, 30)
573 .await
574 .with_context(|| {
575 format!(
576 "browser selector failed to resolve ({resolve_err:#}); also failed to start {}",
577 kind.as_str()
578 )
579 })
580 } else {
581 Err(resolve_err)
582 }
583 }
584 }
585}
586
587async fn resolve_target_by_regex(backend: &TabBackend, regex: &str) -> Result<String> {
594 use crate::errors::SessionError;
595 use regex::Regex;
596 use std::time::Duration;
597 const PROBE: Duration = Duration::from_millis(500);
598
599 let re = Regex::new(regex)?;
600 let targets = backend.live_targets().await?;
601 let matches: Vec<_> = targets.iter().filter(|t| re.is_match(&t.url)).collect();
602 if matches.is_empty() {
603 return Err(anyhow::anyhow!("no target matched URL regex `{regex}`"));
604 }
605 let mut last_id: Option<String> = None;
606 let mut last_url: Option<String> = None;
607 for t in &matches {
608 last_id = Some(t.id.clone());
609 last_url = Some(t.url.clone());
610 let ok = matches!(
611 tokio::time::timeout(PROBE, backend.evaluate(&t.id, "1", false, PROBE)).await,
612 Ok(Ok(_))
613 );
614 if ok {
615 return Ok(t.id.clone());
616 }
617 }
618 Err(SessionError::TabHung {
619 target_id: last_id,
620 url: last_url,
621 timeout_ms: PROBE.as_millis() as u64,
622 hint: "all-matches-hung",
623 }
624 .into())
625}
626
627impl std::fmt::Debug for ServerState {
628 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
629 f.debug_struct("ServerState").finish()
630 }
631}
632
633pub type ToolHandler = std::sync::Arc<
635 dyn Fn(ServerState, Value) -> futures_util::future::BoxFuture<'static, Result<Value>>
636 + Send
637 + Sync,
638>;
639
640pub struct RegisteredTool {
641 pub name: String,
642 pub description: String,
643 pub input_schema: Value,
644 pub handler: ToolHandler,
645}
646
647#[derive(Clone, Default)]
648pub struct ToolRegistry {
649 inner: std::sync::Arc<std::sync::Mutex<Vec<RegisteredTool>>>,
650}
651
652impl ToolRegistry {
653 pub fn new() -> Self {
654 Self::default()
655 }
656
657 pub fn register(&self, t: RegisteredTool) {
658 self.inner.lock().unwrap().push(t);
659 }
660
661 pub fn list(&self) -> Vec<Value> {
662 self.inner
663 .lock()
664 .unwrap()
665 .iter()
666 .map(|t| {
667 json!({
668 "name": t.name,
669 "description": t.description,
670 "inputSchema": t.input_schema,
671 })
672 })
673 .collect()
674 }
675
676 pub fn handler(&self, name: &str) -> Option<ToolHandler> {
677 self.inner
678 .lock()
679 .unwrap()
680 .iter()
681 .find(|t| t.name == name)
682 .map(|t| t.handler.clone())
683 }
684}
685
686pub async fn run(state: ServerState, tools: ToolRegistry) -> Result<()> {
688 run_with_streams(state, tools, tokio::io::stdin(), tokio::io::stdout()).await
689}
690
691pub async fn run_with_streams<I, O>(
705 state: ServerState,
706 tools: ToolRegistry,
707 stdin: I,
708 mut stdout: O,
709) -> Result<()>
710where
711 I: tokio::io::AsyncRead + Unpin,
712 O: tokio::io::AsyncWrite + Unpin + Send + 'static,
713{
714 let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<Vec<u8>>();
718 let writer = tokio::spawn(async move {
719 while let Some(frame) = rx.recv().await {
720 if stdout.write_all(&frame).await.is_err() {
721 break;
722 }
723 if stdout.flush().await.is_err() {
724 break;
725 }
726 }
727 });
728
729 let mut lines = BufReader::new(stdin).lines();
730 while let Some(line) = lines.next_line().await? {
731 if line.trim().is_empty() {
732 continue;
733 }
734 let req: Value = match serde_json::from_str(&line) {
735 Ok(v) => v,
736 Err(e) => {
737 let _ = tx.send(error_frame(
738 Value::Null,
739 -32700,
740 &format!("parse error: {e}"),
741 ));
742 continue;
743 }
744 };
745 let id = req.get("id").cloned().unwrap_or(Value::Null);
746 let method = req.get("method").and_then(|m| m.as_str()).unwrap_or("");
747 let params = req.get("params").cloned().unwrap_or(Value::Null);
748
749 if id.is_null() && method == "notifications/initialized" {
751 continue;
752 }
753
754 match method {
759 "initialize" => {
760 let _ = tx.send(handle_initialize(id));
761 }
762 "ping" => {
763 let _ = tx.send(handle_ping(id));
764 }
765 "tools/list" => {
766 let _ = tx.send(handle_tools_list(id, &tools));
767 }
768 "tools/call" => {
769 handle_tools_call(id, ¶ms, &state, &tools, &tx);
770 }
771 _ => {
772 let _ = tx.send(error_frame(
773 id,
774 -32601,
775 &format!("method not found: {method}"),
776 ));
777 }
778 }
779 }
780 drop(tx);
783 let _ = writer.await;
784 Ok(())
785}
786
787fn handle_initialize(id: Value) -> Vec<u8> {
790 result_frame(
791 id,
792 json!({
793 "protocolVersion": "2024-11-05",
794 "capabilities": {"tools": {}},
795 "serverInfo": {
796 "name": "browser-control",
797 "version": env!("CARGO_PKG_VERSION"),
798 },
799 }),
800 )
801}
802
803fn handle_ping(id: Value) -> Vec<u8> {
805 result_frame(id, json!({}))
806}
807
808fn handle_tools_list(id: Value, tools: &ToolRegistry) -> Vec<u8> {
810 result_frame(id, json!({"tools": tools.list()}))
811}
812
813fn handle_tools_call(
824 id: Value,
825 params: &Value,
826 state: &ServerState,
827 tools: &ToolRegistry,
828 tx: &tokio::sync::mpsc::UnboundedSender<Vec<u8>>,
829) {
830 let name = params
831 .get("name")
832 .and_then(|v| v.as_str())
833 .unwrap_or("")
834 .to_string();
835 let args = params.get("arguments").cloned().unwrap_or(Value::Null);
836 let handler = tools.handler(&name);
837 let state = state.clone();
838 let tx = tx.clone();
839 let barrier = state.op_barrier.clone();
840 let is_exclusive = name == "browser_select";
841 tokio::spawn(async move {
842 let frame = match handler {
843 None => error_frame(id, -32602, &format!("tool not found: {name}")),
846 Some(h) => {
847 if is_exclusive {
848 let _guard = barrier.write().await;
853 match h(state, args).await {
854 Ok(v) => result_frame(id, v),
855 Err(e) => tool_error_frame(id, &e),
856 }
857 } else {
858 let _guard = barrier.read().await;
859 match h(state, args).await {
860 Ok(v) => result_frame(id, v),
861 Err(e) => tool_error_frame(id, &e),
862 }
863 }
864 }
865 };
866 let _ = tx.send(frame);
867 });
868}
869
870fn result_frame(id: Value, result: Value) -> Vec<u8> {
872 let resp = json!({"jsonrpc": "2.0", "id": id, "result": result});
873 let mut s = serde_json::to_vec(&resp).unwrap_or_else(|_| b"{}".to_vec());
874 s.push(b'\n');
875 s
876}
877
878fn tool_error_frame(id: Value, err: &anyhow::Error) -> Vec<u8> {
884 let message = tool_error_message(err);
885 result_frame(
886 id,
887 json!({
888 "content": [{ "type": "text", "text": message }],
889 "isError": true,
890 }),
891 )
892}
893
894fn tool_error_message(err: &anyhow::Error) -> String {
899 use crate::errors::SessionError;
900 use crate::registry::bidi_lock::BidiLockBusy;
901
902 if let Some(se) = err.downcast_ref::<SessionError>() {
903 return se.to_string();
906 }
907 if let Some(busy) = err.downcast_ref::<BidiLockBusy>() {
908 return busy.to_string();
909 }
910 format!("{err:#}")
913}
914
915fn error_frame(id: Value, code: i64, message: &str) -> Vec<u8> {
917 let resp = json!({
918 "jsonrpc": "2.0",
919 "id": id,
920 "error": {"code": code, "message": message},
921 });
922 let mut s = serde_json::to_vec(&resp).unwrap_or_else(|_| b"{}".to_vec());
923 s.push(b'\n');
924 s
925}
926
927#[cfg(test)]
928mod tests {
929 use super::*;
930 use crate::cli::env_resolver::Source;
931 use crate::detect::Engine;
932 use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
933
934 fn dummy_resolved() -> ResolvedBrowser {
935 ResolvedBrowser {
936 endpoint: "ws://localhost:9999".into(),
937 engine: Engine::Cdp,
938 source: Source::External,
939 }
940 }
941
942 fn dummy_state() -> ServerState {
943 ServerState::new(dummy_resolved())
944 }
945
946 async fn send_recv(tools: ToolRegistry, requests: &[Value]) -> Vec<Value> {
947 let (mut client_w, server_r) = tokio::io::duplex(8192);
948 let (server_w, client_r) = tokio::io::duplex(8192);
949 let state = dummy_state();
950 let join = tokio::spawn(async move {
951 let _ = run_with_streams(state, tools, server_r, server_w).await;
952 });
953
954 for req in requests {
955 let mut s = serde_json::to_vec(req).unwrap();
956 s.push(b'\n');
957 client_w.write_all(&s).await.unwrap();
958 }
959 drop(client_w);
961
962 let mut reader = BufReader::new(client_r);
963 let mut responses = Vec::new();
964 loop {
965 let mut line = String::new();
966 let n = reader.read_line(&mut line).await.unwrap();
967 if n == 0 {
968 break;
969 }
970 responses.push(serde_json::from_str(line.trim()).unwrap());
971 }
972 let _ = join.await;
973 responses
974 }
975
976 fn echo_tool() -> RegisteredTool {
977 RegisteredTool {
978 name: "echo".to_string(),
979 description: "Echo arguments back".to_string(),
980 input_schema: json!({"type": "object"}),
981 handler: std::sync::Arc::new(|_state, args| {
982 Box::pin(async move { Ok(json!({"echoed": args})) })
983 }),
984 }
985 }
986
987 fn failing_tool() -> RegisteredTool {
988 RegisteredTool {
989 name: "boom".to_string(),
990 description: "Always fails with a typed SessionError".to_string(),
991 input_schema: json!({"type": "object"}),
992 handler: std::sync::Arc::new(|_state, _args| {
993 Box::pin(async move {
994 Err(anyhow::Error::new(crate::errors::SessionError::TabHung {
995 target_id: Some("T1".into()),
996 url: Some("https://example.test".into()),
997 timeout_ms: 20_000,
998 hint: "renderer wedged",
999 }))
1000 })
1001 }),
1002 }
1003 }
1004
1005 #[tokio::test]
1006 async fn initialize_round_trip() {
1007 let resp = send_recv(
1008 ToolRegistry::new(),
1009 &[json!({"jsonrpc":"2.0","id":1,"method":"initialize","params":{}})],
1010 )
1011 .await;
1012 assert_eq!(resp.len(), 1);
1013 assert_eq!(resp[0]["id"], 1);
1014 assert_eq!(resp[0]["result"]["protocolVersion"], "2024-11-05");
1015 assert_eq!(resp[0]["result"]["serverInfo"]["name"], "browser-control");
1016 }
1017
1018 #[tokio::test]
1019 async fn tools_list_empty() {
1020 let resp = send_recv(
1021 ToolRegistry::new(),
1022 &[json!({"jsonrpc":"2.0","id":2,"method":"tools/list"})],
1023 )
1024 .await;
1025 assert_eq!(resp[0]["result"]["tools"], json!([]));
1026 }
1027
1028 #[tokio::test]
1029 async fn tools_list_after_register() {
1030 let tools = ToolRegistry::new();
1031 tools.register(echo_tool());
1032 let resp = send_recv(
1033 tools,
1034 &[json!({"jsonrpc":"2.0","id":3,"method":"tools/list"})],
1035 )
1036 .await;
1037 let list = resp[0]["result"]["tools"].as_array().unwrap();
1038 assert_eq!(list.len(), 1);
1039 assert_eq!(list[0]["name"], "echo");
1040 }
1041
1042 #[tokio::test]
1043 async fn tools_call_unknown_errors() {
1044 let resp = send_recv(
1045 ToolRegistry::new(),
1046 &[json!({
1047 "jsonrpc":"2.0","id":4,"method":"tools/call",
1048 "params":{"name":"nope","arguments":{}}
1049 })],
1050 )
1051 .await;
1052 assert_eq!(resp[0]["error"]["code"], -32602);
1054 assert!(resp[0]["error"]["message"]
1055 .as_str()
1056 .unwrap()
1057 .contains("nope"));
1058 }
1059
1060 #[tokio::test]
1061 async fn tool_failure_returns_iserror_result_not_protocol_error() {
1062 let tools = ToolRegistry::new();
1063 tools.register(failing_tool());
1064 let resp = send_recv(
1065 tools,
1066 &[json!({
1067 "jsonrpc":"2.0","id":7,"method":"tools/call",
1068 "params":{"name":"boom","arguments":{}}
1069 })],
1070 )
1071 .await;
1072 assert!(resp[0]["error"].is_null());
1075 assert_eq!(resp[0]["result"]["isError"], true);
1076 let text = resp[0]["result"]["content"][0]["text"].as_str().unwrap();
1077 assert!(text.contains("tab hung"), "got: {text}");
1079 assert!(text.contains("renderer wedged"), "got: {text}");
1080 assert!(text.contains("T1"), "got: {text}");
1081 }
1082
1083 #[tokio::test]
1084 async fn tools_call_registered_returns_result() {
1085 let tools = ToolRegistry::new();
1086 tools.register(echo_tool());
1087 let resp = send_recv(
1088 tools,
1089 &[json!({
1090 "jsonrpc":"2.0","id":5,"method":"tools/call",
1091 "params":{"name":"echo","arguments":{"hello":"world"}}
1092 })],
1093 )
1094 .await;
1095 assert_eq!(resp[0]["result"]["echoed"], json!({"hello":"world"}));
1096 }
1097
1098 #[tokio::test]
1099 async fn unknown_method_returns_minus_32601() {
1100 let resp = send_recv(
1101 ToolRegistry::new(),
1102 &[json!({"jsonrpc":"2.0","id":6,"method":"bogus"})],
1103 )
1104 .await;
1105 assert_eq!(resp[0]["error"]["code"], -32601);
1106 }
1107
1108 #[tokio::test]
1109 async fn ping_returns_empty_object() {
1110 let resp = send_recv(
1111 ToolRegistry::new(),
1112 &[json!({"jsonrpc":"2.0","id":7,"method":"ping"})],
1113 )
1114 .await;
1115 assert_eq!(resp[0]["result"], json!({}));
1116 }
1117
1118 use futures_util::{SinkExt, StreamExt};
1127 use std::sync::atomic::{AtomicUsize, Ordering};
1128 use tokio_tungstenite::tungstenite::Message;
1129
1130 async fn spawn_counting_cdp_mock(
1134 live: Vec<String>,
1135 ) -> (String, Arc<AtomicUsize>, tokio::sync::oneshot::Sender<()>) {
1136 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
1137 let addr = listener.local_addr().unwrap();
1138 let (stop_tx, mut stop_rx) = tokio::sync::oneshot::channel::<()>();
1139 let conns = Arc::new(AtomicUsize::new(0));
1140 let conns_srv = conns.clone();
1141 tokio::spawn(async move {
1142 loop {
1143 let accept = tokio::select! {
1144 _ = &mut stop_rx => break,
1145 a = listener.accept() => a,
1146 };
1147 let (stream, _) = match accept {
1148 Ok(s) => s,
1149 Err(_) => break,
1150 };
1151 conns_srv.fetch_add(1, Ordering::SeqCst);
1152 let live = live.clone();
1153 tokio::spawn(async move {
1154 let mut ws = match tokio_tungstenite::accept_async(stream).await {
1155 Ok(w) => w,
1156 Err(_) => return,
1157 };
1158 while let Some(Ok(msg)) = ws.next().await {
1159 if let Message::Text(t) = msg {
1160 let req: Value = match serde_json::from_str(&t) {
1161 Ok(v) => v,
1162 Err(_) => continue,
1163 };
1164 let id = req["id"].as_u64().unwrap_or(0);
1165 let method = req["method"].as_str().unwrap_or("");
1166 let result = match method {
1167 "Target.getTargets" => {
1168 let infos: Vec<Value> = live
1169 .iter()
1170 .map(|tid| {
1171 json!({"targetId": tid, "type": "page", "url": ""})
1172 })
1173 .collect();
1174 json!({"targetInfos": infos})
1175 }
1176 "Target.attachToTarget" => json!({"sessionId": "S1"}),
1177 "Runtime.evaluate" => json!({"result": {"value": 1}}),
1178 _ => json!({}),
1179 };
1180 let resp = json!({"id": id, "result": result});
1181 if ws.send(Message::Text(resp.to_string())).await.is_err() {
1182 break;
1183 }
1184 }
1185 }
1186 });
1187 }
1188 });
1189 (format!("ws://{addr}"), conns, stop_tx)
1190 }
1191
1192 fn registered_state(name: &str, endpoint: &str) -> ServerState {
1195 ServerState::new(ResolvedBrowser {
1196 endpoint: endpoint.to_string(),
1197 engine: Engine::Cdp,
1198 source: Source::Registered { name: name.into() },
1199 })
1200 }
1201
1202 fn register_browser_row(name: &str, endpoint: &str) {
1203 let port: u16 = endpoint
1204 .strip_prefix("ws://127.0.0.1:")
1205 .and_then(|s| s.split('/').next())
1206 .and_then(|s| s.parse().ok())
1207 .expect("parse mock port");
1208 let reg = crate::registry::Registry::open().unwrap();
1209 reg.insert(&crate::registry::BrowserRow {
1210 name: name.to_string(),
1211 kind: crate::detect::Kind::Chrome,
1212 engine: Engine::Cdp,
1213 pid: std::process::id(),
1214 endpoint: endpoint.to_string(),
1215 port,
1216 profile_dir: std::path::PathBuf::from("/tmp/profiles/bx"),
1217 executable: std::path::PathBuf::from("/usr/bin/example"),
1218 headless: false,
1219 started_at: "2024-01-01T00:00:00Z".into(),
1220 })
1221 .unwrap();
1222 }
1223
1224 #[allow(clippy::await_holding_lock)]
1227 #[tokio::test]
1228 async fn resolve_named_tab_live_resolves_and_touches() {
1229 let _g = crate::test_support::ENV_LOCK
1230 .lock()
1231 .unwrap_or_else(|e| e.into_inner());
1232 let tmp = tempfile::TempDir::new().unwrap();
1233 std::env::set_var("BROWSER_CONTROL_DATA_DIR", tmp.path());
1234
1235 {
1237 let reg = crate::registry::Registry::open().unwrap();
1238 reg.tab_upsert("bx", "work", "T1", "about:blank", true)
1239 .unwrap();
1240 }
1241
1242 let (url, _conns, _stop) = spawn_counting_cdp_mock(vec!["T1".into()]).await;
1243 register_browser_row("bx", &url);
1244 let state = registered_state("bx", &url);
1245
1246 let target_id = match state.resolve_target_for_args(&json!({"tab": "work"})).await {
1247 Ok((_backend, tid)) => tid,
1248 Err(e) => panic!("named tab should resolve: {e:#}"),
1249 };
1250 assert_eq!(target_id, "T1");
1251
1252 std::env::remove_var("BROWSER_CONTROL_DATA_DIR");
1253 }
1254
1255 #[allow(clippy::await_holding_lock)]
1257 #[tokio::test]
1258 async fn resolve_named_tab_stale_sweeps_and_errors() {
1259 let _g = crate::test_support::ENV_LOCK
1260 .lock()
1261 .unwrap_or_else(|e| e.into_inner());
1262 let tmp = tempfile::TempDir::new().unwrap();
1263 std::env::set_var("BROWSER_CONTROL_DATA_DIR", tmp.path());
1264
1265 {
1269 let reg = crate::registry::Registry::open().unwrap();
1270 reg.tab_upsert("bx", "gone", "T_DEAD", "about:blank", true)
1271 .unwrap();
1272 }
1273
1274 let (url, _conns, _stop) = spawn_counting_cdp_mock(vec!["T1".into()]).await;
1275 register_browser_row("bx", &url);
1276 let state = registered_state("bx", &url);
1277
1278 let err = match state.resolve_target_for_args(&json!({"tab": "gone"})).await {
1279 Ok(_) => panic!("stale named tab must error"),
1280 Err(e) => e,
1281 };
1282 let typed = err
1283 .downcast_ref::<crate::errors::SessionError>()
1284 .expect("typed SessionError");
1285 assert!(
1286 matches!(typed, crate::errors::SessionError::TabNotFound { .. }),
1287 "expected TabNotFound, got {typed:?}"
1288 );
1289
1290 let reg = crate::registry::Registry::open().unwrap();
1292 assert!(
1293 reg.tab_get("bx", "gone").unwrap().is_none(),
1294 "row not swept"
1295 );
1296
1297 std::env::remove_var("BROWSER_CONTROL_DATA_DIR");
1298 }
1299
1300 #[allow(clippy::await_holding_lock)]
1302 #[tokio::test]
1303 async fn resolve_named_tab_missing_row_errors() {
1304 let _g = crate::test_support::ENV_LOCK
1305 .lock()
1306 .unwrap_or_else(|e| e.into_inner());
1307 let tmp = tempfile::TempDir::new().unwrap();
1308 std::env::set_var("BROWSER_CONTROL_DATA_DIR", tmp.path());
1309
1310 let (url, _conns, _stop) = spawn_counting_cdp_mock(vec!["T1".into()]).await;
1311 register_browser_row("bx", &url);
1312 let state = registered_state("bx", &url);
1313
1314 let err = match state.resolve_target_for_args(&json!({"tab": "nope"})).await {
1315 Ok(_) => panic!("unknown named tab must error"),
1316 Err(e) => e,
1317 };
1318 let typed = err
1319 .downcast_ref::<crate::errors::SessionError>()
1320 .expect("typed SessionError");
1321 assert!(
1322 matches!(typed, crate::errors::SessionError::TabNotFound { .. }),
1323 "expected TabNotFound, got {typed:?}"
1324 );
1325
1326 std::env::remove_var("BROWSER_CONTROL_DATA_DIR");
1327 }
1328
1329 #[tokio::test]
1330 async fn resolve_tab_and_target_mutually_exclusive() {
1331 let state = dummy_state();
1333 let err = match state
1334 .resolve_target_for_args(&json!({"tab": "a", "target": "b"}))
1335 .await
1336 {
1337 Ok(_) => panic!("tab+target must error"),
1338 Err(e) => e,
1339 };
1340 assert!(
1341 err.to_string().contains("mutually exclusive"),
1342 "got: {err:#}"
1343 );
1344 }
1345
1346 #[tokio::test]
1347 async fn concurrent_ensure_backend_opens_one_backend() {
1348 let (url, conns, _stop) = spawn_counting_cdp_mock(vec!["T1".into()]).await;
1352 let state = ServerState::new(ResolvedBrowser {
1353 endpoint: url,
1354 engine: Engine::Cdp,
1355 source: Source::External,
1356 });
1357
1358 let mut handles = Vec::new();
1359 for _ in 0..8 {
1360 let s = state.clone();
1361 handles.push(tokio::spawn(async move { s.ensure_backend().await }));
1362 }
1363 for h in handles {
1364 h.await.unwrap().expect("ensure_backend should succeed");
1365 }
1366 assert_eq!(
1367 conns.load(Ordering::SeqCst),
1368 1,
1369 "expected exactly one backend (one WS connection) under concurrency"
1370 );
1371 }
1372
1373 #[allow(clippy::await_holding_lock)]
1374 #[tokio::test]
1375 async fn external_cdp_backend_does_not_open_registry() {
1376 let _g = crate::test_support::ENV_LOCK
1377 .lock()
1378 .unwrap_or_else(|e| e.into_inner());
1379 let stale_path = {
1380 let tmp = tempfile::TempDir::new().unwrap();
1381 tmp.path().to_path_buf()
1382 };
1383 std::env::set_var("BROWSER_CONTROL_DATA_DIR", &stale_path);
1384
1385 let (url, _conns, _stop) = spawn_counting_cdp_mock(vec!["T1".into()]).await;
1386 let state = ServerState::new(ResolvedBrowser {
1387 endpoint: url,
1388 engine: Engine::Cdp,
1389 source: Source::External,
1390 });
1391
1392 let result = state.ensure_backend().await;
1393
1394 std::env::remove_var("BROWSER_CONTROL_DATA_DIR");
1395 result.expect("external CDP backend must not depend on registry env");
1396 }
1397
1398 #[tokio::test]
1399 async fn initialized_notification_is_silently_ignored() {
1400 let resp = send_recv(
1403 ToolRegistry::new(),
1404 &[
1405 json!({"jsonrpc":"2.0","method":"notifications/initialized"}),
1406 json!({"jsonrpc":"2.0","id":8,"method":"ping"}),
1407 ],
1408 )
1409 .await;
1410 assert_eq!(resp.len(), 1);
1411 assert_eq!(resp[0]["id"], 8);
1412 }
1413}