aljaz/goose
1
0
Fork 0
mirror of https://github.com/aljazceru/goose.git synced 2025-12-18 06:34:26 +01:00
Code Issues Packages Projects Releases Wiki Activity

feat: Handle MCP server notification messages (#2613)

Browse Source 
Co-authored-by: Michael Neale <michael.neale@gmail.com>
This commit is contained in:
Jack Amadeo
2025-05-30 11:50:14 -04:00
committed by GitHub
parent eeb61ace22
commit 03e5549b54
40 changed files with 1186 additions and 443 deletions
Download Patch File Download Diff File Expand all files Collapse all files

45
crates/mcp-server/src/lib.rs
View File

@@ -4,9 +4,13 @@ use std::{
};
use futures::{Future, Stream};
use mcp_core::protocol::{JsonRpcError, JsonRpcMessage, JsonRpcRequest, JsonRpcResponse};
use mcp_core::protocol::{JsonRpcError, JsonRpcMessage, JsonRpcResponse};
use pin_project::pin_project;
use tokio::io::{AsyncBufReadExt, AsyncRead, AsyncWrite, AsyncWriteExt, BufReader};
use router::McpRequest;
use tokio::{
io::{AsyncBufReadExt, AsyncRead, AsyncWrite, AsyncWriteExt, BufReader},
sync::mpsc,
};
use tower_service::Service;
mod errors;
@@ -123,7 +127,7 @@ pub struct Server<S> {
impl<S> Server<S>
where
S: Service<JsonRpcRequest, Response = JsonRpcResponse> + Send,
S: Service<McpRequest, Response = JsonRpcResponse> + Send,
S::Error: Into<BoxError>,
S::Future: Send,
{
@@ -134,8 +138,8 @@ where
// TODO transport trait instead of byte transport if we implement others
pub async fn run<R, W>(self, mut transport: ByteTransport<R, W>) -> Result<(), ServerError>
where
R: AsyncRead + Unpin,
W: AsyncWrite + Unpin,
R: AsyncRead + Unpin + Send + 'static,
W: AsyncWrite + Unpin + Send + 'static,
{
use futures::StreamExt;
let mut service = self.service;
@@ -160,7 +164,22 @@ where
);
// Process the request using our service
let response = match service.call(request).await {
let (notify_tx, mut notify_rx) = mpsc::channel(256);
let mcp_request = McpRequest {
request,
notifier: notify_tx,
};
let transport_fut = tokio::spawn(async move {
while let Some(notification) = notify_rx.recv().await {
if transport.write_message(notification).await.is_err() {
break;
}
}
transport
});
let response = match service.call(mcp_request).await {
Ok(resp) => resp,
Err(e) => {
let error_msg = e.into().to_string();
@@ -178,6 +197,16 @@ where
}
};
transport = match transport_fut.await {
Ok(transport) => transport,
Err(e) => {
tracing::error!(error = %e, "Failed to spawn transport task");
return Err(ServerError::Transport(TransportError::Io(
e.into(),
)));
}
};
// Serialize response for logging
let response_json = serde_json::to_string(&response)
.unwrap_or_else(|_| "Failed to serialize response".to_string());
@@ -247,7 +276,7 @@ where
// Any router implements this
pub trait BoundedService:
Service<
JsonRpcRequest,
McpRequest,
Response = JsonRpcResponse,
Error = BoxError,
Future = Pin<Box<dyn Future<Output = Result<JsonRpcResponse, BoxError>> + Send>>,
@@ -259,7 +288,7 @@ pub trait BoundedService:
// Implement it for any type that meets the bounds
impl<T> BoundedService for T where
T: Service<
JsonRpcRequest,
McpRequest,
Response = JsonRpcResponse,
Error = BoxError,
Future = Pin<Box<dyn Future<Output = Result<JsonRpcResponse, BoxError>> + Send>>,

3
crates/mcp-server/src/main.rs
View File

@@ -2,12 +2,14 @@ use anyhow::Result;
use mcp_core::content::Content;
use mcp_core::handler::{PromptError, ResourceError};
use mcp_core::prompt::{Prompt, PromptArgument};
use mcp_core::protocol::JsonRpcMessage;
use mcp_core::tool::ToolAnnotations;
use mcp_core::{handler::ToolError, protocol::ServerCapabilities, resource::Resource, tool::Tool};
use mcp_server::router::{CapabilitiesBuilder, RouterService};
use mcp_server::{ByteTransport, Router, Server};
use serde_json::Value;
use std::{future::Future, pin::Pin, sync::Arc};
use tokio::sync::mpsc;
use tokio::{
io::{stdin, stdout},
sync::Mutex,
@@ -124,6 +126,7 @@ impl Router for CounterRouter {
&self,
tool_name: &str,
_arguments: Value,
_notifier: mpsc::Sender<JsonRpcMessage>,
) -> Pin<Box<dyn Future<Output = Result<Vec<Content>, ToolError>> + Send + 'static>> {
let this = self.clone();
let tool_name = tool_name.to_string();

38
crates/mcp-server/src/router.rs
View File

@@ -11,14 +11,15 @@ use mcp_core::{
handler::{PromptError, ResourceError, ToolError},
prompt::{Prompt, PromptMessage, PromptMessageRole},
protocol::{
CallToolResult, GetPromptResult, Implementation, InitializeResult, JsonRpcRequest,
JsonRpcResponse, ListPromptsResult, ListResourcesResult, ListToolsResult,
CallToolResult, GetPromptResult, Implementation, InitializeResult, JsonRpcMessage,
JsonRpcRequest, JsonRpcResponse, ListPromptsResult, ListResourcesResult, ListToolsResult,
PromptsCapability, ReadResourceResult, ResourcesCapability, ServerCapabilities,
ToolsCapability,
},
ResourceContents,
};
use serde_json::Value;
use tokio::sync::mpsc;
use tower_service::Service;
use crate::{BoxError, RouterError};
@@ -91,6 +92,7 @@ pub trait Router: Send + Sync + 'static {
&self,
tool_name: &str,
arguments: Value,
notifier: mpsc::Sender<JsonRpcMessage>,
) -> Pin<Box<dyn Future<Output = Result<Vec<Content>, ToolError>> + Send + 'static>>;
fn list_resources(&self) -> Vec<mcp_core::resource::Resource>;
fn read_resource(
@@ -159,6 +161,7 @@ pub trait Router: Send + Sync + 'static {
fn handle_tools_call(
&self,
req: JsonRpcRequest,
notifier: mpsc::Sender<JsonRpcMessage>,
) -> impl Future<Output = Result<JsonRpcResponse, RouterError>> + Send {
async move {
let params = req
@@ -172,7 +175,7 @@ pub trait Router: Send + Sync + 'static {
let arguments = params.get("arguments").cloned().unwrap_or(Value::Null);
let result = match self.call_tool(name, arguments).await {
let result = match self.call_tool(name, arguments, notifier).await {
Ok(result) => CallToolResult {
content: result,
is_error: None,
@@ -394,7 +397,12 @@ pub trait Router: Send + Sync + 'static {
pub struct RouterService<T>(pub T);
impl<T> Service<JsonRpcRequest> for RouterService<T>
pub struct McpRequest {
pub request: JsonRpcRequest,
pub notifier: mpsc::Sender<JsonRpcMessage>,
}
impl<T> Service<McpRequest> for RouterService<T>
where
T: Router + Clone + Send + Sync + 'static,
{
@@ -406,21 +414,21 @@ where
Poll::Ready(Ok(()))
}
fn call(&mut self, req: JsonRpcRequest) -> Self::Future {
fn call(&mut self, req: McpRequest) -> Self::Future {
let this = self.0.clone();
Box::pin(async move {
let result = match req.method.as_str() {
"initialize" => this.handle_initialize(req).await,
"tools/list" => this.handle_tools_list(req).await,
"tools/call" => this.handle_tools_call(req).await,
"resources/list" => this.handle_resources_list(req).await,
"resources/read" => this.handle_resources_read(req).await,
"prompts/list" => this.handle_prompts_list(req).await,
"prompts/get" => this.handle_prompts_get(req).await,
let result = match req.request.method.as_str() {
"initialize" => this.handle_initialize(req.request).await,
"tools/list" => this.handle_tools_list(req.request).await,
"tools/call" => this.handle_tools_call(req.request, req.notifier).await,
"resources/list" => this.handle_resources_list(req.request).await,
"resources/read" => this.handle_resources_read(req.request).await,
"prompts/list" => this.handle_prompts_list(req.request).await,
"prompts/get" => this.handle_prompts_get(req.request).await,
_ => {
let mut response = this.create_response(req.id);
response.error = Some(RouterError::MethodNotFound(req.method).into());
let mut response = this.create_response(req.request.id);
response.error = Some(RouterError::MethodNotFound(req.request.method).into());
Ok(response)
}
};