refactor: cleanup

src-tauri/src/mcp_server.rs

@@ -184,19 +184,17 @@ impl McpServer {
 
   pub async fn start(&self, app_handle: AppHandle) -> Result<u16, String> {
     if !WayfernTermsManager::instance().is_terms_accepted() {
-      return Err(
-        "Wayfern Terms and Conditions must be accepted before starting MCP server".to_string(),
-      );
+      return Err(crate::backend_error("WAYFERN_TERMS_REQUIRED"));
     }
 
     if self.is_running() {
-      return Err("MCP server is already running".to_string());
+      return Err(crate::backend_error("MCP_SERVER_ALREADY_RUNNING"));
     }
 
     let settings_manager = SettingsManager::instance();
     let settings = settings_manager
       .load_settings()
-      .map_err(|e| format!("Failed to load settings: {e}"))?;
+      .map_err(|e| crate::backend_error_with_detail("INTERNAL_ERROR", e))?;
 
     // Get or generate token
     let existing_token = settings_manager
@@ -211,12 +209,16 @@ impl McpServer {
       settings_manager
         .generate_mcp_token(&app_handle)
         .await
-        .map_err(|e| format!("Failed to generate MCP token: {e}"))?
+        .map_err(|e| crate::backend_error_with_detail("INTERNAL_ERROR", e))?
     };
 
     // Determine port (use saved port, or try default, or random)
     let preferred_port = settings.mcp_port.unwrap_or(DEFAULT_MCP_PORT);
-    let actual_port = self.bind_to_available_port(preferred_port).await?;
+    let listener = self.bind_to_available_port(preferred_port).await?;
+    let actual_port = listener
+      .local_addr()
+      .map_err(|e| crate::backend_error_with_detail("INTERNAL_ERROR", e))?
+      .port();
 
     // Save port if it changed
     if settings.mcp_port != Some(actual_port) {
@@ -224,7 +226,7 @@ impl McpServer {
       new_settings.mcp_port = Some(actual_port);
       settings_manager
         .save_settings(&new_settings)
-        .map_err(|e| format!("Failed to save settings: {e}"))?;
+        .map_err(|e| crate::backend_error_with_detail("INTERNAL_ERROR", e))?;
     }
 
     // Store state
@@ -244,31 +246,31 @@ impl McpServer {
       server: McpServer::instance(),
       token,
     };
-    tokio::spawn(Self::run_http_server(actual_port, http_state, shutdown_rx));
+    tokio::spawn(Self::run_http_server(listener, http_state, shutdown_rx));
 
     log::info!("[mcp] Server started on port {}", actual_port);
     Ok(actual_port)
   }
 
-  async fn bind_to_available_port(&self, preferred: u16) -> Result<u16, String> {
+  async fn bind_to_available_port(&self, preferred: u16) -> Result<TcpListener, String> {
     let addr = SocketAddr::from(([127, 0, 0, 1], preferred));
-    if TcpListener::bind(addr).await.is_ok() {
-      return Ok(preferred);
+    if let Ok(listener) = TcpListener::bind(addr).await {
+      return Ok(listener);
     }
 
     for _ in 0..10 {
       let port = 51000 + (rand::random::<u16>() % 1000);
       let addr = SocketAddr::from(([127, 0, 0, 1], port));
-      if TcpListener::bind(addr).await.is_ok() {
-        return Ok(port);
+      if let Ok(listener) = TcpListener::bind(addr).await {
+        return Ok(listener);
       }
     }
 
-    Err("Could not find available port for MCP server".to_string())
+    Err(crate::backend_error("MCP_PORT_UNAVAILABLE"))
   }
 
   async fn run_http_server(
-    port: u16,
+    listener: TcpListener,
     state: McpHttpState,
     shutdown_rx: tokio::sync::oneshot::Receiver<()>,
   ) {
@@ -286,28 +288,17 @@ impl McpServer {
           .delete(Self::handle_mcp_delete),
       )
       .route("/health", get(Self::handle_health))
-      // Inert chokepoint (innermost → runs after auth) for the future per-hour
-      // automation request limit. See rate_limit_middleware.
-      .layer(middleware::from_fn(Self::rate_limit_middleware))
       .layer(middleware::from_fn_with_state(
         state.clone(),
         Self::auth_middleware,
       ))
       .with_state(state);
 
-    let addr = SocketAddr::from(([127, 0, 0, 1], port));
-
-    let server = async {
-      match TcpListener::bind(addr).await {
-        Ok(listener) => {
-          log::info!("[mcp] Server listening on http://127.0.0.1:{}/mcp", port);
-          if let Err(e) = axum::serve(listener, app).await {
-            log::error!("[mcp] Server error: {}", e);
-          }
-        }
-        Err(e) => {
-          log::error!("[mcp] Failed to bind on port {}: {}", port, e);
-        }
+    let port = listener.local_addr().map(|addr| addr.port()).unwrap_or(0);
+    let server = async move {
+      log::info!("[mcp] Server listening on http://127.0.0.1:{}/mcp", port);
+      if let Err(e) = axum::serve(listener, app).await {
+        log::error!("[mcp] Server error: {}", e);
       }
     };
 
@@ -319,17 +310,6 @@ impl McpServer {
     }
   }
 
-  /// Chokepoint for the future per-hour automation request limit, mirroring the
-  /// REST API's. The limit (`requests_per_hour`, default 100) is plumbed through
-  /// entitlements; this is intentionally inert today — it resolves the limit but
-  /// never blocks. To enforce, count authenticated tool calls per rolling hour
-  /// and return StatusCode::TOO_MANY_REQUESTS once the limit (when > 0) is hit.
-  async fn rate_limit_middleware(req: Request<Body>, next: Next) -> Result<Response, StatusCode> {
-    let _requests_per_hour = CLOUD_AUTH.requests_per_hour().await;
-    // TODO(rate-limit): enforce `_requests_per_hour` for MCP tool calls.
-    Ok(next.run(req).await)
-  }
-
   async fn auth_middleware(
     State(state): State<McpHttpState>,
     req: Request<Body>,
@@ -476,14 +456,65 @@ impl McpServer {
         }
       }
 
+      if Self::is_automation_tool_call(&request) {
+        if let crate::automation_rate_limiter::RateLimitOutcome::Limited { retry_after_secs } =
+          crate::automation_rate_limiter::check_automation_rate_limit().await
+        {
+          log::warn!(
+            "[mcp] Rejected tools/call: automation rate limit exceeded; retry in {}s",
+            retry_after_secs
+          );
+          return (
+            StatusCode::TOO_MANY_REQUESTS,
+            [(header::RETRY_AFTER, retry_after_secs.to_string())],
+            "automation request rate limit exceeded",
+          )
+            .into_response();
+        }
+      }
+
       let response = state.server.handle_request(request).await;
       Json(response).into_response()
     }
   }
 
+  fn is_automation_tool_call(request: &McpRequest) -> bool {
+    if request.method != "tools/call" {
+      return false;
+    }
+
+    let Some(tool_name) = request
+      .params
+      .as_ref()
+      .and_then(|params| params.get("name"))
+      .and_then(|name| name.as_str())
+    else {
+      return false;
+    };
+
+    matches!(
+      tool_name,
+      "run_profile"
+        | "kill_profile"
+        | "batch_run_profiles"
+        | "batch_stop_profiles"
+        | "start_sync_session"
+        | "navigate"
+        | "screenshot"
+        | "evaluate_javascript"
+        | "click_element"
+        | "type_text"
+        | "get_page_content"
+        | "get_page_info"
+        | "get_interactive_elements"
+        | "click_by_index"
+        | "type_by_index"
+    )
+  }
+
   pub async fn stop(&self) -> Result<(), String> {
     if !self.is_running() {
-      return Err("MCP server is not running".to_string());
+      return Err(crate::backend_error("MCP_SERVER_NOT_RUNNING"));
     }
 
     let mut inner = self.inner.lock().await;
@@ -5528,4 +5559,46 @@ mod tests {
     let server = McpServer::new();
     assert!(!server.is_running());
   }
+
+  #[test]
+  fn rate_limit_only_classifies_browser_automation_tools() {
+    let request = |method: &str, name: Option<&str>| McpRequest {
+      jsonrpc: "2.0".to_string(),
+      id: Some(serde_json::json!(1)),
+      method: method.to_string(),
+      params: name.map(|name| serde_json::json!({ "name": name, "arguments": {} })),
+    };
+
+    for name in [
+      "run_profile",
+      "kill_profile",
+      "batch_run_profiles",
+      "batch_stop_profiles",
+      "start_sync_session",
+      "navigate",
+      "screenshot",
+      "evaluate_javascript",
+      "click_element",
+      "type_text",
+      "get_page_content",
+      "get_page_info",
+      "get_interactive_elements",
+      "click_by_index",
+      "type_by_index",
+    ] {
+      assert!(
+        McpServer::is_automation_tool_call(&request("tools/call", Some(name))),
+        "automation tool was not limited: {name}"
+      );
+    }
+
+    assert!(!McpServer::is_automation_tool_call(&request(
+      "tools/call",
+      Some("list_profiles")
+    )));
+    assert!(!McpServer::is_automation_tool_call(&request(
+      "tools/list",
+      None
+    )));
+  }
 }