tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

feat: upgrade rmcp (#3738)

Browse Source
This commit is contained in:
Andrew Harvard
2025-07-31 08:57:51 -04:00
committed by GitHub
parent 86871ad855
commit 001bcebcee
6 changed files with 392 additions and 281 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-server/src/openapi.rs
+219 -155
View File
@@ -12,12 +12,11 @@ use goose::providers::base::{ConfigKey, ModelInfo, ProviderMetadata};
use goose::session::info::SessionInfo;
use goose::session::SessionMetadata;
use rmcp::model::{
Annotations, Content, EmbeddedResource, ImageContent, ResourceContents, Role, TextContent,
Tool, ToolAnnotations,
Annotations, Content, EmbeddedResource, ImageContent, RawEmbeddedResource, RawImageContent,
RawTextContent, ResourceContents, Role, TextContent, Tool, ToolAnnotations,
};
use utoipa::{OpenApi, ToSchema};
use rmcp::schemars::schema::{InstanceType, SchemaObject, SingleOrVec};
use utoipa::openapi::schema::{
AdditionalProperties, AnyOfBuilder, ArrayBuilder, ObjectBuilder, OneOfBuilder, Schema,
SchemaFormat, SchemaType,
@@ -30,7 +29,7 @@ macro_rules! derive_utoipa {
impl<'__s> ToSchema<'__s> for $schema_name {
fn schema() -> (&'__s str, utoipa::openapi::RefOr<utoipa::openapi::Schema>) {
let settings = rmcp::schemars::gen::SchemaSettings::openapi3();
let settings = rmcp::schemars::generate::SchemaSettings::openapi3();
let generator = settings.into_generator();
let schema = generator.into_root_schema_for::<$inner_type>();
let schema = convert_schemars_to_utoipa(schema);
@@ -44,119 +43,128 @@ macro_rules! derive_utoipa {
};
}
fn convert_schemars_to_utoipa(schema: rmcp::schemars::schema::RootSchema) -> RefOr<Schema> {
convert_schema_object(&rmcp::schemars::schema::Schema::Object(
schema.schema.clone(),
))
}
fn convert_schema_object(schema: &rmcp::schemars::schema::Schema) -> RefOr<Schema> {
match schema {
rmcp::schemars::schema::Schema::Object(schema_object) => {
convert_schema_object_inner(schema_object)
}
rmcp::schemars::schema::Schema::Bool(true) => {
RefOr::T(Schema::Object(ObjectBuilder::new().build()))
}
rmcp::schemars::schema::Schema::Bool(false) => {
RefOr::T(Schema::Object(ObjectBuilder::new().build()))
}
fn convert_schemars_to_utoipa(schema: rmcp::schemars::Schema) -> RefOr<Schema> {
// For schemars 1.0+, we need to work with the public API
// The schema is now a wrapper around a JSON Value that can be either an object or bool
if let Some(true) = schema.as_bool() {
return RefOr::T(Schema::Object(ObjectBuilder::new().build()));
}
if let Some(false) = schema.as_bool() {
return RefOr::T(Schema::Object(ObjectBuilder::new().build()));
}
// For object schemas, we'll need to work with the JSON Value directly
if let Some(obj) = schema.as_object() {
return convert_json_object_to_utoipa(obj);
}
// Fallback
RefOr::T(Schema::Object(ObjectBuilder::new().build()))
}
fn convert_schema_object_inner(schema: &SchemaObject) -> RefOr<Schema> {
// Handle references first
if let Some(reference) = &schema.reference {
fn convert_json_object_to_utoipa(
obj: &serde_json::Map<String, serde_json::Value>,
) -> RefOr<Schema> {
use serde_json::Value;
// Handle $ref
if let Some(Value::String(reference)) = obj.get("$ref") {
return RefOr::Ref(Ref::new(reference.clone()));
}
// Handle subschemas (oneOf, allOf, anyOf)
if let Some(subschemas) = &schema.subschemas {
if let Some(one_of) = &subschemas.one_of {
let schemas: Vec<RefOr<Schema>> = one_of.iter().map(convert_schema_object).collect();
let mut builder = OneOfBuilder::new();
for schema in schemas {
builder = builder.item(schema);
// Handle oneOf, allOf, anyOf
if let Some(Value::Array(one_of)) = obj.get("oneOf") {
let mut builder = OneOfBuilder::new();
for item in one_of {
if let Ok(schema) = rmcp::schemars::Schema::try_from(item.clone()) {
builder = builder.item(convert_schemars_to_utoipa(schema));
}
return RefOr::T(Schema::OneOf(builder.build()));
}
if let Some(all_of) = &subschemas.all_of {
let schemas: Vec<RefOr<Schema>> = all_of.iter().map(convert_schema_object).collect();
let mut all_of = AllOfBuilder::new();
for schema in schemas {
all_of = all_of.item(schema);
}
return RefOr::T(Schema::AllOf(all_of.build()));
}
if let Some(any_of) = &subschemas.any_of {
let schemas: Vec<RefOr<Schema>> = any_of.iter().map(convert_schema_object).collect();
let mut any_of = AnyOfBuilder::new();
for schema in schemas {
any_of = any_of.item(schema);
}
return RefOr::T(Schema::AnyOf(any_of.build()));
}
return RefOr::T(Schema::OneOf(builder.build()));
}
// Handle based on instance type
match &schema.instance_type {
Some(SingleOrVec::Single(instance_type)) => {
convert_single_instance_type(instance_type, schema)
}
Some(SingleOrVec::Vec(instance_types)) => {
// Multiple types - use AnyOf
let schemas: Vec<RefOr<Schema>> = instance_types
.iter()
.map(|instance_type| convert_single_instance_type(instance_type, schema))
.collect();
let mut any_of = AnyOfBuilder::new();
for schema in schemas {
any_of = any_of.item(schema);
if let Some(Value::Array(all_of)) = obj.get("allOf") {
let mut builder = AllOfBuilder::new();
for item in all_of {
if let Ok(schema) = rmcp::schemars::Schema::try_from(item.clone()) {
builder = builder.item(convert_schemars_to_utoipa(schema));
}
RefOr::T(Schema::AnyOf(any_of.build()))
}
None => {
// No type specified - create a generic schema
RefOr::T(Schema::Object(ObjectBuilder::new().build()))
return RefOr::T(Schema::AllOf(builder.build()));
}
if let Some(Value::Array(any_of)) = obj.get("anyOf") {
let mut builder = AnyOfBuilder::new();
for item in any_of {
if let Ok(schema) = rmcp::schemars::Schema::try_from(item.clone()) {
builder = builder.item(convert_schemars_to_utoipa(schema));
}
}
return RefOr::T(Schema::AnyOf(builder.build()));
}
// Handle type-based schemas
match obj.get("type") {
Some(Value::String(type_str)) => convert_typed_schema(type_str, obj),
Some(Value::Array(types)) => {
// Multiple types - use AnyOf
let mut builder = AnyOfBuilder::new();
for type_val in types {
if let Value::String(type_str) = type_val {
builder = builder.item(convert_typed_schema(type_str, obj));
}
}
RefOr::T(Schema::AnyOf(builder.build()))
}
None => RefOr::T(Schema::Object(ObjectBuilder::new().build())),
_ => RefOr::T(Schema::Object(ObjectBuilder::new().build())), // Handle other value types
}
}
fn convert_single_instance_type(
instance_type: &InstanceType,
schema: &SchemaObject,
fn convert_typed_schema(
type_str: &str,
obj: &serde_json::Map<String, serde_json::Value>,
) -> RefOr<Schema> {
match instance_type {
InstanceType::Object => {
use serde_json::Value;
match type_str {
"object" => {
let mut object_builder = ObjectBuilder::new();
if let Some(object_validation) = &schema.object {
// Add properties
for (name, prop_schema) in &object_validation.properties {
let prop = convert_schema_object(prop_schema);
object_builder = object_builder.property(name, prop);
// Add properties
if let Some(Value::Object(properties)) = obj.get("properties") {
for (name, prop_value) in properties {
if let Ok(prop_schema) = rmcp::schemars::Schema::try_from(prop_value.clone()) {
let prop = convert_schemars_to_utoipa(prop_schema);
object_builder = object_builder.property(name, prop);
}
}
}
// Add required fields
for required_field in &object_validation.required {
object_builder = object_builder.required(required_field);
// Add required fields
if let Some(Value::Array(required)) = obj.get("required") {
for req in required {
if let Value::String(field_name) = req {
object_builder = object_builder.required(field_name);
}
}
}
// Handle additional properties
if let Some(additional) = &object_validation.additional_properties {
match &**additional {
rmcp::schemars::schema::Schema::Bool(false) => {
object_builder = object_builder
.additional_properties(Some(AdditionalProperties::FreeForm(false)));
}
rmcp::schemars::schema::Schema::Bool(true) => {
object_builder = object_builder
.additional_properties(Some(AdditionalProperties::FreeForm(true)));
}
rmcp::schemars::schema::Schema::Object(obj) => {
let schema = convert_schema_object(
&rmcp::schemars::schema::Schema::Object(obj.clone()),
);
// Handle additional properties
if let Some(additional) = obj.get("additionalProperties") {
match additional {
Value::Bool(false) => {
object_builder = object_builder
.additional_properties(Some(AdditionalProperties::FreeForm(false)));
}
Value::Bool(true) => {
object_builder = object_builder
.additional_properties(Some(AdditionalProperties::FreeForm(true)));
}
_ => {
if let Ok(schema) = rmcp::schemars::Schema::try_from(additional.clone()) {
let schema = convert_schemars_to_utoipa(schema);
object_builder = object_builder
.additional_properties(Some(AdditionalProperties::RefOr(schema)));
}
@@ -166,117 +174,140 @@ fn convert_single_instance_type(
RefOr::T(Schema::Object(object_builder.build()))
}
InstanceType::Array => {
"array" => {
let mut array_builder = ArrayBuilder::new();
if let Some(array_validation) = &schema.array {
// Add items schema
if let Some(items) = &array_validation.items {
match items {
rmcp::schemars::schema::SingleOrVec::Single(item_schema) => {
let item_schema = convert_schema_object(item_schema);
// Add items schema
if let Some(items) = obj.get("items") {
match items {
Value::Object(_) | Value::Bool(_) => {
if let Ok(item_schema) = rmcp::schemars::Schema::try_from(items.clone()) {
let item_schema = convert_schemars_to_utoipa(item_schema);
array_builder = array_builder.items(item_schema);
}
rmcp::schemars::schema::SingleOrVec::Vec(item_schemas) => {
// Multiple item types - use AnyOf
let schemas: Vec<RefOr<Schema>> =
item_schemas.iter().map(convert_schema_object).collect();
let mut any_of = AnyOfBuilder::new();
for schema in schemas {
any_of = any_of.item(schema);
}
let any_of_schema = RefOr::T(Schema::AnyOf(any_of.build()));
array_builder = array_builder.items(any_of_schema);
}
}
Value::Array(item_schemas) => {
// Multiple item types - use AnyOf
let mut any_of = AnyOfBuilder::new();
for item in item_schemas {
if let Ok(schema) = rmcp::schemars::Schema::try_from(item.clone()) {
any_of = any_of.item(convert_schemars_to_utoipa(schema));
}
}
let any_of_schema = RefOr::T(Schema::AnyOf(any_of.build()));
array_builder = array_builder.items(any_of_schema);
}
_ => {}
}
}
// Add constraints
if let Some(min_items) = array_validation.min_items {
array_builder = array_builder.min_items(Some(min_items as usize));
// Add constraints
if let Some(Value::Number(min_items)) = obj.get("minItems") {
if let Some(min) = min_items.as_u64() {
array_builder = array_builder.min_items(Some(min as usize));
}
if let Some(max_items) = array_validation.max_items {
array_builder = array_builder.max_items(Some(max_items as usize));
}
if let Some(Value::Number(max_items)) = obj.get("maxItems") {
if let Some(max) = max_items.as_u64() {
array_builder = array_builder.max_items(Some(max as usize));
}
}
RefOr::T(Schema::Array(array_builder.build()))
}
InstanceType::String => {
"string" => {
let mut object_builder = ObjectBuilder::new().schema_type(SchemaType::String);
if let Some(string_validation) = &schema.string {
if let Some(min_length) = string_validation.min_length {
object_builder = object_builder.min_length(Some(min_length as usize));
}
if let Some(max_length) = string_validation.max_length {
object_builder = object_builder.max_length(Some(max_length as usize));
}
if let Some(pattern) = &string_validation.pattern {
object_builder = object_builder.pattern(Some(pattern.clone()));
if let Some(Value::Number(min_length)) = obj.get("minLength") {
if let Some(min) = min_length.as_u64() {
object_builder = object_builder.min_length(Some(min as usize));
}
}
if let Some(format) = &schema.format {
if let Some(Value::Number(max_length)) = obj.get("maxLength") {
if let Some(max) = max_length.as_u64() {
object_builder = object_builder.max_length(Some(max as usize));
}
}
if let Some(Value::String(pattern)) = obj.get("pattern") {
object_builder = object_builder.pattern(Some(pattern.clone()));
}
if let Some(Value::String(format)) = obj.get("format") {
object_builder = object_builder.format(Some(SchemaFormat::Custom(format.clone())));
}
RefOr::T(Schema::Object(object_builder.build()))
}
InstanceType::Number => {
"number" => {
let mut object_builder = ObjectBuilder::new().schema_type(SchemaType::Number);
if let Some(number_validation) = &schema.number {
if let Some(minimum) = number_validation.minimum {
object_builder = object_builder.minimum(Some(minimum));
if let Some(Value::Number(minimum)) = obj.get("minimum") {
if let Some(min) = minimum.as_f64() {
object_builder = object_builder.minimum(Some(min));
}
if let Some(maximum) = number_validation.maximum {
object_builder = object_builder.maximum(Some(maximum));
}
if let Some(Value::Number(maximum)) = obj.get("maximum") {
if let Some(max) = maximum.as_f64() {
object_builder = object_builder.maximum(Some(max));
}
if let Some(exclusive_minimum) = number_validation.exclusive_minimum {
object_builder = object_builder.exclusive_minimum(Some(exclusive_minimum));
}
if let Some(Value::Number(exclusive_minimum)) = obj.get("exclusiveMinimum") {
if let Some(min) = exclusive_minimum.as_f64() {
object_builder = object_builder.exclusive_minimum(Some(min));
}
if let Some(exclusive_maximum) = number_validation.exclusive_maximum {
object_builder = object_builder.exclusive_maximum(Some(exclusive_maximum));
}
if let Some(Value::Number(exclusive_maximum)) = obj.get("exclusiveMaximum") {
if let Some(max) = exclusive_maximum.as_f64() {
object_builder = object_builder.exclusive_maximum(Some(max));
}
if let Some(multiple_of) = number_validation.multiple_of {
object_builder = object_builder.multiple_of(Some(multiple_of));
}
if let Some(Value::Number(multiple_of)) = obj.get("multipleOf") {
if let Some(mult) = multiple_of.as_f64() {
object_builder = object_builder.multiple_of(Some(mult));
}
}
RefOr::T(Schema::Object(object_builder.build()))
}
InstanceType::Integer => {
"integer" => {
let mut object_builder = ObjectBuilder::new().schema_type(SchemaType::Integer);
if let Some(number_validation) = &schema.number {
if let Some(minimum) = number_validation.minimum {
object_builder = object_builder.minimum(Some(minimum));
if let Some(Value::Number(minimum)) = obj.get("minimum") {
if let Some(min) = minimum.as_f64() {
object_builder = object_builder.minimum(Some(min));
}
if let Some(maximum) = number_validation.maximum {
object_builder = object_builder.maximum(Some(maximum));
}
if let Some(Value::Number(maximum)) = obj.get("maximum") {
if let Some(max) = maximum.as_f64() {
object_builder = object_builder.maximum(Some(max));
}
if let Some(exclusive_minimum) = number_validation.exclusive_minimum {
object_builder = object_builder.exclusive_minimum(Some(exclusive_minimum));
}
if let Some(Value::Number(exclusive_minimum)) = obj.get("exclusiveMinimum") {
if let Some(min) = exclusive_minimum.as_f64() {
object_builder = object_builder.exclusive_minimum(Some(min));
}
if let Some(exclusive_maximum) = number_validation.exclusive_maximum {
object_builder = object_builder.exclusive_maximum(Some(exclusive_maximum));
}
if let Some(Value::Number(exclusive_maximum)) = obj.get("exclusiveMaximum") {
if let Some(max) = exclusive_maximum.as_f64() {
object_builder = object_builder.exclusive_maximum(Some(max));
}
if let Some(multiple_of) = number_validation.multiple_of {
object_builder = object_builder.multiple_of(Some(multiple_of));
}
if let Some(Value::Number(multiple_of)) = obj.get("multipleOf") {
if let Some(mult) = multiple_of.as_f64() {
object_builder = object_builder.multiple_of(Some(mult));
}
}
RefOr::T(Schema::Object(object_builder.build()))
}
InstanceType::Boolean => RefOr::T(Schema::Object(
"boolean" => RefOr::T(Schema::Object(
ObjectBuilder::new()
.schema_type(SchemaType::Boolean)
.build(),
)),
InstanceType::Null => RefOr::T(Schema::Object(
"null" => RefOr::T(Schema::Object(
ObjectBuilder::new().schema_type(SchemaType::String).build(),
)),
_ => RefOr::T(Schema::Object(ObjectBuilder::new().build())),
}
}
@@ -285,11 +316,40 @@ derive_utoipa!(Content as ContentSchema);
derive_utoipa!(EmbeddedResource as EmbeddedResourceSchema);
derive_utoipa!(ImageContent as ImageContentSchema);
derive_utoipa!(TextContent as TextContentSchema);
derive_utoipa!(RawTextContent as RawTextContentSchema);
derive_utoipa!(RawImageContent as RawImageContentSchema);
derive_utoipa!(RawEmbeddedResource as RawEmbeddedResourceSchema);
derive_utoipa!(Tool as ToolSchema);
derive_utoipa!(ToolAnnotations as ToolAnnotationsSchema);
derive_utoipa!(Annotations as AnnotationsSchema);
derive_utoipa!(ResourceContents as ResourceContentsSchema);
// Create a manual schema for the generic Annotated type
// We manually define this to avoid circular references from RawContent::Audio(AudioContent)
// where AudioContent = Annotated<RawAudioContent>
struct AnnotatedSchema {}
impl<'__s> ToSchema<'__s> for AnnotatedSchema {
fn schema() -> (&'__s str, utoipa::openapi::RefOr<utoipa::openapi::Schema>) {
// Create a oneOf schema with only the variants we actually use in the API
// This avoids the circular reference from RawContent::Audio(AudioContent)
let schema = Schema::OneOf(
OneOfBuilder::new()
.item(RefOr::Ref(Ref::new("#/components/schemas/RawTextContent")))
.item(RefOr::Ref(Ref::new("#/components/schemas/RawImageContent")))
.item(RefOr::Ref(Ref::new(
"#/components/schemas/RawEmbeddedResource",
)))
.build(),
);
("Annotated", RefOr::T(schema))
}
fn aliases() -> Vec<(&'__s str, utoipa::openapi::schema::Schema)> {
Vec::new()
}
}
#[allow(dead_code)] // Used by utoipa for OpenAPI generation
#[derive(OpenApi)]
#[openapi(
@@ -349,6 +409,10 @@ derive_utoipa!(ResourceContents as ResourceContentsSchema);
ImageContentSchema,
AnnotationsSchema,
TextContentSchema,
RawTextContentSchema,
RawImageContentSchema,
RawEmbeddedResourceSchema,
AnnotatedSchema,
ToolResponse,
ToolRequest,
ToolConfirmationRequest,