chore: using tokio::select and adjusting proxy config

Author: alocay

Cargo.lock

@@ -16,6 +16,7 @@ rmcp = { workspace = true, features = [
 tracing.workspace = true
 tracing-appender = "0.2.3"
 tracing-subscriber.workspace = true
+serde_json = "1.0.140"
 
 [lints]
 workspace = true

crates/apollo-mcp-proxy/Cargo.toml

@@ -30,7 +30,7 @@ pub async fn start_proxy_client(url: &str) -> Result<(), Box<dyn Error>> {
     };
 
     let server_info = client.peer_info();
-    info!("[Proxy] Connected to server");
+    info!("[Proxy] Connected to server at {}", url);
     debug!("{server_info:#?}");
 
     let proxy_server = ProxyServer::new(client.peer().clone(), client.peer_info());

crates/apollo-mcp-proxy/src/client.rs

@@ -11,6 +11,7 @@ use apollo_mcp_server::server::Transport;
 use clap::Parser;
 use clap::builder::Styles;
 use clap::builder::styling::{AnsiColor, Effects};
+use tokio::signal;
 use runtime::IdOrDefault;
 use tracing::{Level, info, warn};
 use tracing_subscriber::EnvFilter;
@@ -159,15 +160,19 @@ async fn main() -> anyhow::Result<()> {
         .start();
 
     match config.transport {
-        Transport::StreamableHttp { proxy, proxy_endpoint, .. } => {
+        Transport::StreamableHttp { proxy, proxy_url, address, port } => {
             if proxy {
-                let mut endpoint = proxy_endpoint;
-                if !endpoint.starts_with("http") {
-                    endpoint = format!("http://{endpoint}");
-                }
-                
-                let proxy_client = start_proxy_client(endpoint.as_str());
-                let (_, _ ) = tokio::join!(mcp_server, proxy_client);
+                let url = Transport::proxy_url(&proxy_url, &address, &port);
+                let proxy_client = async {
+                    tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
+                    _ = start_proxy_client(url.as_str()).await;
+                };
+
+                tokio::select! {
+                    _ = mcp_server => {},
+                    _ = proxy_client => {},
+                    _ = signal::ctrl_c() => {}
+                };
             } else {
                 mcp_server.await?;
             }

crates/apollo-mcp-server/src/main.rs

@@ -5,7 +5,7 @@ use bon::bon;
 use reqwest::header::{CONTENT_TYPE, HeaderMap, HeaderValue};
 use schemars::JsonSchema;
 use serde::Deserialize;
-use url::Url;
+use url::{Url};
 
 use crate::custom_scalar_map::CustomScalarMap;
 use crate::errors::ServerError;
@@ -71,13 +71,12 @@ pub enum Transport {
 
         /// Flag indicating whether to spin up a proxy server
         /// [default: false]
-        #[serde(default = "Transport::default_proxy")]
+        #[serde(default)]
         proxy: bool,
 
         /// Proxy server endpoint to connect to if proxy is enabled
-        /// [default: http://<address>:<port>/mcp]
-        #[serde(default = "Transport::default_proxy_endpoint")]
-        proxy_endpoint: String,
+        #[serde(default)]
+        proxy_url: Option<Url>,
     },
 }
 
@@ -90,12 +89,16 @@ impl Transport {
         5000
     }
 
-    fn default_proxy() -> bool {
-        false
-    }
+    pub fn proxy_url(proxy_url: &Option<Url>, address: &IpAddr, port: &u16) -> Url {
+        match proxy_url {
+            Some(proxy_url) => { proxy_url.clone() }
+            None => {
+                let address = format!("http://{address}:{port}/mcp");
+                #[allow(clippy::unwrap_used)]
+                Url::parse(address.as_str()).unwrap()
+            }
+        }
 
-    fn default_proxy_endpoint() -> String {
-        format!("http://{}:{}/mcp", Self::default_address(), Self::default_port())
     }
 }