Skip to content

Commit f0e350a

Browse files
committed
add message interceptor return action
1 parent 1fd7d49 commit f0e350a

File tree

4 files changed

+47
-5
lines changed

4 files changed

+47
-5
lines changed

mcp-guardian-core/src/message_interceptor.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ use crate::message::{
1616
pub enum MessageInterceptorAction {
1717
Send(Message),
1818
Drop,
19+
Return(Message),
1920
}
2021

2122
#[async_trait]

mcp-guardian-core/src/message_interceptor/chain.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ use crate::{
77
message::{Message, MessageDirection},
88
message_interceptor::{
99
MessageInterceptor, MessageInterceptorAction,
10-
MessageInterceptorAction::{Drop, Send},
10+
MessageInterceptorAction::{Drop, Return, Send},
1111
},
1212
};
1313
pub struct ChainInterceptor {
@@ -35,6 +35,7 @@ impl MessageInterceptor for ChainInterceptor {
3535
match interceptor.intercept_message(direction, message).await? {
3636
Send(new_message) => message = new_message,
3737
Drop => return Ok(Drop),
38+
Return(return_message) => return Ok(Return(return_message)),
3839
}
3940
}
4041

mcp-guardian-core/src/message_interceptor/manual_approval.rs

Lines changed: 27 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,12 @@
11
use anyhow::Result;
22
use async_trait::async_trait;
3+
use serde_json::json;
34
use tokio::time::{sleep, Duration};
45
use uuid::Uuid;
5-
use MessageInterceptorAction::{Drop, Send};
6+
use MessageInterceptorAction::{Return, Send};
67

78
use crate::{
8-
message::{Message, MessageDirection},
9+
message::{Message, MessageDirection, MessageType},
910
message_approval::{request_approval, MessageStatus},
1011
message_interceptor::{MessageInterceptor, MessageInterceptorAction},
1112
};
@@ -43,7 +44,30 @@ impl MessageInterceptor for ManualApprovalInterceptor {
4344
return Ok(Send(message));
4445
}
4546
MessageStatus::Denied | MessageStatus::Unknown => {
46-
return Ok(Drop);
47+
let id = message
48+
.raw_msg
49+
.get("id")
50+
.ok_or_else(|| anyhow::anyhow!("Request message did not contain an ID"))?
51+
.to_owned();
52+
53+
let return_message = Message {
54+
type_: MessageType::ResponseSuccess,
55+
raw_msg: json!({
56+
"id": id,
57+
"jsonrpc": "2.0",
58+
"result": {
59+
"content": [
60+
{
61+
"text": "Acess approval was denied.",
62+
"type": "text"
63+
}
64+
],
65+
"isError": false
66+
}
67+
}),
68+
};
69+
70+
return Ok(Return(return_message));
4771
}
4872
}
4973
}

mcp-guardian-core/src/proxy.rs

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ use crate::{
1313
message::Message,
1414
message_interceptor::{
1515
MessageInterceptor,
16-
MessageInterceptorAction::{Drop, Send},
16+
MessageInterceptorAction::{Drop, Return, Send},
1717
},
1818
};
1919

@@ -128,6 +128,14 @@ pub async fn proxy_mcp_server(
128128
}
129129
}
130130
Ok(Drop) => {}
131+
Ok(Return(message)) => {
132+
if let Err(e) = writeln!(io::stdout(), "{}", message.raw_msg) {
133+
log::error!("Failed to write to stdout: {e}");
134+
}
135+
if let Err(e) = io::stdout().flush() {
136+
log::error!("Failed to flush stdout: {e}");
137+
}
138+
}
131139
Err(e) => {
132140
log::error!("Failed to intercept outbound message properly: {e}");
133141
}
@@ -161,6 +169,14 @@ pub async fn proxy_mcp_server(
161169
}
162170
}
163171
Ok(Drop) => {}
172+
Ok(Return(message)) => {
173+
if let Err(e) = writeln!(io::stdout(), "{}", message.raw_msg) {
174+
log::error!("Failed to write to stdout: {e}");
175+
}
176+
if let Err(e) = io::stdout().flush() {
177+
log::error!("Failed to flush stdout: {e}");
178+
}
179+
}
164180
Err(e) => {
165181
log::error!("Failed to intercept outbound message properly: {e}");
166182
}

0 commit comments

Comments
 (0)