457 lines
16 KiB
Rust
457 lines
16 KiB
Rust
//! 通过已授权的本机实例进行串行 MCP 会话。 Host 批准途径仍必须建立用户同意和当前安装/信任。
|
|
use crate::{
|
|
extension_io::{Event, Pump},
|
|
extension_process::Running,
|
|
extension_stdio::HostIo,
|
|
workspace::{HostError, Result},
|
|
};
|
|
use serde::Deserialize;
|
|
use serde_json::{json, Value};
|
|
use std::{
|
|
sync::atomic::{AtomicBool, Ordering},
|
|
time::{Duration, Instant},
|
|
};
|
|
pub const PROTOCOL_VERSION: &str = "2025-11-25";
|
|
const VERSIONS: [&str; 4] = [PROTOCOL_VERSION, "2025-06-18", "2025-03-26", "2024-11-05"];
|
|
fn present<'de, D: serde::Deserializer<'de>, T: Deserialize<'de>>(
|
|
deserializer: D,
|
|
) -> std::result::Result<Option<T>, D::Error> {
|
|
T::deserialize(deserializer).map(Some)
|
|
}
|
|
#[derive(Deserialize)]
|
|
#[serde(deny_unknown_fields)]
|
|
struct Envelope {
|
|
jsonrpc: String,
|
|
#[serde(default, deserialize_with = "present")]
|
|
id: Option<Value>,
|
|
#[serde(default, deserialize_with = "present")]
|
|
method: Option<String>,
|
|
#[serde(default, deserialize_with = "present")]
|
|
params: Option<Value>,
|
|
#[serde(default, deserialize_with = "present")]
|
|
result: Option<Value>,
|
|
#[serde(default, deserialize_with = "present")]
|
|
error: Option<RemoteError>,
|
|
}
|
|
#[derive(Deserialize)]
|
|
#[serde(deny_unknown_fields)]
|
|
struct RemoteError {
|
|
code: i64,
|
|
message: String,
|
|
data: Option<Value>,
|
|
}
|
|
fn invalid() -> HostError {
|
|
HostError::new("EXTENSION_MCP_PROTOCOL_INVALID")
|
|
}
|
|
fn decode(bytes: &[u8]) -> Result<Envelope> {
|
|
let value: Envelope = serde_json::from_slice(bytes).map_err(|_| invalid())?;
|
|
if value.jsonrpc != "2.0" || value.params.as_ref().is_some_and(|v| !v.is_object()) {
|
|
return Err(invalid());
|
|
}
|
|
if let Some(id) = &value.id {
|
|
if id.as_i64().is_none() && !id.as_str().is_some_and(|s| !s.is_empty() && s.len() <= 128) {
|
|
return Err(invalid());
|
|
}
|
|
}
|
|
if let Some(method) = &value.method {
|
|
if method.is_empty()
|
|
|| method.len() > 256
|
|
|| method.chars().any(char::is_control)
|
|
|| value.result.is_some()
|
|
|| value.error.is_some()
|
|
{
|
|
return Err(invalid());
|
|
}
|
|
} else if value.id.is_none()
|
|
|| value.params.is_some()
|
|
|| value.result.is_some() == value.error.is_some()
|
|
{
|
|
return Err(invalid());
|
|
}
|
|
if value.result.as_ref().is_some_and(|v| !v.is_object()) {
|
|
return Err(invalid());
|
|
}
|
|
if let Some(error) = &value.error {
|
|
if error.message.len() > 4096
|
|
|| error.code < i32::MIN as i64
|
|
|| error.code > i32::MAX as i64
|
|
{
|
|
return Err(invalid());
|
|
}
|
|
// 错误数据可能包含秘密，因此绝不传播或记录。
|
|
let _ = &error.data;
|
|
}
|
|
Ok(value)
|
|
}
|
|
pub struct Session<'a, 'p> {
|
|
process: &'a Running<'p>,
|
|
pump: Pump,
|
|
next: u64,
|
|
ready: bool,
|
|
tools: bool,
|
|
failed: bool,
|
|
tools_changed: bool,
|
|
catalog: Option<crate::extension_mcp_tools::Catalog>,
|
|
calls: crate::extension_call_authorization::Gate,
|
|
}
|
|
impl<'a, 'p> Session<'a, 'p> {
|
|
pub fn new(process: &'a Running<'p>, io: HostIo) -> Result<Self> {
|
|
Ok(Self {
|
|
process,
|
|
pump: process.start_io(io)?,
|
|
next: 1,
|
|
ready: false,
|
|
tools: false,
|
|
failed: false,
|
|
tools_changed: false,
|
|
catalog: None,
|
|
calls: crate::extension_call_authorization::Gate::new(process.call_identity()?),
|
|
})
|
|
}
|
|
pub fn initialize(&mut self, cancel: &AtomicBool) -> Result<String> {
|
|
if self.ready {
|
|
return Err(HostError::new("EXTENSION_MCP_ALREADY_INITIALIZED"));
|
|
}
|
|
let value = self.request("initialize", json!({"protocolVersion":PROTOCOL_VERSION,"capabilities":{},"clientInfo":{"name":"OpenNexus","version":env!("CARGO_PKG_VERSION")}}), Duration::from_secs(10), cancel)?;
|
|
let version = value["protocolVersion"]
|
|
.as_str()
|
|
.filter(|v| VERSIONS.contains(v));
|
|
let info = value["serverInfo"].as_object();
|
|
let capabilities = value["capabilities"].as_object();
|
|
if version.is_none()
|
|
|| info.is_none()
|
|
|| capabilities.is_none()
|
|
|| !value["serverInfo"]["name"]
|
|
.as_str()
|
|
.is_some_and(|s| !s.is_empty() && s.len() <= 256)
|
|
|| !value["serverInfo"]["version"]
|
|
.as_str()
|
|
.is_some_and(|s| !s.is_empty() && s.len() <= 128)
|
|
|| value["capabilities"]
|
|
.get("tools")
|
|
.is_some_and(|v| !v.is_object())
|
|
{
|
|
self.abort();
|
|
return Err(HostError::new("EXTENSION_MCP_INITIALIZATION_INVALID"));
|
|
}
|
|
self.tools = value["capabilities"].get("tools").is_some();
|
|
if let Err(error) = self.send(json!({"jsonrpc":"2.0","method":"notifications/initialized"}))
|
|
{
|
|
self.abort();
|
|
return Err(error);
|
|
}
|
|
self.ready = true;
|
|
Ok(version.unwrap().into())
|
|
}
|
|
pub fn refresh_tools(
|
|
&mut self,
|
|
cancel: &AtomicBool,
|
|
) -> Result<Vec<crate::extension_mcp_tools::Description>> {
|
|
self.require_tools()?;
|
|
self.drain_pending()?;
|
|
self.catalog = None;
|
|
self.calls.invalidate();
|
|
self.tools_changed = false;
|
|
let started = Instant::now();
|
|
let catalog = crate::extension_mcp_tools::Catalog::discover(|cursor| {
|
|
let remaining = Duration::from_secs(20).saturating_sub(started.elapsed());
|
|
if remaining.is_zero() {
|
|
return Err(HostError::new("EXTENSION_MCP_CATALOG_TIMEOUT"));
|
|
}
|
|
self.list_tools_page(cursor, cancel, remaining.min(Duration::from_secs(10)))
|
|
});
|
|
let catalog = match catalog {
|
|
Ok(catalog) => catalog,
|
|
Err(error) => {
|
|
self.abort();
|
|
return Err(error);
|
|
}
|
|
};
|
|
self.drain_pending()?;
|
|
if self.tools_changed {
|
|
return Err(HostError::new("EXTENSION_MCP_CATALOG_CHANGED"));
|
|
}
|
|
let descriptions = catalog.descriptions();
|
|
self.catalog = Some(catalog);
|
|
Ok(descriptions)
|
|
}
|
|
fn list_tools_page(
|
|
&mut self,
|
|
cursor: Option<&str>,
|
|
cancel: &AtomicBool,
|
|
budget: Duration,
|
|
) -> Result<Value> {
|
|
self.require_tools()?;
|
|
if cursor.is_some_and(|s| s.is_empty() || s.len() > 1024) {
|
|
return Err(invalid());
|
|
}
|
|
let value = self.request(
|
|
"tools/list",
|
|
cursor.map_or_else(|| json!({}), |c| json!({"cursor":c})),
|
|
budget,
|
|
cancel,
|
|
)?;
|
|
if !value["tools"]
|
|
.as_array()
|
|
.is_some_and(|tools| tools.len() <= 500)
|
|
|| value
|
|
.get("nextCursor")
|
|
.is_some_and(|v| !v.as_str().is_some_and(|s| !s.is_empty() && s.len() <= 1024))
|
|
{
|
|
self.abort();
|
|
return Err(invalid());
|
|
}
|
|
Ok(value)
|
|
}
|
|
pub fn review_call(
|
|
&mut self,
|
|
name: &str,
|
|
arguments: Value,
|
|
) -> Result<crate::extension_call_authorization::Review> {
|
|
self.require_tools()?;
|
|
self.drain_pending()?;
|
|
let tool = self
|
|
.catalog
|
|
.as_ref()
|
|
.ok_or_else(|| HostError::new("EXTENSION_MCP_CATALOG_REQUIRED"))?
|
|
.tool(name)?;
|
|
self.calls.review(&tool, arguments)
|
|
}
|
|
/// 仅在用户批准完全一致的审查后由已认证 Host 路由调用。
|
|
/// 此方法不会注册为 renderer/Core 命令。
|
|
pub fn confirm_call(
|
|
&mut self,
|
|
review_id: &str,
|
|
) -> Result<crate::extension_call_authorization::ApprovedCall> {
|
|
self.require_tools()?;
|
|
self.drain_pending()?;
|
|
self.calls.confirm(review_id)
|
|
}
|
|
pub fn call_tool(
|
|
&mut self,
|
|
approved: crate::extension_call_authorization::ApprovedCall,
|
|
cancel: &AtomicBool,
|
|
) -> Result<Value> {
|
|
self.require_tools()?;
|
|
self.drain_pending()?;
|
|
let call = self.calls.consume(approved)?;
|
|
let tool = self
|
|
.catalog
|
|
.as_ref()
|
|
.ok_or_else(|| HostError::new("EXTENSION_MCP_CATALOG_REQUIRED"))?
|
|
.tool(&call.name)?;
|
|
if crate::extension_call_authorization::contract_digest(&tool.description())?
|
|
!= call.contract_digest
|
|
{
|
|
return Err(HostError::new("EXTENSION_CALL_CATALOG_CHANGED"));
|
|
}
|
|
self.invoke_tool(&call.name, call.arguments, cancel)
|
|
}
|
|
#[cfg(test)]
|
|
pub(crate) fn test_call_tool(
|
|
&mut self,
|
|
name: &str,
|
|
arguments: Value,
|
|
cancel: &AtomicBool,
|
|
) -> Result<Value> {
|
|
let review = self.review_call(name, arguments)?;
|
|
let approved = self.confirm_call(&review.review_id)?;
|
|
self.call_tool(approved, cancel)
|
|
}
|
|
fn invoke_tool(&mut self, name: &str, arguments: Value, cancel: &AtomicBool) -> Result<Value> {
|
|
self.require_tools()?;
|
|
self.drain_pending()?;
|
|
if name.is_empty()
|
|
|| name.len() > 256
|
|
|| name.chars().any(char::is_control)
|
|
|| !arguments.is_object()
|
|
{
|
|
return Err(invalid());
|
|
}
|
|
let tool = self
|
|
.catalog
|
|
.as_ref()
|
|
.ok_or_else(|| HostError::new("EXTENSION_MCP_CATALOG_REQUIRED"))?
|
|
.tool(name)?;
|
|
tool.validate_arguments(&arguments)?;
|
|
let value = self.request(
|
|
"tools/call",
|
|
json!({"name":name,"arguments":arguments}),
|
|
Duration::from_secs(60),
|
|
cancel,
|
|
)?;
|
|
if let Err(error) = tool.validate_result(&value) {
|
|
self.abort();
|
|
return Err(error);
|
|
}
|
|
Ok(value)
|
|
}
|
|
fn require_tools(&self) -> Result<()> {
|
|
if self.failed {
|
|
return Err(HostError::new("EXTENSION_MCP_SESSION_FAILED"));
|
|
}
|
|
if !self.ready {
|
|
return Err(HostError::new("EXTENSION_MCP_NOT_INITIALIZED"));
|
|
}
|
|
if !self.tools {
|
|
return Err(HostError::new("EXTENSION_MCP_TOOLS_UNAVAILABLE"));
|
|
}
|
|
Ok(())
|
|
}
|
|
fn server_message(&mut self, message: &Envelope) -> Result<bool> {
|
|
let Some(method) = &message.method else {
|
|
return Ok(false);
|
|
};
|
|
if let Some(id) = &message.id {
|
|
self.send(if method == "ping" { json!({"jsonrpc":"2.0","id":id,"result":{}}) }
|
|
else { json!({"jsonrpc":"2.0","id":id,"error":{"code":-32601,"message":"Method not supported"}}) })?;
|
|
} else if method == "notifications/tools/list_changed" {
|
|
self.tools_changed = true;
|
|
self.catalog = None;
|
|
self.calls.invalidate();
|
|
}
|
|
Ok(true)
|
|
}
|
|
/// 选择缓存契约前应用已收到的通知；最终注册表循环在实例空闲时也必须调用此方法。
|
|
pub fn drain_pending(&mut self) -> Result<()> {
|
|
let result = (|| {
|
|
for _ in 0..128 {
|
|
self.process.check_authorization()?;
|
|
match self.pump.receive(Duration::ZERO) {
|
|
Err(error) if error.code == "EXTENSION_IO_TIMEOUT" => return Ok(()),
|
|
Err(error) => return Err(error),
|
|
Ok(Event::Closed) => {
|
|
return Err(HostError::new("EXTENSION_MCP_CONNECTION_CLOSED"))
|
|
}
|
|
Ok(Event::Frame(bytes)) => {
|
|
let message = decode(&bytes)?;
|
|
if !self.server_message(&message)? {
|
|
return Err(HostError::new("EXTENSION_MCP_UNEXPECTED_RESPONSE"));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
Err(HostError::new("EXTENSION_IO_RATE_LIMITED"))
|
|
})();
|
|
if result.is_err() {
|
|
self.abort();
|
|
}
|
|
result
|
|
}
|
|
fn send(&self, value: Value) -> Result<()> {
|
|
self.pump.send_wait(
|
|
serde_json::to_vec(&value).map_err(|_| invalid())?,
|
|
Duration::from_millis(50),
|
|
)
|
|
}
|
|
fn abort(&mut self) {
|
|
self.failed = true;
|
|
self.calls.invalidate();
|
|
let _ = self.process.terminate();
|
|
}
|
|
fn request(
|
|
&mut self,
|
|
method: &str,
|
|
params: Value,
|
|
budget: Duration,
|
|
cancel: &AtomicBool,
|
|
) -> Result<Value> {
|
|
if self.failed {
|
|
return Err(HostError::new("EXTENSION_MCP_SESSION_FAILED"));
|
|
}
|
|
if cancel.load(Ordering::Acquire) {
|
|
return Err(HostError::new("EXTENSION_MCP_CANCELLED"));
|
|
}
|
|
let id = format!("opennexus.{}", self.next);
|
|
self.next = self.next.checked_add(1).ok_or_else(invalid)?;
|
|
let deadline = self.process.start_tool_call()?;
|
|
let started = Instant::now();
|
|
let result = (|| {
|
|
self.send(json!({"jsonrpc":"2.0","id":id,"method":method,"params":params}))?;
|
|
loop {
|
|
self.process.check_authorization()?;
|
|
deadline.check()?;
|
|
if cancel.load(Ordering::Acquire) || started.elapsed() >= budget {
|
|
// initialize 无法在协议层取消。
|
|
if method != "initialize" {
|
|
let _ = self.send(json!({"jsonrpc":"2.0","method":"notifications/cancelled","params":{"requestId":id}}));
|
|
}
|
|
return Err(HostError::new(if cancel.load(Ordering::Acquire) {
|
|
"EXTENSION_MCP_CANCELLED"
|
|
} else {
|
|
"EXTENSION_MCP_TIMEOUT"
|
|
}));
|
|
}
|
|
let received = self.pump.receive(Duration::from_millis(20));
|
|
self.process.check_authorization()?;
|
|
deadline.check()?;
|
|
let bytes = match received {
|
|
Ok(Event::Frame(bytes)) => bytes,
|
|
Ok(Event::Closed) => {
|
|
return Err(HostError::new("EXTENSION_MCP_CONNECTION_CLOSED"))
|
|
}
|
|
Err(error) if error.code == "EXTENSION_IO_TIMEOUT" => continue,
|
|
Err(error) => return Err(error),
|
|
};
|
|
let message = decode(&bytes)?;
|
|
if self.server_message(&message)? {
|
|
continue;
|
|
}
|
|
if message.id != Some(json!(id)) {
|
|
return Err(HostError::new("EXTENSION_MCP_RESPONSE_ID_MISMATCH"));
|
|
}
|
|
return Ok(message.result);
|
|
}
|
|
})();
|
|
match result {
|
|
Ok(value) => {
|
|
if let Err(error) = deadline.finish() {
|
|
self.abort();
|
|
return Err(error);
|
|
}
|
|
value.ok_or_else(|| HostError::new("EXTENSION_MCP_REMOTE_ERROR"))
|
|
}
|
|
Err(error) => {
|
|
self.abort();
|
|
drop(deadline);
|
|
Err(error)
|
|
}
|
|
}
|
|
}
|
|
pub(crate) fn is_failed(&self) -> bool {
|
|
self.failed
|
|
}
|
|
pub fn take_tools_changed(&mut self) -> bool {
|
|
std::mem::take(&mut self.tools_changed)
|
|
}
|
|
pub fn shutdown(mut self) -> Result<()> {
|
|
self.pump.close_input();
|
|
let _ = self.process.wait(Duration::from_millis(200));
|
|
self.pump.shutdown()
|
|
}
|
|
}
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
#[test]
|
|
fn malformed_and_ambiguous_envelopes_are_rejected() {
|
|
for input in [
|
|
r#"[]"#,
|
|
r#"{"jsonrpc":"2.0","id":null,"method":"ping"}"#,
|
|
r#"{"jsonrpc":"2.0","method":"ping","params":null}"#,
|
|
r#"{"jsonrpc":"2.0","id":1,"id":2,"result":{}}"#,
|
|
r#"{"jsonrpc":"1.0","id":1,"result":{}}"#,
|
|
r#"{"jsonrpc":"2.0","id":null,"result":{}}"#,
|
|
r#"{"jsonrpc":"2.0","id":1,"result":{},"error":{"code":1,"message":"bad"}}"#,
|
|
r#"{"jsonrpc":"2.0","id":1,"method":"ping","result":{}}"#,
|
|
] {
|
|
assert!(decode(input.as_bytes()).is_err(), "{input}");
|
|
}
|
|
assert!(decode(br#"{"jsonrpc":"2.0","id":"request","result":{}}"#).is_ok());
|
|
assert!(decode(br#"{"jsonrpc":"2.0","id":1,"method":"ping"}"#).is_ok());
|
|
assert!(
|
|
decode(br#"{"jsonrpc":"2.0","method":"notifications/tools/list_changed"}"#).is_ok()
|
|
);
|
|
}
|
|
}
|