diff --git a/golem-worker-service/src/mcp/agent_mcp_capability.rs b/golem-worker-service/src/mcp/agent_mcp_capability.rs index 621d2356f4..d194530d09 100644 --- a/golem-worker-service/src/mcp/agent_mcp_capability.rs +++ b/golem-worker-service/src/mcp/agent_mcp_capability.rs @@ -14,7 +14,7 @@ use crate::mcp::agent_mcp_resource::AgentMcpResource; use crate::mcp::agent_mcp_tool::AgentMcpTool; -use crate::mcp::mcp_schema::{GetMcpSchema, GetMcpToolSchema, McpToolSchema}; +use crate::mcp::schema::{GetMcpSchema, GetMcpToolSchema, McpSchema, McpToolSchema}; use golem_common::base_model::account::AccountId; use golem_common::base_model::agent::{AgentMethod, AgentTypeName, DataSchema}; use golem_common::base_model::component::ComponentId; @@ -50,8 +50,10 @@ impl McpAgentCapability { ); let constructor_schema = constructor.input_schema.get_mcp_schema(); + let mut tool_schema = method.get_mcp_tool_schema(); - tool_schema.merge_input_schema(constructor_schema); + + tool_schema.prepend_input_schema(constructor_schema); let McpToolSchema { input_schema, @@ -62,8 +64,9 @@ impl McpAgentCapability { name: Cow::from(get_tool_name(agent_type_name, method)), title: None, description: Some(method.description.clone().into()), - input_schema: Arc::new(input_schema), - output_schema: output_schema.map(Arc::new), + input_schema: Arc::new(McpSchema::from(input_schema)), + output_schema: output_schema + .map(|internal| Arc::new(McpSchema::from(internal))), annotations: None, execution: None, icons: None, diff --git a/golem-worker-service/src/mcp/agent_mcp_server.rs b/golem-worker-service/src/mcp/agent_mcp_server.rs index 48be6e0c3c..7e2a08c573 100644 --- a/golem-worker-service/src/mcp/agent_mcp_server.rs +++ b/golem-worker-service/src/mcp/agent_mcp_server.rs @@ -23,6 +23,7 @@ use golem_common::base_model::agent::{ use golem_common::base_model::domain_registration::Domain; use golem_common::model::WorkerId; use golem_common::model::agent::{NamedElementSchema, UntypedDataValue, UntypedElementValue}; +use golem_wasm::analysis::AnalysedType; use golem_wasm::json::ValueAndTypeJsonExtensions; use golem_wasm::{IntoValue, ValueAndType}; use poem::http; @@ -186,12 +187,19 @@ where { match elem_schema { ElementSchema::ComponentModel(ComponentModelElementSchema { element_type }) => { - let value = args_map - .get(name) - .ok_or_else(|| format!("Missing parameter: {}", name))?; + let json_value = match args_map.get(name) { + Some(value) => value.clone(), + None => { + if matches!(element_type, AnalysedType::Option(_)) { + serde_json::Value::Null + } else { + return Err(format!("Missing parameter: {}", name)); + } + } + }; let value_and_type = - golem_wasm::ValueAndType::parse_with_type(value, element_type) + golem_wasm::ValueAndType::parse_with_type(&json_value, element_type) .map_err(|errs| { format!( "Failed to parse parameter '{}': {}", diff --git a/golem-worker-service/src/mcp/mcp_schema.rs b/golem-worker-service/src/mcp/mcp_schema.rs deleted file mode 100644 index df1adc0fe4..0000000000 --- a/golem-worker-service/src/mcp/mcp_schema.rs +++ /dev/null @@ -1,116 +0,0 @@ -// Copyright 2024-2025 Golem Cloud -// -// Licensed under the Golem Source License v1.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://license.golem.cloud/LICENSE -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use golem_common::base_model::agent::{ - AgentMethod, ComponentModelElementSchema, DataSchema, ElementSchema, NamedElementSchema, -}; -use golem_wasm::analysis::AnalysedType; -use rmcp::model::JsonObject; -use serde_json::{Map, Value, json}; - -pub trait GetMcpSchema { - fn get_mcp_schema(&self) -> JsonObject; -} - -impl GetMcpSchema for DataSchema { - fn get_mcp_schema(&self) -> JsonObject { - match self { - DataSchema::Tuple(schemas) => { - let properties = get_json_schema(&schemas.elements); - json!({ - "type": "object", - "properties": properties, - }) - .as_object() - .unwrap() - .clone() - } - DataSchema::Multimodal(_) => { - todo!("Multimodal schema is not supported in this example") - } - } - } -} - -pub trait GetMcpToolSchema { - fn get_mcp_tool_schema(&self) -> McpToolSchema; -} - -pub struct McpToolSchema { - pub input_schema: JsonObject, - pub output_schema: Option, // TODO; may be avoid Option -} - -impl McpToolSchema { - pub fn merge_input_schema(&mut self, input_schema: JsonObject) { - let mut new_properties = input_schema - .get("properties") - .and_then(|props| props.as_object()) - .cloned() - .unwrap_or_default(); - - if let Some(existing_properties) = self - .input_schema - .get("properties") - .and_then(|props| props.as_object()) - { - for (key, value) in existing_properties { - new_properties.insert(key.clone(), value.clone()); - } - } - - self.input_schema = json!({ - "type": "object", - "properties": new_properties, - }) - .as_object() - .unwrap() - .clone(); - } -} - -impl GetMcpToolSchema for AgentMethod { - fn get_mcp_tool_schema(&self) -> McpToolSchema { - let input_schema: JsonObject = self.input_schema.get_mcp_schema(); - let output_schema: JsonObject = self.output_schema.get_mcp_schema(); - - McpToolSchema { - input_schema, - output_schema: Some(output_schema), - } - } -} - -fn get_json_schema(schemas: &Vec) -> Map { - let mut properties = Map::new(); - - for NamedElementSchema { name, schema } in schemas { - let json_schema = match schema { - ElementSchema::ComponentModel(ComponentModelElementSchema { element_type }) => { - match element_type { - AnalysedType::Str(_) => json!({"type": "string"}), - AnalysedType::U32(_) => json!({"type": "integer"}), - AnalysedType::Bool(_) => json!({"type": "boolean"}), - _ => { - todo!("Unsupported component model element type in schema mapping") - } - } - } - _ => todo!("Unsupported component model element type in schema mapping"), - }; - properties.insert(name.clone(), json_schema); - } - - properties -} diff --git a/golem-worker-service/src/mcp/mod.rs b/golem-worker-service/src/mcp/mod.rs index c7478bb09a..d919eb39b2 100644 --- a/golem-worker-service/src/mcp/mod.rs +++ b/golem-worker-service/src/mcp/mod.rs @@ -7,4 +7,4 @@ mod agent_mcp_resource; mod agent_mcp_server; mod agent_mcp_tool; mod mcp_capabilities_lookup; -mod mcp_schema; +mod schema; diff --git a/golem-worker-service/src/mcp/schema/internal.rs b/golem-worker-service/src/mcp/schema/internal.rs new file mode 100644 index 0000000000..21f59b134f --- /dev/null +++ b/golem-worker-service/src/mcp/schema/internal.rs @@ -0,0 +1,221 @@ +// Copyright 2024-2025 Golem Cloud +// +// Licensed under the Golem Source License v1.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://license.golem.cloud/LICENSE +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use crate::mcp::schema::McpSchema; +use golem_common::base_model::agent::{ + ComponentModelElementSchema, ElementSchema, NamedElementSchema, +}; +use golem_wasm::analysis::AnalysedType; +use serde_json::{Map, Value, json}; + +// A better internal representation of McpSchema for object types + +#[derive(Default)] +pub struct McpSchemaInternal { + pub properties: Map, + pub required: Vec, +} + +impl From for McpSchema { + fn from(value: McpSchemaInternal) -> Self { + let json_value = json!({ + "type": "object", + "properties": value.properties, + "required": value.required, + }); + + rmcp::model::object(json_value) + } +} + +impl McpSchemaInternal { + pub fn prepend_schema(&mut self, mut new_schema: McpSchemaInternal) { + new_schema + .properties + .extend(std::mem::take(&mut self.properties)); + + new_schema + .required + .extend(std::mem::take(&mut self.required)); + + *self = new_schema; + } + + pub fn from_named_element_schemas(schemas: &Vec) -> McpSchemaInternal { + let named_types: Vec<(&str, &AnalysedType)> = schemas + .iter() + .map(|s| match &s.schema { + ElementSchema::ComponentModel(ComponentModelElementSchema { element_type }) => { + (s.name.as_str(), element_type) + } + _ => todo!("Unsupported element schema type in MCP schema mapping"), + }) + .collect(); + + Self::from_record_fields(&named_types) + } + + pub fn from_record_fields(fields: &[(&str, &AnalysedType)]) -> McpSchemaInternal { + let mut properties: Map = Map::new(); + let mut required = Vec::new(); + + for (name, typ) in fields { + properties.insert(name.to_string(), analysed_type_to_json_schema(typ)); + if !matches!(typ, AnalysedType::Option(_)) { + required.push(name.to_string()); + } + } + + McpSchemaInternal { + properties, + required, + } + } +} + +pub type JsonTypeDescription = Value; +pub type FieldName = String; + +// Based on https://modelcontextprotocol.io/specification/2025-11-25/server/tools and +// https://json-schema.org/draft/2020-12/json-schema-core (Example: oneOf) +// while ensuring how golem-wasm treats json values +fn analysed_type_to_json_schema(analysed_type: &AnalysedType) -> JsonTypeDescription { + match analysed_type { + AnalysedType::Bool(_) => json!({"type": "boolean"}), + AnalysedType::Str(_) => json!({"type": "string"}), + AnalysedType::Chr(_) => json!({"type": "integer"}), + AnalysedType::U8(_) + | AnalysedType::U16(_) + | AnalysedType::U32(_) + | AnalysedType::U64(_) + | AnalysedType::S8(_) + | AnalysedType::S16(_) + | AnalysedType::S32(_) + | AnalysedType::S64(_) => json!({"type": "integer"}), + AnalysedType::F32(_) | AnalysedType::F64(_) => json!({"type": "number"}), + + AnalysedType::List(type_list) => { + let items = analysed_type_to_json_schema(&type_list.inner); + json!({"type": "array", "items": items}) + } + + AnalysedType::Tuple(type_tuple) => { + let prefix_items: Vec = type_tuple + .items + .iter() + .map(analysed_type_to_json_schema) + .collect(); + + json!({ + "type": "array", + "prefixItems": prefix_items, + "items": false + }) + } + + AnalysedType::Record(type_record) => { + let fields: Vec<(&str, &AnalysedType)> = type_record + .fields + .iter() + .map(|f| (f.name.as_str(), &f.typ)) + .collect(); + + let schema = McpSchemaInternal::from_record_fields(&fields); + + json!({ + "type": "object", + "properties": schema.properties, + "required": schema.required + }) + } + + AnalysedType::Option(type_option) => { + let inner = analysed_type_to_json_schema(&type_option.inner); + json!({ + "oneOf": [ + inner, + {"type": "null"} + ] + }) + } + + AnalysedType::Enum(type_enum) => { + json!({"type": "string", "enum": type_enum.cases}) + } + + // Flags → JSON array of enabled flag name strings + AnalysedType::Flags(type_flags) => { + json!({ + "type": "array", + "items": {"type": "string", "enum": type_flags.names}, + "uniqueItems": true + }) + } + + // Variant → object with a single key being the case name + // Aligned with to_json_value: {"case_name": value} or {"case_name": null} + AnalysedType::Variant(type_variant) => { + let one_of: Vec = type_variant + .cases + .iter() + .map(|case| { + let value_schema = match &case.typ { + Some(payload_type) => analysed_type_to_json_schema(payload_type), + None => json!({"type": "null"}), + }; + json!({ + "type": "object", + "properties": { + case.name.clone(): value_schema, + }, + "required": [case.name], + "additionalProperties": false + }) + }) + .collect(); + json!({"oneOf": one_of}) + } + + AnalysedType::Result(type_result) => { + let ok_schema = match &type_result.ok { + Some(ok_type) => analysed_type_to_json_schema(ok_type), + None => json!({"type": "null"}), + }; + let err_schema = match &type_result.err { + Some(err_type) => analysed_type_to_json_schema(err_type), + None => json!({"type": "null"}), + }; + json!({ + "oneOf": [ + { + "type": "object", + "properties": {"ok": ok_schema}, + "required": ["ok"], + "additionalProperties": false + }, + { + "type": "object", + "properties": {"err": err_schema}, + "required": ["err"], + "additionalProperties": false + } + ] + }) + } + + AnalysedType::Handle(_) => { + json!({"type": "string"}) + } + } +} diff --git a/golem-worker-service/src/mcp/schema/mcp_schema.rs b/golem-worker-service/src/mcp/schema/mcp_schema.rs new file mode 100644 index 0000000000..6724673dc1 --- /dev/null +++ b/golem-worker-service/src/mcp/schema/mcp_schema.rs @@ -0,0 +1,36 @@ +// Copyright 2024-2025 Golem Cloud +// +// Licensed under the Golem Source License v1.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://license.golem.cloud/LICENSE +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use crate::mcp::schema::internal::McpSchemaInternal; +use golem_common::base_model::agent::DataSchema; +use rmcp::model::JsonObject; + +pub type McpSchema = JsonObject; + +pub trait GetMcpSchema { + fn get_mcp_schema(&self) -> McpSchemaInternal; +} + +impl GetMcpSchema for DataSchema { + fn get_mcp_schema(&self) -> McpSchemaInternal { + match self { + DataSchema::Tuple(schemas) => { + McpSchemaInternal::from_named_element_schemas(&schemas.elements) + } + DataSchema::Multimodal(_) => { + todo!("Multimodal schema is not supported in this example") + } + } + } +} diff --git a/golem-worker-service/src/mcp/schema/mcp_tool_schema.rs b/golem-worker-service/src/mcp/schema/mcp_tool_schema.rs new file mode 100644 index 0000000000..5e44391c8a --- /dev/null +++ b/golem-worker-service/src/mcp/schema/mcp_tool_schema.rs @@ -0,0 +1,45 @@ +// Copyright 2024-2025 Golem Cloud +// +// Licensed under the Golem Source License v1.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://license.golem.cloud/LICENSE +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use crate::mcp::schema::McpSchema; +use crate::mcp::schema::internal::McpSchemaInternal; +use crate::mcp::schema::mcp_schema::GetMcpSchema; +use golem_common::base_model::agent::AgentMethod; + +pub trait GetMcpToolSchema { + fn get_mcp_tool_schema(&self) -> McpToolSchema; +} + +impl GetMcpToolSchema for AgentMethod { + fn get_mcp_tool_schema(&self) -> McpToolSchema { + let input_schema: McpSchemaInternal = self.input_schema.get_mcp_schema(); + let output_schema: McpSchemaInternal = self.output_schema.get_mcp_schema(); + + McpToolSchema { + input_schema, + output_schema: Some(output_schema), + } + } +} + +pub struct McpToolSchema { + pub input_schema: McpSchemaInternal, + pub output_schema: Option, +} + +impl McpToolSchema { + pub fn prepend_input_schema(&mut self, input_schema: McpSchemaInternal) { + self.input_schema.prepend_schema(input_schema); + } +} diff --git a/golem-worker-service/src/mcp/schema/mod.rs b/golem-worker-service/src/mcp/schema/mod.rs new file mode 100644 index 0000000000..a5bd2980ea --- /dev/null +++ b/golem-worker-service/src/mcp/schema/mod.rs @@ -0,0 +1,20 @@ +// Copyright 2024-2025 Golem Cloud +// +// Licensed under the Golem Source License v1.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://license.golem.cloud/LICENSE +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +pub use mcp_schema::*; +pub use mcp_tool_schema::*; + +mod internal; +mod mcp_schema; +mod mcp_tool_schema;