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 refs: Arc<Mutex<HashMap<String, crate::a11y::RefTable>>>,
77 pub capture: Arc<crate::session::capture::CaptureHub>,
80 pub op_barrier: Arc<RwLock<()>>,
90}
91
92#[derive(Default)]
96pub enum BidiLockState {
97 #[default]
98 Pending,
99 Acquired(crate::registry::BidiLockGuard),
100 NotApplicable,
101}
102
103impl ServerState {
104 pub fn new(browser: ResolvedBrowser) -> Self {
105 Self::with_sidecar_config(browser, crate::sidecar::SidecarConfig::default())
106 }
107
108 pub fn with_sidecar_config(
111 browser: ResolvedBrowser,
112 sidecar_config: crate::sidecar::SidecarConfig,
113 ) -> Self {
114 Self {
115 browser: Arc::new(RwLock::new(browser)),
116 bidi: Arc::new(Mutex::new(None)),
117 bidi_lock: Arc::new(Mutex::new(BidiLockState::Pending)),
118 backend: Arc::new(Mutex::new(None)),
119 active_target_id: Arc::new(Mutex::new(None)),
120 origin_target_ids: Arc::new(Mutex::new(HashMap::new())),
121 sidecar: Arc::new(Mutex::new(None)),
122 sidecar_config,
123 refs: Arc::new(Mutex::new(HashMap::new())),
124 capture: Arc::new(crate::session::capture::CaptureHub::new()),
125 op_barrier: Arc::new(RwLock::new(())),
126 }
127 }
128
129 pub async fn ensure_body_capture_supported(&self, tool_name: &str) -> Result<()> {
133 self.ensure_cdp_engine(tool_name, crate::session::capture::BIDI_NO_BODIES_HINT)
134 .await
135 }
136
137 pub async fn ensure_cdp_engine(&self, tool_name: &str, hint: &'static 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,
148 }
149 .into());
150 }
151 Ok(())
152 }
153
154 pub async fn ensure_native_ready(&self, _tool_name: &str) -> Result<()> {
160 self.ensure_active_browser_alive().await
161 }
162
163 pub async fn ensure_sidecar(&self, tool_name: &str) -> Result<crate::sidecar::Sidecar> {
170 self.ensure_sidecar_supported(tool_name).await?;
171 let resolved = self.browser_snapshot().await;
172 let mut guard = self.sidecar.lock().await;
173 if let Some(sc) = guard.as_ref() {
174 return Ok(sc.clone());
175 }
176 let sc = crate::sidecar::Sidecar::start(self.sidecar_config.clone()).await?;
177 sc.connect(&resolved.endpoint).await?;
178 *guard = Some(sc.clone());
179 Ok(sc)
180 }
181
182 pub async fn ensure_sidecar_supported(&self, tool_name: &str) -> Result<()> {
185 self.ensure_cdp_engine(
186 tool_name,
187 "use engine-agnostic tools such as browser_get_html, browser_fetch, browser_take_screenshot, or switch to a Chromium browser via browser_select",
188 )
189 .await
190 }
191
192 pub async fn reset_sidecar(&self) {
196 let mut sidecar = self.sidecar.lock().await;
197 let _ = sidecar.take();
198 }
199
200 pub async fn browser_snapshot(&self) -> ResolvedBrowser {
202 self.browser.read().await.clone()
203 }
204
205 pub async fn ensure_bidi_lock(&self) -> Result<()> {
209 use crate::cli::env_resolver::Source;
210 use crate::cli::mcp::acquire_bidi_lock_if_needed;
211 use crate::detect::Engine;
212 use crate::registry::Registry;
213 let mut guard = self.bidi_lock.lock().await;
214 if matches!(*guard, BidiLockState::Pending) {
215 let resolved = self.browser_snapshot().await;
216 if resolved.engine != Engine::Bidi || matches!(resolved.source, Source::External) {
217 *guard = BidiLockState::NotApplicable;
218 return Ok(());
219 }
220 let acquired = tokio::task::spawn_blocking(move || {
226 let registry = Registry::open()?;
227 acquire_bidi_lock_if_needed(®istry, &resolved)
228 })
229 .await??;
230 *guard = match acquired {
231 Some(lock) => BidiLockState::Acquired(lock),
232 None => BidiLockState::NotApplicable,
233 };
234 }
235 Ok(())
236 }
237
238 pub async fn ensure_backend(&self) -> Result<TabBackend> {
243 self.ensure_active_browser_alive().await?;
244 self.ensure_bidi_lock().await?;
245 let mut guard = self.backend.lock().await;
246 if let Some(b) = guard.as_ref() {
247 return Ok(b.clone());
248 }
249 let resolved = self.browser_snapshot().await;
250 let b = crate::session::backend::open_backend(&resolved.endpoint, resolved.engine).await?;
251 *guard = Some(b.clone());
252 Ok(b)
253 }
254
255 pub async fn ensure_active_browser_alive(&self) -> Result<()> {
263 use crate::cli::env_resolver::Source;
264 use crate::registry::BrowserLiveness;
265
266 let resolved = self.browser_snapshot().await;
267 let Source::Registered { name } = resolved.source else {
268 return Ok(());
269 };
270 sync_registry_op(move |registry| {
271 let Some(row) = registry
272 .get_by_name(&name)
273 .with_context(|| format!("checking active browser `{name}`"))?
274 else {
275 let hint = "call `browser_start` to start a browser, or `browser_select` to switch to another live browser";
276 let kind_hint = generated_kind_hint(&name);
277 anyhow::bail!(
278 "active browser `{name}` is no longer registered{kind_hint}; {hint}"
279 );
280 };
281 match crate::registry::liveness(&row) {
282 BrowserLiveness::Alive => Ok(()),
283 BrowserLiveness::DeadPid => {
284 registry
285 .delete(&row.name)
286 .with_context(|| format!("pruning terminated browser {}", row.name))?;
287 anyhow::bail!(
288 "active browser `{}` has exited (pid {}); call `browser_start` with `{}` to launch/reuse a browser, or `browser_select` another live browser",
289 row.name,
290 row.pid,
291 row.kind.as_str()
292 );
293 }
294 BrowserLiveness::EndpointUnreachable => {
295 anyhow::bail!(
296 "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",
297 row.name,
298 row.endpoint,
299 row.pid,
300 row.kind.as_str()
301 );
302 }
303 }
304 })
305 .await
306 }
307
308 pub async fn current_tab(&self) -> Result<(TabBackend, String)> {
327 let backend = self.ensure_backend().await?;
328 let mut pointer = self.active_target_id.lock().await;
329 if let Some(tid) = pointer.as_ref() {
330 let live = backend.live_target_ids().await?;
331 if live.contains(tid) {
332 self.capture.touch(&backend, tid);
333 return Ok((backend, tid.clone()));
334 }
335 }
337 let new_tid = backend.create_tab("about:blank").await?;
338 *pointer = Some(new_tid.clone());
339 self.capture.touch(&backend, &new_tid);
340 Ok((backend, new_tid))
341 }
342
343 pub async fn resolve_or_create_for_origin(&self, url: &str) -> Result<(TabBackend, String)> {
352 let want =
353 url::Url::parse(url).map_err(|e| anyhow::anyhow!("invalid fetch URL `{url}`: {e}"))?;
354 let origin_root = crate::session::attach::origin_root_url(&want);
355 let backend = self.ensure_backend().await?;
356 let mut origin_targets = self.origin_target_ids.lock().await;
357 let live_targets = backend.live_targets().await?;
358 let live_ids: std::collections::HashSet<&str> =
359 live_targets.iter().map(|t| t.id.as_str()).collect();
360
361 if let Some(cached) = origin_targets.get(&origin_root) {
362 if live_ids.contains(cached.as_str()) {
363 self.capture.touch(&backend, cached);
364 return Ok((backend, cached.clone()));
365 }
366 origin_targets.remove(&origin_root);
367 }
368
369 if let Some(existing) = live_targets.iter().find(|t| {
370 url::Url::parse(&t.url)
371 .map(|parsed| crate::session::attach::same_origin(&parsed, &want))
372 .unwrap_or(false)
373 }) {
374 origin_targets.insert(origin_root, existing.id.clone());
375 self.capture.touch(&backend, &existing.id);
376 return Ok((backend, existing.id.clone()));
377 }
378
379 let new_tid = backend.create_tab(&origin_root).await?;
380 origin_targets.insert(origin_root, new_tid.clone());
381 self.capture.touch(&backend, &new_tid);
382 Ok((backend, new_tid))
383 }
384
385 pub async fn resolve_target_for_args(&self, args: &Value) -> Result<(TabBackend, String)> {
400 let tab = args.get("tab").and_then(|v| v.as_str()).map(String::from);
401 let target = args
402 .get("target")
403 .and_then(|v| v.as_str())
404 .map(String::from);
405 match (tab, target) {
406 (Some(_), Some(_)) => Err(anyhow::anyhow!("`tab` and `target` are mutually exclusive")),
407 (Some(name), None) => {
408 let backend = self.ensure_backend().await?;
409 let browser_name = self.registered_browser_name().await?;
410 let bn = browser_name.clone();
414 let n = name.clone();
415 let row = sync_registry_op(move |reg| reg.tab_get(&bn, &n))
416 .await?
417 .ok_or_else(|| crate::errors::SessionError::TabNotFound {
418 browser: browser_name.clone(),
419 name: name.clone(),
420 })?;
421 let live = backend.live_target_ids().await?;
422 if !live.contains(&row.target_id) {
423 let bn = browser_name.clone();
425 let n = name.clone();
426 sync_registry_op(move |reg| reg.tab_delete(&bn, &n)).await?;
427 return Err(crate::errors::SessionError::TabNotFound {
428 browser: browser_name,
429 name,
430 }
431 .into());
432 }
433 let bn = browser_name.clone();
434 let n = name.clone();
435 sync_registry_op(move |reg| reg.tab_touch(&bn, &n)).await?;
436 self.capture.touch(&backend, &row.target_id);
437 Ok((backend, row.target_id))
438 }
439 (None, Some(regex)) => {
440 let backend = self.ensure_backend().await?;
441 let target_id = resolve_target_by_regex(&backend, ®ex).await?;
442 self.capture.touch(&backend, &target_id);
443 Ok((backend, target_id))
444 }
445 (None, None) => self.current_tab().await,
446 }
447 }
448
449 pub async fn registered_browser_name(&self) -> Result<String> {
452 use crate::cli::env_resolver::Source;
453 let resolved = self.browser_snapshot().await;
454 match resolved.source {
455 Source::Registered { name } => Ok(name),
456 Source::External => Err(anyhow::anyhow!(
457 "operation requires a registered browser; external URL endpoints \
458 don't have a stable identity"
459 )),
460 }
461 }
462
463 pub async fn switch_browser(&self, new_browser: ResolvedBrowser) -> Result<()> {
475 self.capture.reset();
478 {
480 let mut bidi = self.bidi.lock().await;
481 if let Some(client) = bidi.take() {
482 let _ = client.session_end().await;
483 }
484 }
485 {
488 let old = self.backend.lock().await.take();
489 if let Some(b) = old {
490 b.shutdown().await;
491 }
492 }
493 {
495 let mut lock = self.bidi_lock.lock().await;
496 *lock = BidiLockState::Pending;
497 }
498 {
500 let mut pointer = self.active_target_id.lock().await;
501 *pointer = None;
502 }
503 {
505 let mut origins = self.origin_target_ids.lock().await;
506 origins.clear();
507 }
508 {
510 let mut refs = self.refs.lock().await;
511 refs.clear();
512 }
513 {
517 let mut sidecar = self.sidecar.lock().await;
518 if let Some(sc) = sidecar.take() {
519 let _ = sc.call("dispose", serde_json::json!({})).await;
520 drop(sc);
522 }
523 }
524 {
526 let mut br = self.browser.write().await;
527 *br = new_browser;
528 }
529 self.ensure_bidi_lock().await?;
532 Ok(())
533 }
534}
535
536pub(crate) async fn sync_registry_op<T, F>(f: F) -> Result<T>
547where
548 F: FnOnce(&crate::registry::Registry) -> Result<T> + Send + 'static,
549 T: Send + 'static,
550{
551 tokio::task::spawn_blocking(move || {
552 let reg = crate::registry::Registry::open()?;
553 f(®)
554 })
555 .await?
556}
557
558fn generated_kind_hint(name: &str) -> String {
559 let Some((prefix, _)) = name.split_once('-') else {
560 return String::new();
561 };
562 if crate::detect::Kind::parse(prefix).is_some() {
563 format!("; it looks like a generated `{prefix}` browser name")
564 } else {
565 String::new()
566 }
567}
568
569pub(crate) async fn resolve_browser_send(
575 selector: crate::cli::env_resolver::BrowserSelector,
576) -> Result<ResolvedBrowser> {
577 use crate::cli::env_resolver::{BrowserSelector, DefaultResolver, Resolver};
578 let startable_kind = crate::cli::mcp::startable_kind_from_selector(&selector);
579 let resolved = match selector {
580 BrowserSelector::Url(u) => match u.scheme() {
581 "ws" | "wss" => Ok(ResolvedBrowser {
582 engine: if u.path().contains("/session") {
583 crate::detect::Engine::Bidi
584 } else {
585 crate::detect::Engine::Cdp
586 },
587 endpoint: u.to_string(),
588 source: crate::cli::env_resolver::Source::External,
589 }),
590 "http" | "https" => {
591 let base = u.as_str().trim_end_matches('/').to_string();
592 let ws = DefaultResolver.fetch_version(&base).await?;
593 let ws_url = url::Url::parse(&ws)?;
594 Ok(ResolvedBrowser {
595 engine: if ws_url.path().contains("/session") {
596 crate::detect::Engine::Bidi
597 } else {
598 crate::detect::Engine::Cdp
599 },
600 endpoint: ws,
601 source: crate::cli::env_resolver::Source::External,
602 })
603 }
604 other => anyhow::bail!("unsupported URL scheme: {other}"),
605 },
606 other => {
607 tokio::task::spawn_blocking(move || {
608 let reg = crate::registry::Registry::open()?;
609 let rt = tokio::runtime::Builder::new_current_thread().build()?;
614 rt.block_on(crate::cli::env_resolver::resolve_with(
615 other,
616 ®,
617 &DefaultResolver,
618 ))
619 })
620 .await?
621 }
622 };
623 match resolved {
624 Ok(browser) => Ok(browser),
625 Err(resolve_err) => {
626 if let Some(kind) = startable_kind {
627 crate::cli::mcp::start_and_resolve(Some(kind.as_str().to_string()), false, 30)
628 .await
629 .with_context(|| {
630 format!(
631 "browser selector failed to resolve ({resolve_err:#}); also failed to start {}",
632 kind.as_str()
633 )
634 })
635 } else {
636 Err(resolve_err)
637 }
638 }
639 }
640}
641
642async fn resolve_target_by_regex(backend: &TabBackend, regex: &str) -> Result<String> {
649 use crate::errors::SessionError;
650 use regex::Regex;
651 use std::time::Duration;
652 const PROBE: Duration = Duration::from_millis(500);
653
654 let re = Regex::new(regex)?;
655 let targets = backend.live_targets().await?;
656 let matches: Vec<_> = targets.iter().filter(|t| re.is_match(&t.url)).collect();
657 if matches.is_empty() {
658 return Err(anyhow::anyhow!("no target matched URL regex `{regex}`"));
659 }
660 let mut last_id: Option<String> = None;
661 let mut last_url: Option<String> = None;
662 for t in &matches {
663 last_id = Some(t.id.clone());
664 last_url = Some(t.url.clone());
665 let ok = matches!(
666 tokio::time::timeout(PROBE, backend.evaluate(&t.id, "1", false, PROBE)).await,
667 Ok(Ok(_))
668 );
669 if ok {
670 return Ok(t.id.clone());
671 }
672 }
673 Err(SessionError::TabHung {
674 target_id: last_id,
675 url: last_url,
676 timeout_ms: PROBE.as_millis() as u64,
677 hint: "all-matches-hung",
678 }
679 .into())
680}
681
682impl std::fmt::Debug for ServerState {
683 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
684 f.debug_struct("ServerState").finish()
685 }
686}
687
688pub type ToolHandler = std::sync::Arc<
690 dyn Fn(ServerState, Value) -> futures_util::future::BoxFuture<'static, Result<Value>>
691 + Send
692 + Sync,
693>;
694
695pub struct RegisteredTool {
696 pub name: String,
697 pub description: String,
698 pub input_schema: Value,
699 pub handler: ToolHandler,
700}
701
702#[derive(Clone, Default)]
703pub struct ToolRegistry {
704 inner: std::sync::Arc<std::sync::Mutex<Vec<RegisteredTool>>>,
705}
706
707impl ToolRegistry {
708 pub fn new() -> Self {
709 Self::default()
710 }
711
712 pub fn register(&self, t: RegisteredTool) {
713 self.inner.lock().unwrap().push(t);
714 }
715
716 pub fn list(&self) -> Vec<Value> {
717 self.inner
718 .lock()
719 .unwrap()
720 .iter()
721 .map(|t| {
722 json!({
723 "name": t.name,
724 "description": t.description,
725 "inputSchema": t.input_schema,
726 })
727 })
728 .collect()
729 }
730
731 pub fn handler(&self, name: &str) -> Option<ToolHandler> {
732 self.inner
733 .lock()
734 .unwrap()
735 .iter()
736 .find(|t| t.name == name)
737 .map(|t| t.handler.clone())
738 }
739}
740
741pub async fn run(state: ServerState, tools: ToolRegistry) -> Result<()> {
743 run_with_streams(state, tools, tokio::io::stdin(), tokio::io::stdout()).await
744}
745
746pub async fn run_with_streams<I, O>(
760 state: ServerState,
761 tools: ToolRegistry,
762 stdin: I,
763 mut stdout: O,
764) -> Result<()>
765where
766 I: tokio::io::AsyncRead + Unpin,
767 O: tokio::io::AsyncWrite + Unpin + Send + 'static,
768{
769 let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<Vec<u8>>();
773 let writer = tokio::spawn(async move {
774 while let Some(frame) = rx.recv().await {
775 if stdout.write_all(&frame).await.is_err() {
776 break;
777 }
778 if stdout.flush().await.is_err() {
779 break;
780 }
781 }
782 });
783
784 let mut lines = BufReader::new(stdin).lines();
785 while let Some(line) = lines.next_line().await? {
786 if line.trim().is_empty() {
787 continue;
788 }
789 let req: Value = match serde_json::from_str(&line) {
790 Ok(v) => v,
791 Err(e) => {
792 let _ = tx.send(error_frame(
793 Value::Null,
794 -32700,
795 &format!("parse error: {e}"),
796 ));
797 continue;
798 }
799 };
800 let id = req.get("id").cloned().unwrap_or(Value::Null);
801 let method = req.get("method").and_then(|m| m.as_str()).unwrap_or("");
802 let params = req.get("params").cloned().unwrap_or(Value::Null);
803
804 if id.is_null() && method == "notifications/initialized" {
806 continue;
807 }
808
809 match method {
814 "initialize" => {
815 let _ = tx.send(handle_initialize(id));
816 }
817 "ping" => {
818 let _ = tx.send(handle_ping(id));
819 }
820 "tools/list" => {
821 let _ = tx.send(handle_tools_list(id, &tools));
822 }
823 "tools/call" => {
824 handle_tools_call(id, ¶ms, &state, &tools, &tx);
825 }
826 _ => {
827 let _ = tx.send(error_frame(
828 id,
829 -32601,
830 &format!("method not found: {method}"),
831 ));
832 }
833 }
834 }
835 {
840 let _exclusive = state.op_barrier.write().await;
841 let backend = state.backend.lock().await.take();
842 if let Some(b) = backend {
843 b.shutdown().await;
844 }
845 }
846 drop(tx);
847 let _ = writer.await;
848 Ok(())
849}
850
851fn handle_initialize(id: Value) -> Vec<u8> {
854 result_frame(
855 id,
856 json!({
857 "protocolVersion": "2024-11-05",
858 "capabilities": {"tools": {}},
859 "serverInfo": {
860 "name": "browser-control",
861 "version": env!("CARGO_PKG_VERSION"),
862 },
863 }),
864 )
865}
866
867fn handle_ping(id: Value) -> Vec<u8> {
869 result_frame(id, json!({}))
870}
871
872fn handle_tools_list(id: Value, tools: &ToolRegistry) -> Vec<u8> {
874 result_frame(id, json!({"tools": tools.list()}))
875}
876
877fn handle_tools_call(
888 id: Value,
889 params: &Value,
890 state: &ServerState,
891 tools: &ToolRegistry,
892 tx: &tokio::sync::mpsc::UnboundedSender<Vec<u8>>,
893) {
894 let name = params
895 .get("name")
896 .and_then(|v| v.as_str())
897 .unwrap_or("")
898 .to_string();
899 let args = params.get("arguments").cloned().unwrap_or(Value::Null);
900 let handler = tools.handler(&name);
901 let state = state.clone();
902 let tx = tx.clone();
903 let barrier = state.op_barrier.clone();
904 let is_exclusive = name == "browser_select";
905 tokio::spawn(async move {
906 let frame = match handler {
907 None => error_frame(id, -32602, &format!("tool not found: {name}")),
910 Some(h) => {
911 if is_exclusive {
912 let _guard = barrier.write().await;
917 match h(state, args).await {
918 Ok(v) => result_frame(id, v),
919 Err(e) => tool_error_frame(id, &e),
920 }
921 } else {
922 let _guard = barrier.read().await;
923 match h(state, args).await {
924 Ok(v) => result_frame(id, v),
925 Err(e) => tool_error_frame(id, &e),
926 }
927 }
928 }
929 };
930 let _ = tx.send(frame);
931 });
932}
933
934fn result_frame(id: Value, result: Value) -> Vec<u8> {
936 let resp = json!({"jsonrpc": "2.0", "id": id, "result": result});
937 let mut s = serde_json::to_vec(&resp).unwrap_or_else(|_| b"{}".to_vec());
938 s.push(b'\n');
939 s
940}
941
942fn tool_error_frame(id: Value, err: &anyhow::Error) -> Vec<u8> {
948 let message = tool_error_message(err);
949 result_frame(
950 id,
951 json!({
952 "content": [{ "type": "text", "text": message }],
953 "isError": true,
954 }),
955 )
956}
957
958fn tool_error_message(err: &anyhow::Error) -> String {
963 use crate::errors::SessionError;
964 use crate::registry::bidi_lock::BidiLockBusy;
965
966 if let Some(se) = err.downcast_ref::<SessionError>() {
967 return se.to_string();
970 }
971 if let Some(busy) = err.downcast_ref::<BidiLockBusy>() {
972 return busy.to_string();
973 }
974 format!("{err:#}")
977}
978
979fn error_frame(id: Value, code: i64, message: &str) -> Vec<u8> {
981 let resp = json!({
982 "jsonrpc": "2.0",
983 "id": id,
984 "error": {"code": code, "message": message},
985 });
986 let mut s = serde_json::to_vec(&resp).unwrap_or_else(|_| b"{}".to_vec());
987 s.push(b'\n');
988 s
989}
990
991#[cfg(test)]
992mod tests {
993 use super::*;
994 use crate::cli::env_resolver::Source;
995 use crate::detect::Engine;
996 use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
997
998 fn dummy_resolved() -> ResolvedBrowser {
999 ResolvedBrowser {
1000 endpoint: "ws://localhost:9999".into(),
1001 engine: Engine::Cdp,
1002 source: Source::External,
1003 }
1004 }
1005
1006 fn dummy_state() -> ServerState {
1007 ServerState::new(dummy_resolved())
1008 }
1009
1010 async fn send_recv(tools: ToolRegistry, requests: &[Value]) -> Vec<Value> {
1011 let (mut client_w, server_r) = tokio::io::duplex(8192);
1012 let (server_w, client_r) = tokio::io::duplex(8192);
1013 let state = dummy_state();
1014 let join = tokio::spawn(async move {
1015 let _ = run_with_streams(state, tools, server_r, server_w).await;
1016 });
1017
1018 for req in requests {
1019 let mut s = serde_json::to_vec(req).unwrap();
1020 s.push(b'\n');
1021 client_w.write_all(&s).await.unwrap();
1022 }
1023 drop(client_w);
1025
1026 let mut reader = BufReader::new(client_r);
1027 let mut responses = Vec::new();
1028 loop {
1029 let mut line = String::new();
1030 let n = reader.read_line(&mut line).await.unwrap();
1031 if n == 0 {
1032 break;
1033 }
1034 responses.push(serde_json::from_str(line.trim()).unwrap());
1035 }
1036 let _ = join.await;
1037 responses
1038 }
1039
1040 fn echo_tool() -> RegisteredTool {
1041 RegisteredTool {
1042 name: "echo".to_string(),
1043 description: "Echo arguments back".to_string(),
1044 input_schema: json!({"type": "object"}),
1045 handler: std::sync::Arc::new(|_state, args| {
1046 Box::pin(async move { Ok(json!({"echoed": args})) })
1047 }),
1048 }
1049 }
1050
1051 fn failing_tool() -> RegisteredTool {
1052 RegisteredTool {
1053 name: "boom".to_string(),
1054 description: "Always fails with a typed SessionError".to_string(),
1055 input_schema: json!({"type": "object"}),
1056 handler: std::sync::Arc::new(|_state, _args| {
1057 Box::pin(async move {
1058 Err(anyhow::Error::new(crate::errors::SessionError::TabHung {
1059 target_id: Some("T1".into()),
1060 url: Some("https://example.test".into()),
1061 timeout_ms: 20_000,
1062 hint: "renderer wedged",
1063 }))
1064 })
1065 }),
1066 }
1067 }
1068
1069 #[tokio::test]
1070 async fn initialize_round_trip() {
1071 let resp = send_recv(
1072 ToolRegistry::new(),
1073 &[json!({"jsonrpc":"2.0","id":1,"method":"initialize","params":{}})],
1074 )
1075 .await;
1076 assert_eq!(resp.len(), 1);
1077 assert_eq!(resp[0]["id"], 1);
1078 assert_eq!(resp[0]["result"]["protocolVersion"], "2024-11-05");
1079 assert_eq!(resp[0]["result"]["serverInfo"]["name"], "browser-control");
1080 }
1081
1082 #[tokio::test]
1083 async fn tools_list_empty() {
1084 let resp = send_recv(
1085 ToolRegistry::new(),
1086 &[json!({"jsonrpc":"2.0","id":2,"method":"tools/list"})],
1087 )
1088 .await;
1089 assert_eq!(resp[0]["result"]["tools"], json!([]));
1090 }
1091
1092 #[tokio::test]
1093 async fn tools_list_after_register() {
1094 let tools = ToolRegistry::new();
1095 tools.register(echo_tool());
1096 let resp = send_recv(
1097 tools,
1098 &[json!({"jsonrpc":"2.0","id":3,"method":"tools/list"})],
1099 )
1100 .await;
1101 let list = resp[0]["result"]["tools"].as_array().unwrap();
1102 assert_eq!(list.len(), 1);
1103 assert_eq!(list[0]["name"], "echo");
1104 }
1105
1106 #[tokio::test]
1107 async fn tools_call_unknown_errors() {
1108 let resp = send_recv(
1109 ToolRegistry::new(),
1110 &[json!({
1111 "jsonrpc":"2.0","id":4,"method":"tools/call",
1112 "params":{"name":"nope","arguments":{}}
1113 })],
1114 )
1115 .await;
1116 assert_eq!(resp[0]["error"]["code"], -32602);
1118 assert!(resp[0]["error"]["message"]
1119 .as_str()
1120 .unwrap()
1121 .contains("nope"));
1122 }
1123
1124 #[tokio::test]
1125 async fn tool_failure_returns_iserror_result_not_protocol_error() {
1126 let tools = ToolRegistry::new();
1127 tools.register(failing_tool());
1128 let resp = send_recv(
1129 tools,
1130 &[json!({
1131 "jsonrpc":"2.0","id":7,"method":"tools/call",
1132 "params":{"name":"boom","arguments":{}}
1133 })],
1134 )
1135 .await;
1136 assert!(resp[0]["error"].is_null());
1139 assert_eq!(resp[0]["result"]["isError"], true);
1140 let text = resp[0]["result"]["content"][0]["text"].as_str().unwrap();
1141 assert!(text.contains("tab hung"), "got: {text}");
1143 assert!(text.contains("renderer wedged"), "got: {text}");
1144 assert!(text.contains("T1"), "got: {text}");
1145 }
1146
1147 #[tokio::test]
1148 async fn tools_call_registered_returns_result() {
1149 let tools = ToolRegistry::new();
1150 tools.register(echo_tool());
1151 let resp = send_recv(
1152 tools,
1153 &[json!({
1154 "jsonrpc":"2.0","id":5,"method":"tools/call",
1155 "params":{"name":"echo","arguments":{"hello":"world"}}
1156 })],
1157 )
1158 .await;
1159 assert_eq!(resp[0]["result"]["echoed"], json!({"hello":"world"}));
1160 }
1161
1162 #[tokio::test]
1163 async fn unknown_method_returns_minus_32601() {
1164 let resp = send_recv(
1165 ToolRegistry::new(),
1166 &[json!({"jsonrpc":"2.0","id":6,"method":"bogus"})],
1167 )
1168 .await;
1169 assert_eq!(resp[0]["error"]["code"], -32601);
1170 }
1171
1172 #[tokio::test]
1173 async fn ping_returns_empty_object() {
1174 let resp = send_recv(
1175 ToolRegistry::new(),
1176 &[json!({"jsonrpc":"2.0","id":7,"method":"ping"})],
1177 )
1178 .await;
1179 assert_eq!(resp[0]["result"], json!({}));
1180 }
1181
1182 use futures_util::{SinkExt, StreamExt};
1191 use std::sync::atomic::{AtomicUsize, Ordering};
1192 use tokio_tungstenite::tungstenite::Message;
1193
1194 async fn spawn_counting_cdp_mock(
1198 live: Vec<String>,
1199 ) -> (String, Arc<AtomicUsize>, tokio::sync::oneshot::Sender<()>) {
1200 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
1201 let addr = listener.local_addr().unwrap();
1202 let (stop_tx, mut stop_rx) = tokio::sync::oneshot::channel::<()>();
1203 let conns = Arc::new(AtomicUsize::new(0));
1204 let conns_srv = conns.clone();
1205 tokio::spawn(async move {
1206 loop {
1207 let accept = tokio::select! {
1208 _ = &mut stop_rx => break,
1209 a = listener.accept() => a,
1210 };
1211 let (stream, _) = match accept {
1212 Ok(s) => s,
1213 Err(_) => break,
1214 };
1215 conns_srv.fetch_add(1, Ordering::SeqCst);
1216 let live = live.clone();
1217 tokio::spawn(async move {
1218 let mut ws = match tokio_tungstenite::accept_async(stream).await {
1219 Ok(w) => w,
1220 Err(_) => return,
1221 };
1222 while let Some(Ok(msg)) = ws.next().await {
1223 if let Message::Text(t) = msg {
1224 let req: Value = match serde_json::from_str(&t) {
1225 Ok(v) => v,
1226 Err(_) => continue,
1227 };
1228 let id = req["id"].as_u64().unwrap_or(0);
1229 let method = req["method"].as_str().unwrap_or("");
1230 let result = match method {
1231 "Target.getTargets" => {
1232 let infos: Vec<Value> = live
1233 .iter()
1234 .map(|tid| {
1235 json!({"targetId": tid, "type": "page", "url": ""})
1236 })
1237 .collect();
1238 json!({"targetInfos": infos})
1239 }
1240 "Target.attachToTarget" => json!({"sessionId": "S1"}),
1241 "Runtime.evaluate" => json!({"result": {"value": 1}}),
1242 _ => json!({}),
1243 };
1244 let resp = json!({"id": id, "result": result});
1245 if ws.send(Message::Text(resp.to_string())).await.is_err() {
1246 break;
1247 }
1248 }
1249 }
1250 });
1251 }
1252 });
1253 (format!("ws://{addr}"), conns, stop_tx)
1254 }
1255
1256 fn registered_state(name: &str, endpoint: &str) -> ServerState {
1259 ServerState::new(ResolvedBrowser {
1260 endpoint: endpoint.to_string(),
1261 engine: Engine::Cdp,
1262 source: Source::Registered { name: name.into() },
1263 })
1264 }
1265
1266 fn register_browser_row(name: &str, endpoint: &str) {
1267 let port: u16 = endpoint
1268 .strip_prefix("ws://127.0.0.1:")
1269 .and_then(|s| s.split('/').next())
1270 .and_then(|s| s.parse().ok())
1271 .expect("parse mock port");
1272 let reg = crate::registry::Registry::open().unwrap();
1273 reg.insert(&crate::registry::BrowserRow {
1274 name: name.to_string(),
1275 kind: crate::detect::Kind::Chrome,
1276 engine: Engine::Cdp,
1277 pid: std::process::id(),
1278 endpoint: endpoint.to_string(),
1279 port,
1280 profile_dir: std::path::PathBuf::from("/tmp/profiles/bx"),
1281 executable: std::path::PathBuf::from("/usr/bin/example"),
1282 headless: false,
1283 started_at: "2024-01-01T00:00:00Z".into(),
1284 })
1285 .unwrap();
1286 }
1287
1288 #[allow(clippy::await_holding_lock)]
1291 #[tokio::test]
1292 async fn resolve_named_tab_live_resolves_and_touches() {
1293 let _g = crate::test_support::ENV_LOCK
1294 .lock()
1295 .unwrap_or_else(|e| e.into_inner());
1296 let tmp = tempfile::TempDir::new().unwrap();
1297 std::env::set_var("BROWSER_CONTROL_DATA_DIR", tmp.path());
1298
1299 {
1301 let reg = crate::registry::Registry::open().unwrap();
1302 reg.tab_upsert("bx", "work", "T1", "about:blank", true)
1303 .unwrap();
1304 }
1305
1306 let (url, _conns, _stop) = spawn_counting_cdp_mock(vec!["T1".into()]).await;
1307 register_browser_row("bx", &url);
1308 let state = registered_state("bx", &url);
1309
1310 let target_id = match state.resolve_target_for_args(&json!({"tab": "work"})).await {
1311 Ok((_backend, tid)) => tid,
1312 Err(e) => panic!("named tab should resolve: {e:#}"),
1313 };
1314 assert_eq!(target_id, "T1");
1315
1316 std::env::remove_var("BROWSER_CONTROL_DATA_DIR");
1317 }
1318
1319 #[allow(clippy::await_holding_lock)]
1321 #[tokio::test]
1322 async fn resolve_named_tab_stale_sweeps_and_errors() {
1323 let _g = crate::test_support::ENV_LOCK
1324 .lock()
1325 .unwrap_or_else(|e| e.into_inner());
1326 let tmp = tempfile::TempDir::new().unwrap();
1327 std::env::set_var("BROWSER_CONTROL_DATA_DIR", tmp.path());
1328
1329 {
1333 let reg = crate::registry::Registry::open().unwrap();
1334 reg.tab_upsert("bx", "gone", "T_DEAD", "about:blank", true)
1335 .unwrap();
1336 }
1337
1338 let (url, _conns, _stop) = spawn_counting_cdp_mock(vec!["T1".into()]).await;
1339 register_browser_row("bx", &url);
1340 let state = registered_state("bx", &url);
1341
1342 let err = match state.resolve_target_for_args(&json!({"tab": "gone"})).await {
1343 Ok(_) => panic!("stale named tab must error"),
1344 Err(e) => e,
1345 };
1346 let typed = err
1347 .downcast_ref::<crate::errors::SessionError>()
1348 .expect("typed SessionError");
1349 assert!(
1350 matches!(typed, crate::errors::SessionError::TabNotFound { .. }),
1351 "expected TabNotFound, got {typed:?}"
1352 );
1353
1354 let reg = crate::registry::Registry::open().unwrap();
1356 assert!(
1357 reg.tab_get("bx", "gone").unwrap().is_none(),
1358 "row not swept"
1359 );
1360
1361 std::env::remove_var("BROWSER_CONTROL_DATA_DIR");
1362 }
1363
1364 #[allow(clippy::await_holding_lock)]
1366 #[tokio::test]
1367 async fn resolve_named_tab_missing_row_errors() {
1368 let _g = crate::test_support::ENV_LOCK
1369 .lock()
1370 .unwrap_or_else(|e| e.into_inner());
1371 let tmp = tempfile::TempDir::new().unwrap();
1372 std::env::set_var("BROWSER_CONTROL_DATA_DIR", tmp.path());
1373
1374 let (url, _conns, _stop) = spawn_counting_cdp_mock(vec!["T1".into()]).await;
1375 register_browser_row("bx", &url);
1376 let state = registered_state("bx", &url);
1377
1378 let err = match state.resolve_target_for_args(&json!({"tab": "nope"})).await {
1379 Ok(_) => panic!("unknown named tab must error"),
1380 Err(e) => e,
1381 };
1382 let typed = err
1383 .downcast_ref::<crate::errors::SessionError>()
1384 .expect("typed SessionError");
1385 assert!(
1386 matches!(typed, crate::errors::SessionError::TabNotFound { .. }),
1387 "expected TabNotFound, got {typed:?}"
1388 );
1389
1390 std::env::remove_var("BROWSER_CONTROL_DATA_DIR");
1391 }
1392
1393 #[tokio::test]
1394 async fn resolve_tab_and_target_mutually_exclusive() {
1395 let state = dummy_state();
1397 let err = match state
1398 .resolve_target_for_args(&json!({"tab": "a", "target": "b"}))
1399 .await
1400 {
1401 Ok(_) => panic!("tab+target must error"),
1402 Err(e) => e,
1403 };
1404 assert!(
1405 err.to_string().contains("mutually exclusive"),
1406 "got: {err:#}"
1407 );
1408 }
1409
1410 #[tokio::test]
1411 async fn concurrent_ensure_backend_opens_one_backend() {
1412 let (url, conns, _stop) = spawn_counting_cdp_mock(vec!["T1".into()]).await;
1416 let state = ServerState::new(ResolvedBrowser {
1417 endpoint: url,
1418 engine: Engine::Cdp,
1419 source: Source::External,
1420 });
1421
1422 let mut handles = Vec::new();
1423 for _ in 0..8 {
1424 let s = state.clone();
1425 handles.push(tokio::spawn(async move { s.ensure_backend().await }));
1426 }
1427 for h in handles {
1428 h.await.unwrap().expect("ensure_backend should succeed");
1429 }
1430 assert_eq!(
1431 conns.load(Ordering::SeqCst),
1432 1,
1433 "expected exactly one backend (one WS connection) under concurrency"
1434 );
1435 }
1436
1437 #[allow(clippy::await_holding_lock)]
1438 #[tokio::test]
1439 async fn external_cdp_backend_does_not_open_registry() {
1440 let _g = crate::test_support::ENV_LOCK
1441 .lock()
1442 .unwrap_or_else(|e| e.into_inner());
1443 let stale_path = {
1444 let tmp = tempfile::TempDir::new().unwrap();
1445 tmp.path().to_path_buf()
1446 };
1447 std::env::set_var("BROWSER_CONTROL_DATA_DIR", &stale_path);
1448
1449 let (url, _conns, _stop) = spawn_counting_cdp_mock(vec!["T1".into()]).await;
1450 let state = ServerState::new(ResolvedBrowser {
1451 endpoint: url,
1452 engine: Engine::Cdp,
1453 source: Source::External,
1454 });
1455
1456 let result = state.ensure_backend().await;
1457
1458 std::env::remove_var("BROWSER_CONTROL_DATA_DIR");
1459 result.expect("external CDP backend must not depend on registry env");
1460 }
1461
1462 #[tokio::test]
1463 async fn initialized_notification_is_silently_ignored() {
1464 let resp = send_recv(
1467 ToolRegistry::new(),
1468 &[
1469 json!({"jsonrpc":"2.0","method":"notifications/initialized"}),
1470 json!({"jsonrpc":"2.0","id":8,"method":"ping"}),
1471 ],
1472 )
1473 .await;
1474 assert_eq!(resp.len(), 1);
1475 assert_eq!(resp[0]["id"], 8);
1476 }
1477}