Skip to content

Commit 8f8f776

Browse files
committed
chore: using tokio::select and adjusting proxy config
1 parent 7a35974 commit 8f8f776

File tree

5 files changed

+29
-19
lines changed

5 files changed

+29
-19
lines changed

Cargo.lock

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

crates/apollo-mcp-proxy/Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ rmcp = { workspace = true, features = [
1616
tracing.workspace = true
1717
tracing-appender = "0.2.3"
1818
tracing-subscriber.workspace = true
19+
serde_json = "1.0.140"
1920

2021
[lints]
2122
workspace = true

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ pub async fn start_proxy_client(url: &str) -> Result<(), Box<dyn Error>> {
3030
};
3131

3232
let server_info = client.peer_info();
33-
info!("[Proxy] Connected to server");
33+
info!("[Proxy] Connected to server at {}", url);
3434
debug!("{server_info:#?}");
3535

3636
let proxy_server = ProxyServer::new(client.peer().clone(), client.peer_info());

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

Lines changed: 13 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ use apollo_mcp_server::server::Transport;
1111
use clap::Parser;
1212
use clap::builder::Styles;
1313
use clap::builder::styling::{AnsiColor, Effects};
14+
use tokio::signal;
1415
use runtime::IdOrDefault;
1516
use tracing::{Level, info, warn};
1617
use tracing_subscriber::EnvFilter;
@@ -154,15 +155,19 @@ async fn main() -> anyhow::Result<()> {
154155
.start();
155156

156157
match config.transport {
157-
Transport::StreamableHttp { proxy, proxy_endpoint, .. } => {
158+
Transport::StreamableHttp { proxy, proxy_url, address, port } => {
158159
if proxy {
159-
let mut endpoint = proxy_endpoint;
160-
if !endpoint.starts_with("http") {
161-
endpoint = format!("http://{endpoint}");
162-
}
163-
164-
let proxy_client = start_proxy_client(endpoint.as_str());
165-
let (_, _ ) = tokio::join!(mcp_server, proxy_client);
160+
let url = Transport::proxy_url(&proxy_url, &address, &port);
161+
let proxy_client = async {
162+
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
163+
_ = start_proxy_client(url.as_str()).await;
164+
};
165+
166+
tokio::select! {
167+
_ = mcp_server => {},
168+
_ = proxy_client => {},
169+
_ = signal::ctrl_c() => {}
170+
};
166171
} else {
167172
mcp_server.await?;
168173
}

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

Lines changed: 13 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ use bon::bon;
55
use reqwest::header::{CONTENT_TYPE, HeaderMap, HeaderValue};
66
use schemars::JsonSchema;
77
use serde::Deserialize;
8-
use url::Url;
8+
use url::{Url};
99

1010
use crate::custom_scalar_map::CustomScalarMap;
1111
use crate::errors::ServerError;
@@ -66,13 +66,12 @@ pub enum Transport {
6666

6767
/// Flag indicating whether to spin up a proxy server
6868
/// [default: false]
69-
#[serde(default = "Transport::default_proxy")]
69+
#[serde(default)]
7070
proxy: bool,
7171

7272
/// Proxy server endpoint to connect to if proxy is enabled
73-
/// [default: http://<address>:<port>/mcp]
74-
#[serde(default = "Transport::default_proxy_endpoint")]
75-
proxy_endpoint: String,
73+
#[serde(default)]
74+
proxy_url: Option<Url>,
7675
},
7776
}
7877

@@ -85,12 +84,16 @@ impl Transport {
8584
5000
8685
}
8786

88-
fn default_proxy() -> bool {
89-
false
90-
}
87+
pub fn proxy_url(proxy_url: &Option<Url>, address: &IpAddr, port: &u16) -> Url {
88+
match proxy_url {
89+
Some(proxy_url) => { proxy_url.clone() }
90+
None => {
91+
let address = format!("http://{address}:{port}/mcp");
92+
#[allow(clippy::unwrap_used)]
93+
Url::parse(address.as_str()).unwrap()
94+
}
95+
}
9196

92-
fn default_proxy_endpoint() -> String {
93-
format!("http://{}:{}/mcp", Self::default_address(), Self::default_port())
9497
}
9598
}
9699

0 commit comments

Comments
 (0)