Refactor
This commit is contained in:
Danielle Jenkins
parent
5fe28f7fe7
commit
4171c923e6
4 changed files with 0 additions and 222 deletions
|
|
@ -1,139 +0,0 @@
|
||||||
use axum::{
|
|
||||||
body::Body,
|
|
||||||
extract::{Query, State},
|
|
||||||
http::StatusCode,
|
|
||||||
response::sse::{Event, Sse},
|
|
||||||
routing::get,
|
|
||||||
Router,
|
|
||||||
};
|
|
||||||
|
|
||||||
#[cfg(test)]
|
|
||||||
mod tests;
|
|
||||||
use futures::{stream::Stream, StreamExt, TryStreamExt};
|
|
||||||
use mcp_server::{ByteTransport, Server};
|
|
||||||
use std::collections::HashMap;
|
|
||||||
use tokio_util::codec::FramedRead;
|
|
||||||
|
|
||||||
use anyhow::Result;
|
|
||||||
use mcp_server::router::RouterService;
|
|
||||||
use crate::{jsonrpc_frame_codec::JsonRpcFrameCodec, DocRouter};
|
|
||||||
use std::sync::Arc;
|
|
||||||
use tokio::{
|
|
||||||
io::{self, AsyncWriteExt},
|
|
||||||
sync::Mutex,
|
|
||||||
};
|
|
||||||
|
|
||||||
type C2SWriter = Arc<Mutex<io::WriteHalf<io::SimplexStream>>>;
|
|
||||||
type SessionId = Arc<str>;
|
|
||||||
|
|
||||||
#[derive(Clone, Default)]
|
|
||||||
pub struct App {
|
|
||||||
txs: Arc<tokio::sync::RwLock<HashMap<SessionId, C2SWriter>>>,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl App {
|
|
||||||
pub fn new() -> Self {
|
|
||||||
Self {
|
|
||||||
txs: Default::default(),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
pub fn router(&self) -> Router {
|
|
||||||
Router::new()
|
|
||||||
.route("/sse", get(sse_handler).post(post_event_handler))
|
|
||||||
.with_state(self.clone())
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn session_id() -> SessionId {
|
|
||||||
let id = format!("{:016x}", rand::random::<u128>());
|
|
||||||
Arc::from(id)
|
|
||||||
}
|
|
||||||
|
|
||||||
#[derive(Debug, serde::Deserialize)]
|
|
||||||
#[serde(rename_all = "camelCase")]
|
|
||||||
pub struct PostEventQuery {
|
|
||||||
pub session_id: String,
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn post_event_handler(
|
|
||||||
State(app): State<App>,
|
|
||||||
Query(PostEventQuery { session_id }): Query<PostEventQuery>,
|
|
||||||
body: Body,
|
|
||||||
) -> Result<StatusCode, StatusCode> {
|
|
||||||
const BODY_BYTES_LIMIT: usize = 1 << 22;
|
|
||||||
let write_stream = {
|
|
||||||
let rg = app.txs.read().await;
|
|
||||||
rg.get(session_id.as_str())
|
|
||||||
.ok_or(StatusCode::NOT_FOUND)?
|
|
||||||
.clone()
|
|
||||||
};
|
|
||||||
let mut write_stream = write_stream.lock().await;
|
|
||||||
let mut body = body.into_data_stream();
|
|
||||||
if let (_, Some(size)) = body.size_hint() {
|
|
||||||
if size > BODY_BYTES_LIMIT {
|
|
||||||
return Err(StatusCode::PAYLOAD_TOO_LARGE);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
// calculate the body size
|
|
||||||
let mut size = 0;
|
|
||||||
while let Some(chunk) = body.next().await {
|
|
||||||
let Ok(chunk) = chunk else {
|
|
||||||
return Err(StatusCode::BAD_REQUEST);
|
|
||||||
};
|
|
||||||
size += chunk.len();
|
|
||||||
if size > BODY_BYTES_LIMIT {
|
|
||||||
return Err(StatusCode::PAYLOAD_TOO_LARGE);
|
|
||||||
}
|
|
||||||
write_stream
|
|
||||||
.write_all(&chunk)
|
|
||||||
.await
|
|
||||||
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
|
|
||||||
}
|
|
||||||
write_stream
|
|
||||||
.write_u8(b'\n')
|
|
||||||
.await
|
|
||||||
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
|
|
||||||
Ok(StatusCode::ACCEPTED)
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn sse_handler(State(app): State<App>) -> Sse<impl Stream<Item = Result<Event, io::Error>>> {
|
|
||||||
// it's 4KB
|
|
||||||
const BUFFER_SIZE: usize = 1 << 12;
|
|
||||||
let session = session_id();
|
|
||||||
tracing::info!(%session, "sse connection");
|
|
||||||
let (c2s_read, c2s_write) = tokio::io::simplex(BUFFER_SIZE);
|
|
||||||
let (s2c_read, s2c_write) = tokio::io::simplex(BUFFER_SIZE);
|
|
||||||
app.txs
|
|
||||||
.write()
|
|
||||||
.await
|
|
||||||
.insert(session.clone(), Arc::new(Mutex::new(c2s_write)));
|
|
||||||
{
|
|
||||||
let app_clone = app.clone();
|
|
||||||
let session = session.clone();
|
|
||||||
tokio::spawn(async move {
|
|
||||||
let router = RouterService(DocRouter::new());
|
|
||||||
let server = Server::new(router);
|
|
||||||
let bytes_transport = ByteTransport::new(c2s_read, s2c_write);
|
|
||||||
let _result = server
|
|
||||||
.run(bytes_transport)
|
|
||||||
.await
|
|
||||||
.inspect_err(|e| tracing::error!(?e, "server run error"));
|
|
||||||
app_clone.txs.write().await.remove(&session);
|
|
||||||
});
|
|
||||||
}
|
|
||||||
|
|
||||||
let stream = futures::stream::once(futures::future::ok(
|
|
||||||
Event::default()
|
|
||||||
.event("endpoint")
|
|
||||||
.data(format!("?sessionId={session}")),
|
|
||||||
))
|
|
||||||
.chain(
|
|
||||||
FramedRead::new(s2c_read, JsonRpcFrameCodec)
|
|
||||||
.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
|
|
||||||
.and_then(move |bytes| match std::str::from_utf8(bytes.as_ref()) {
|
|
||||||
Ok(message) => futures::future::ok(Event::default().event("message").data(message)),
|
|
||||||
Err(e) => futures::future::err(io::Error::new(io::ErrorKind::InvalidData, e)),
|
|
||||||
}),
|
|
||||||
);
|
|
||||||
Sse::new(stream)
|
|
||||||
}
|
|
||||||
|
|
@ -1,63 +0,0 @@
|
||||||
use super::*;
|
|
||||||
use axum::{
|
|
||||||
body::Body,
|
|
||||||
http::{Method, Request},
|
|
||||||
};
|
|
||||||
use tokio::sync::RwLock;
|
|
||||||
// Comment out tower imports for now, as we'll handle router testing differently
|
|
||||||
// use tower::Service;
|
|
||||||
// use tower::util::ServiceExt;
|
|
||||||
|
|
||||||
// Helper function to create an App with an empty state
|
|
||||||
fn create_test_app() -> App {
|
|
||||||
App {
|
|
||||||
txs: Arc::new(RwLock::new(HashMap::new())),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn test_app_initialization() {
|
|
||||||
let app = App::new();
|
|
||||||
// App should be created with an empty hashmap
|
|
||||||
assert_eq!(app.txs.read().await.len(), 0);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn test_router_setup() {
|
|
||||||
let app = App::new();
|
|
||||||
let _router = app.router();
|
|
||||||
|
|
||||||
// Check if the router is constructed properly
|
|
||||||
// This is a basic test to ensure the router is created without panics
|
|
||||||
// Just check that the router exists, no need to invoke methods
|
|
||||||
assert!(true);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn test_session_id_generation() {
|
|
||||||
// Generate two session IDs and ensure they're different
|
|
||||||
let id1 = session_id();
|
|
||||||
let id2 = session_id();
|
|
||||||
|
|
||||||
assert_ne!(id1, id2);
|
|
||||||
assert_eq!(id1.len(), 32); // Should be 32 hex chars
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn test_post_event_handler_not_found() {
|
|
||||||
let app = create_test_app();
|
|
||||||
let _router = app.router();
|
|
||||||
|
|
||||||
// Create a request with a session ID that doesn't exist
|
|
||||||
let _request = Request::builder()
|
|
||||||
.method(Method::POST)
|
|
||||||
.uri("/sse?sessionId=nonexistent")
|
|
||||||
.body(Body::empty())
|
|
||||||
.unwrap();
|
|
||||||
|
|
||||||
// Since we can't use oneshot without tower imports,
|
|
||||||
// we'll skip the actual request handling for now
|
|
||||||
|
|
||||||
// Just check that the handler would have been called
|
|
||||||
assert!(true);
|
|
||||||
}
|
|
||||||
|
|
@ -1,4 +0,0 @@
|
||||||
pub mod http_sse_server;
|
|
||||||
|
|
||||||
#[cfg(test)]
|
|
||||||
mod tests;
|
|
||||||
|
|
@ -1,16 +0,0 @@
|
||||||
use mcp_server::router::RouterService;
|
|
||||||
use mcp_server::{Server};
|
|
||||||
use crate::DocRouter;
|
|
||||||
use mcp_server::Router;
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn test_server_initialization() {
|
|
||||||
// Basic test to ensure the server initializes properly
|
|
||||||
let doc_router = DocRouter::new();
|
|
||||||
let router_name = doc_router.name();
|
|
||||||
let router = RouterService(doc_router);
|
|
||||||
let _server = Server::new(router);
|
|
||||||
|
|
||||||
// Server should be created successfully without panics
|
|
||||||
assert!(router_name.contains("rust-docs"));
|
|
||||||
}
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue