Ything-RepoMirrors/poem
1
0
Fork 0
mirror of https://github.com/poem-web/poem.git synced 2026-01-25 04:18:25 +00:00
Code Issues Packages Projects Releases Wiki Activity

poem-mcp: add support for output schema in tools.

Browse Source
This commit is contained in:
Sunli
2025-10-10 14:49:11 +08:00
parent 3cf007a197
commit 91c9ce36de
11 changed files with 163 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
Cargo.toml
View File

@@ -27,7 +27,7 @@ poem = { path = "poem", version = "3.1.12", default-features = false }
poem-derive = { path = "poem-derive", version = "3.1.12" }
poem-openapi-derive = { path = "poem-openapi-derive", version = "5.1.15" }
poem-grpc-build = { path = "poem-grpc-build", version = "0.5.6" }
poem-mcpserver-macros = { path = "poem-mcpserver-macros", version = "0.2.8" }
poem-mcpserver-macros = { path = "poem-mcpserver-macros", version = "0.2.9" }
proc-macro-crate = "3.0.0"
proc-macro2 = "1.0.29"

2
poem-mcpserver-macros/Cargo.toml
View File

@@ -1,6 +1,6 @@
[package]
name = "poem-mcpserver-macros"
version = "0.2.8"
version = "0.2.9"
authors.workspace = true
edition.workspace = true
license.workspace = true

11
poem-mcpserver-macros/src/tools.rs
View File

@@ -83,14 +83,23 @@ pub(crate) fn generate(_args: ToolsArgs, mut item_impl: ItemImpl) -> Result<Toke
});
}
let resp_ty = match &method.sig.output {
syn::ReturnType::Default => {
return Err(Error::custom("missing return type").with_span(&method.sig.ident));
}
syn::ReturnType::Type(_, ty) => ty,
};
tools_descriptions.push(quote! {
#crate_name::protocol::tool::Tool {
name: #tool_name,
description: #tool_description,
input_schema: {
let schema = schemars::SchemaGenerator::default().into_root_schema_for::<#request_type>();
#crate_name::private::serde_json::to_value(schema).expect("serialize schema")
#crate_name::private::serde_json::to_value(schema).expect("serialize input schema")
},
output_schema: std::option::Option::map(<#resp_ty as #crate_name::tool::IntoToolResponse>::output_schema(), |schema| {
#crate_name::private::serde_json::to_value(schema).expect("serialize output schema")
}),
},
});

3
poem-mcpserver/CHANGELOG.md
View File

@@ -4,9 +4,10 @@ All notable changes to this project will be documented in this file.
The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
# [Unreleased]
# [0.2.9]
- bump `schemars` to 1.0
- add support for output schema in tools.
# [0.2.4] 20250-06-06

2
poem-mcpserver/Cargo.toml
View File

@@ -1,6 +1,6 @@
[package]
name = "poem-mcpserver"
version = "0.2.8"
version = "0.2.9"
authors.workspace = true
edition.workspace = true
license.workspace = true

16
poem-mcpserver/src/content.rs
View File

@@ -3,7 +3,6 @@
use std::fmt::Display;
use base64::{Engine, engine::general_purpose::STANDARD};
use serde::Serialize;
use crate::protocol::content::Content;
@@ -101,18 +100,3 @@ where
}
}
}
/// A Json response.
#[derive(Debug, Clone, Copy)]
pub struct Json<T>(pub T);
impl<T> IntoContent for Json<T>
where
T: Serialize,
{
fn into_content(self) -> Content {
Content::Text {
text: serde_json::to_string(&self.0).unwrap_or_default(),
}
}
}

1
poem-mcpserver/src/lib.rs
View File

@@ -15,6 +15,7 @@ pub mod stdio;
pub mod streamable_http;
pub mod tool;
pub use poem_mcpserver_macros::Tools;
pub use schemars::JsonSchema;
pub use server::McpServer;
#[doc(hidden)]

2
poem-mcpserver/src/protocol/mod.rs
View File

@@ -12,4 +12,4 @@ pub mod tool;
pub const JSON_RPC_VERSION: &str = "2.0";
/// The MCP protocol version.
pub const MCP_PROTOCOL_VERSION: time::Date = time::macros::date!(2025 - 03 - 26);
pub const MCP_PROTOCOL_VERSION: time::Date = time::macros::date!(2025 - 06 - 18);

6
poem-mcpserver/src/protocol/tool.rs
View File

@@ -23,6 +23,9 @@ pub struct Tool {
pub description: &'static str,
/// The input schema of the tool.
pub input_schema: Value,
/// The output schema of the tool, if any.
#[serde(skip_serializing_if = "Option::is_none")]
pub output_schema: Option<Value>,
}
/// A response to a tools/list request.
@@ -50,6 +53,9 @@ pub struct ToolsCallRequest {
pub struct ToolsCallResponse {
/// Response content.
pub content: Vec<Content>,
/// Structured content (if any).
#[serde(skip_serializing_if = "Option::is_none")]
pub structured_content: Option<Value>,
/// Whether the response is an error.
pub is_error: bool,
}

70
poem-mcpserver/src/tool.rs
View File

@@ -2,6 +2,8 @@
use std::{fmt::Display, future::Future};
use schemars::{JsonSchema, Schema};
use serde::Serialize;
use serde_json::Value;
use crate::{
@@ -15,6 +17,9 @@ use crate::{
/// Represents the result of a tool call.
pub trait IntoToolResponse {
/// Returns the output schema of the tool response, if any.
fn output_schema() -> Option<Schema>;
/// Consumes the object and converts it into a tool response.
fn into_tool_response(self) -> ToolsCallResponse;
}
@@ -23,9 +28,14 @@ impl<T> IntoToolResponse for T
where
T: IntoContents,
{
fn output_schema() -> Option<Schema> {
None
}
fn into_tool_response(self) -> ToolsCallResponse {
ToolsCallResponse {
content: self.into_contents(),
structured_content: None,
is_error: false,
}
}
@@ -36,22 +46,82 @@ where
T: IntoContents,
E: Display,
{
fn output_schema() -> Option<Schema> {
None
}
fn into_tool_response(self) -> ToolsCallResponse {
match self {
Ok(value) => ToolsCallResponse {
content: value.into_contents(),
structured_content: None,
is_error: false,
},
Err(error) => ToolsCallResponse {
content: vec![Content::Text {
text: error.to_string(),
}],
structured_content: None,
is_error: true,
},
}
}
}
/// A Structured content.
#[derive(Debug, Clone, Copy)]
pub struct StructuredContent<T>(pub T);
impl<T> IntoToolResponse for StructuredContent<T>
where
T: Serialize + JsonSchema,
{
fn output_schema() -> Option<Schema> {
Some(schemars::SchemaGenerator::default().into_root_schema_for::<T>())
}
fn into_tool_response(self) -> ToolsCallResponse {
ToolsCallResponse {
content: vec![Content::Text {
text: serde_json::to_string(&self.0).unwrap_or_default(),
}],
structured_content: Some(serde_json::to_value(&self.0).unwrap_or_default()),
is_error: false,
}
}
}
impl<T, E> IntoToolResponse for Result<StructuredContent<T>, E>
where
T: Serialize + JsonSchema,
E: Display,
{
fn output_schema() -> Option<Schema> {
Some(schemars::SchemaGenerator::default().into_root_schema_for::<T>())
}
fn into_tool_response(self) -> ToolsCallResponse {
match self {
Ok(value) => ToolsCallResponse {
content: vec![Content::Text {
text: serde_json::to_string(&value.0).unwrap_or_default(),
}],
structured_content: Some(serde_json::to_value(&value.0).unwrap_or_default()),
is_error: false,
},
Err(error) => ToolsCallResponse {
content: vec![Content::Text {
text: error.to_string(),
}],
structured_content: None,
is_error: true,
},
}
}
}
// impl IntoToolResponse for Json
/// Represents a tools collection.
pub trait Tools {
/// Returns the instructions for the tools.

71
poem-mcpserver/tests/tools.rs
View File

@@ -6,7 +6,10 @@ use poem_mcpserver::{
rpc::{Request, RequestId, Requests},
tool::{ToolsCallRequest, ToolsListRequest},
},
tool::StructuredContent,
};
use schemars::JsonSchema;
use serde::Serialize;
struct TestTools {
value: i32,
@@ -18,6 +21,12 @@ impl TestTools {
}
}
#[derive(Debug, JsonSchema, Serialize)]
struct CustomResponse {
a: i32,
b: String,
}
#[Tools]
impl TestTools {
/// Add a value to the current value.
@@ -30,6 +39,14 @@ impl TestTools {
async fn get_value(&self) -> Text<i32> {
Text(self.value)
}
/// Get a structured response.
async fn structured_response(&self) -> StructuredContent<CustomResponse> {
StructuredContent(CustomResponse {
a: self.value,
b: format!("Value is {}", self.value),
})
}
}
#[tokio::test]
@@ -112,6 +129,31 @@ async fn call_tool() {
},
})
);
let resp = server
.handle_request(Request {
jsonrpc: JSON_RPC_VERSION.to_string(),
id: Some(RequestId::Int(9)),
body: Requests::ToolsCall {
params: ToolsCallRequest {
name: "structured_response".to_string(),
arguments: serde_json::json!({}),
},
},
})
.await;
assert_eq!(
serde_json::to_value(&resp).unwrap(),
serde_json::json!({
"jsonrpc": "2.0",
"id": 9,
"result": {
"content": [{"type": "text", "text": "{\"a\":40,\"b\":\"Value is 40\"}"}],
"structuredContent": {"a": 40, "b": "Value is 40"},
"isError": false,
},
})
);
}
#[tokio::test]
@@ -161,6 +203,31 @@ async fn tool_list() {
"title": "get_value_Request",
},
},
{
"name": "structured_response",
"description": "Get a structured response.",
"inputSchema": {
"$schema": "https://json-schema.org/draft/2020-12/schema",
"type": "object",
"properties": {},
"title": "structured_response_Request",
},
"outputSchema": {
"$schema": "https://json-schema.org/draft/2020-12/schema",
"type": "object",
"properties": {
"a": {
"format": "int32",
"type": "integer",
},
"b": {
"type": "string",
},
},
"required": ["a", "b"],
"title": "CustomResponse",
},
},
],
},
})
@@ -170,7 +237,9 @@ async fn tool_list() {
#[tokio::test]
async fn disable_tools() {
let tools = TestTools::new();
let mut server = McpServer::new().tools(tools).disable_tools(["get_value"]);
let mut server = McpServer::new()
.tools(tools)
.disable_tools(["get_value", "structured_response"]);
let resp = server
.handle_request(Request {