use std::marker::PhantomData;
use std::pin::Pin;
use std::task::{Context, Poll};
use futures::{Sink, Stream};
use pin_project_lite::pin_project;
#[cfg(feature = "mcp")]
use rmcp::service::ServiceRole;
#[derive(Debug, Clone)]
pub enum EnhancedTransportConfig {
WebSocket {
url: String,
headers: std::collections::HashMap<String, String>,
},
Tcp {
host: String,
port: u16,
},
#[cfg(unix)]
UnixSocket {
path: std::path::PathBuf,
},
#[cfg(windows)]
NamedPipe {
name: String,
},
}
pub trait EnhancedTransport<R>: Send + Sync {
fn description(&self) -> String;
fn is_connected(&self) -> bool;
fn stats(&self) -> TransportStats;
}
#[derive(Debug, Clone, Default)]
pub struct TransportStats {
pub messages_sent: u64,
pub messages_received: u64,
pub connection_errors: u64,
pub last_error: Option<String>,
}
pin_project! {
pub struct GenericTransport<R, S, E> {
#[pin]
stream: S,
stats: TransportStats,
description: String,
marker: PhantomData<(fn() -> E, fn() -> R)>,
}
}
impl<R, S, E> GenericTransport<R, S, E> {
pub fn new(stream: S, description: String) -> Self {
Self {
stream,
stats: TransportStats::default(),
description,
marker: PhantomData,
}
}
fn stats_mut(&mut self) -> &mut TransportStats {
&mut self.stats
}
}
impl<R, S, E> EnhancedTransport<R> for GenericTransport<R, S, E>
where
R: Send + Sync,
S: Send + Sync,
E: Send + Sync,
{
fn description(&self) -> String {
self.description.clone()
}
fn is_connected(&self) -> bool {
true
}
fn stats(&self) -> TransportStats {
self.stats.clone()
}
}
pub struct TransportFactory;
impl TransportFactory {
pub async fn create_transport<R>(
config: &EnhancedTransportConfig,
) -> Result<Box<dyn EnhancedTransport<R>>, TransportError>
where
R: Send + Sync + 'static + rmcp::service::ServiceRole,
{
match config {
EnhancedTransportConfig::WebSocket { url, headers: _ } => {
#[cfg(feature = "websocket")]
{
let transport = super::websocket::WebSocketTransport::connect(url).await?;
Ok(Box::new(transport) as Box<dyn EnhancedTransport<R>>)
}
#[cfg(not(feature = "websocket"))]
{
Err(TransportError::UnsupportedTransport("WebSocket feature not enabled".to_string()))
}
}
EnhancedTransportConfig::Tcp { host: _, port: _ } => {
Err(TransportError::UnsupportedTransport("TCP transport not yet implemented".to_string()))
}
#[cfg(unix)]
EnhancedTransportConfig::UnixSocket { path } => {
let transport = UnixSocketTransport::connect(path).await?;
Ok(Box::new(transport))
}
#[cfg(windows)]
EnhancedTransportConfig::NamedPipe { name: _ } => {
Err(TransportError::UnsupportedTransport("Named pipe transport not yet implemented".to_string()))
}
}
}
}
#[derive(Debug, thiserror::Error)]
pub enum TransportError {
#[error("IO error: {0}")]
Io(#[from] std::io::Error),
#[error("WebSocket error: {0}")]
#[cfg(feature = "websocket")]
WebSocket(#[from] tokio_tungstenite::tungstenite::Error),
#[error("Serialization error: {0}")]
Serialization(#[from] serde_json::Error),
#[error("Unsupported transport: {0}")]
UnsupportedTransport(String),
#[error("Connection failed: {0}")]
ConnectionFailed(String),
}
pub struct TcpTransport<R> {
inner: GenericTransport<R, tokio::net::TcpStream, std::io::Error>,
}
impl<R> TcpTransport<R>
where
R: ServiceRole,
{
pub async fn connect(host: &str, port: u16) -> Result<Self, TransportError> {
let addr = format!("{}:{}", host, port);
let stream = tokio::net::TcpStream::connect(&addr).await?;
let description = format!("TCP connection to {}", addr);
todo!("TCP transport implementation needs proper framing")
}
}
#[cfg(feature = "websocket")]
pub struct WebSocketTransport<R> {
inner: GenericTransport<R, tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>, tokio_tungstenite::tungstenite::Error>,
}
#[cfg(feature = "websocket")]
impl<R> WebSocketTransport<R>
where
R: ServiceRole,
{
pub async fn connect(url: &str) -> Result<Self, TransportError> {
let (stream, _) = tokio_tungstenite::connect_async(url).await?;
let description = format!("WebSocket connection to {}", url);
todo!("WebSocket transport implementation needs proper message handling")
}
}
#[cfg(unix)]
pub struct UnixSocketTransport<R> {
inner: GenericTransport<R, tokio::net::UnixStream, std::io::Error>,
}
#[cfg(unix)]
impl<R> UnixSocketTransport<R>
where
R: ServiceRole,
{
pub async fn connect(path: &std::path::Path) -> Result<Self, TransportError> {
let stream = tokio::net::UnixStream::connect(path).await?;
let description = format!("Unix socket connection to {}", path.display());
todo!("Unix socket transport implementation needs proper framing")
}
}
#[cfg(windows)]
pub struct NamedPipeTransport<R> {
_marker: PhantomData<R>,
}
#[cfg(windows)]
impl<R> NamedPipeTransport<R>
where
R: ServiceRole,
{
pub async fn connect(name: &str) -> Result<Self, TransportError> {
todo!("Named pipe transport implementation")
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_transport_config() {
let config = EnhancedTransportConfig::WebSocket {
url: "ws://localhost:8080".to_string(),
headers: std::collections::HashMap::new(),
};
match config {
EnhancedTransportConfig::WebSocket { url, .. } => {
assert_eq!(url, "ws://localhost:8080");
}
_ => panic!("Expected WebSocket config"),
}
}
#[test]
fn test_transport_stats() {
let stats = TransportStats::default();
assert_eq!(stats.messages_sent, 0);
assert_eq!(stats.messages_received, 0);
assert_eq!(stats.connection_errors, 0);
assert!(stats.last_error.is_none());
}
}