Support MCP resources, multimodal, prompts, unstructured-* with integration tests and other fixes

Author: afsalthaj

Cargo.lock

@@ -547,15 +547,45 @@ pub enum UntypedDataValue {
 }
 
 #[derive(Debug, Clone, PartialEq)]
-#[cfg_attr(
-    feature = "full",
-    derive(IntoValue, FromValue, desert_rust::BinaryCodec)
-)]
+#[cfg_attr(feature = "full", derive(desert_rust::BinaryCodec))]
 pub struct UntypedNamedElementValue {
     pub name: String,
     pub value: UntypedElementValue,
 }
 
+#[cfg(feature = "full")]
+impl golem_wasm::FromValue for UntypedNamedElementValue {
+    fn from_value(value: Value) -> Result<Self, String> {
+        match value {
+            Value::Tuple(fields) if fields.len() == 2 => {
+                let mut iter = fields.into_iter();
+                let name = String::from_value(iter.next().unwrap())?;
+                let value = UntypedElementValue::from_value(iter.next().unwrap())?;
+                Ok(UntypedNamedElementValue { name, value })
+            }
+            _ => Err(format!(
+                "Expected Tuple with 2 fields for UntypedNamedElementValue, got {:?}",
+                value
+            )),
+        }
+    }
+}
+
+#[cfg(feature = "full")]
+impl golem_wasm::IntoValue for UntypedNamedElementValue {
+    fn into_value(self) -> Value {
+        Value::Tuple(vec![self.name.into_value(), self.value.into_value()])
+    }
+
+    fn get_type() -> AnalysedType {
+        AnalysedType::Tuple(golem_wasm::analysis::TypeTuple {
+            name: None,
+            owner: None,
+            items: vec![String::get_type(), UntypedElementValue::get_type()],
+        })
+    }
+}
+
 #[derive(Debug, Clone, PartialEq)]
 #[cfg_attr(
     feature = "full",

golem-common/src/base_model/agent.rs

@@ -38,6 +38,7 @@ golem-wasm = { workspace = true, default-features = true }
 golem-wasm-derive = { workspace = true }
 
 anyhow = { workspace = true }
+base64 = { workspace = true }
 async-trait = { workspace = true }
 bigdecimal = { workspace = true }
 bytes = { workspace = true }

golem-worker-service/Cargo.toml

@@ -12,93 +12,157 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use crate::mcp::agent_mcp_resource::AgentMcpResource;
+use crate::mcp::agent_mcp_resource::{AgentMcpResource, AgentMcpResourceKind};
 use crate::mcp::agent_mcp_tool::AgentMcpTool;
-use crate::mcp::schema::{McpToolSchema, get_mcp_schema, get_mcp_tool_schema};
+use crate::mcp::schema::{McpToolSchema, get_input_mcp_schema, get_mcp_tool_schema};
 use golem_common::base_model::account::AccountId;
-use golem_common::base_model::agent::{AgentMethod, AgentTypeName, DataSchema};
+use golem_common::base_model::agent::{
+    AgentMethod, AgentTypeName, DataSchema, ElementSchema, NamedElementSchemas,
+};
 use golem_common::base_model::component::ComponentId;
 use golem_common::base_model::environment::EnvironmentId;
 use golem_common::model::agent::AgentConstructor;
-use rmcp::model::Tool;
+use rmcp::model::{Annotated, RawResource, RawResourceTemplate, Tool};
 use std::borrow::Cow;
 use std::sync::Arc;
 
 #[derive(Clone)]
 pub enum McpAgentCapability {
     Tool(Box<AgentMcpTool>),
-    #[allow(unused)]
-    Resource(AgentMcpResource),
+    Resource(Box<AgentMcpResource>),
 }
 
 impl McpAgentCapability {
-    pub fn from(
+    pub fn from_agent_method(
         account_id: &AccountId,
         environment_id: &EnvironmentId,
         agent_type_name: &AgentTypeName,
         method: &AgentMethod,
         constructor: &AgentConstructor,
         component_id: ComponentId,
     ) -> Self {
-        match &method.input_schema {
-            DataSchema::Tuple(schemas) => {
-                if !schemas.elements.is_empty() {
-                    tracing::debug!(
-                        "Method {} of agent type {} has input parameters, exposing as tool",
-                        method.name,
-                        agent_type_name.0
-                    );
-
-                    let constructor_schema = get_mcp_schema(&constructor.input_schema);
-
-                    let McpToolSchema {
-                        mut input_schema,
-                        output_schema,
-                    } = get_mcp_tool_schema(method);
-
-                    input_schema.prepend_schema(constructor_schema);
-
-                    let tool = Tool {
-                        name: Cow::from(get_tool_name(agent_type_name, method)),
+        let schemas = match &method.input_schema {
+            DataSchema::Tuple(schemas) | DataSchema::Multimodal(schemas) => schemas,
+        };
+
+        if !schemas.elements.is_empty() {
+            tracing::debug!(
+                "Method {} of agent type {} has input parameters, exposing as tool",
+                method.name,
+                agent_type_name.0
+            );
+
+            let constructor_schema = get_input_mcp_schema(&constructor.input_schema);
+
+            let McpToolSchema {
+                mut input_schema,
+                output_schema,
+            } = get_mcp_tool_schema(method);
+
+            input_schema.prepend_schema(constructor_schema);
+
+            let tool = Tool {
+                name: Cow::from(get_tool_name(agent_type_name, method)),
+                title: None,
+                description: Some(method.description.clone().into()),
+                input_schema: Arc::new(rmcp::model::JsonObject::from(input_schema)),
+                output_schema: output_schema
+                    .map(|internal| Arc::new(rmcp::model::JsonObject::from(internal))),
+                annotations: None,
+                execution: None,
+                icons: None,
+                meta: None,
+            };
+
+            Self::Tool(Box::new(AgentMcpTool {
+                environment_id: *environment_id,
+                account_id: *account_id,
+                constructor: constructor.clone(),
+                raw_method: method.clone(),
+                tool,
+                component_id,
+                agent_type_name: agent_type_name.clone(),
+            }))
+        } else {
+            tracing::debug!(
+                "Method {} of agent type {} has no input parameters, exposing as resource",
+                method.name,
+                agent_type_name.0
+            );
+
+            let constructor_param_names = AgentMcpResource::constructor_param_names(constructor);
+            let name = AgentMcpResource::resource_name(agent_type_name, method);
+
+            let mime_type = output_resource_mime_type(&method.output_schema);
+
+            let kind = if constructor_param_names.is_empty() {
+                let uri = AgentMcpResource::static_uri(agent_type_name, method);
+                AgentMcpResourceKind::Static(Annotated::new(
+                    RawResource {
+                        uri,
+                        name,
                         title: None,
-                        description: Some(method.description.clone().into()),
-                        input_schema: Arc::new(rmcp::model::JsonObject::from(input_schema)),
-                        output_schema: output_schema
-                            .map(|internal| Arc::new(rmcp::model::JsonObject::from(internal))),
-                        annotations: None,
-                        execution: None,
+                        description: Some(method.description.clone()),
+                        mime_type,
+                        size: None,
                         icons: None,
                         meta: None,
-                    };
-
-                    Self::Tool(Box::new(AgentMcpTool {
-                        environment_id: *environment_id,
-                        account_id: *account_id,
-                        constructor: constructor.clone(),
-                        raw_method: method.clone(),
-                        tool,
-                        component_id,
-                        agent_type_name: agent_type_name.clone(),
-                    }))
-                } else {
-                    tracing::debug!(
-                        "Method {} of agent type {} has no input parameters, exposing as resource",
-                        method.name,
-                        agent_type_name.0
-                    );
-
-                    Self::Resource(AgentMcpResource {
-                        resource: method.clone(),
-                    })
+                    },
+                    None,
+                ))
+            } else {
+                let uri_template = AgentMcpResource::template_uri(
+                    agent_type_name,
+                    method,
+                    &constructor_param_names,
+                );
+                AgentMcpResourceKind::Template {
+                    template: Annotated::new(
+                        RawResourceTemplate {
+                            uri_template,
+                            name,
+                            title: None,
+                            description: Some(method.description.clone()),
+                            mime_type,
+                            icons: None,
+                        },
+                        None,
+                    ),
+                    constructor_param_names,
                 }
-            }
-            DataSchema::Multimodal(_) => {
-                todo!("Multimodal schema handling not implemented yet")
-            }
+            };
+
+            Self::Resource(Box::new(AgentMcpResource {
+                kind,
+                environment_id: *environment_id,
+                account_id: *account_id,
+                constructor: constructor.clone(),
+                raw_method: method.clone(),
+                component_id,
+                agent_type_name: agent_type_name.clone(),
+            }))
         }
     }
 }
 
 fn get_tool_name(agent_type_name: &AgentTypeName, method: &AgentMethod) -> String {
     format!("{}-{}", agent_type_name.0, method.name)
 }
+
+fn output_resource_mime_type(output_schema: &DataSchema) -> Option<String> {
+    match output_schema {
+        DataSchema::Tuple(NamedElementSchemas { elements }) => match elements.as_slice() {
+            [single] => match &single.schema {
+                ElementSchema::ComponentModel(_) => Some("application/json".to_string()),
+                ElementSchema::UnstructuredText(_) => Some("text/plain".to_string()),
+                // The actual mime type
+                ElementSchema::UnstructuredBinary(_) => None,
+            },
+            _ => None,
+        },
+
+        // Each individual resource contents could have its own mime type, so we can't assign a single mime type to the whole output
+        // when it comes to multimodal output schemas.
+        DataSchema::Multimodal(_) => None,
+    }
+}