clean ups and optimisation in mcp (#3012)

afsalthaj · web-flow · commit 09b4b06f4cbb · 2026-03-18T20:56:40.000+11:00
diff --git a/golem-api-grpc/proto/golem/mcp/core.proto b/golem-api-grpc/proto/golem/mcp/core.proto
@@ -14,4 +14,5 @@ message CompiledMcp {
   string domain = 4;
   map<string, golem.registry.RegisteredAgentTypeImplementer> agent_type_implementers = 5;
   golem.customapi.SecuritySchemeDetails security_scheme = 6;
+  repeated golem.registry.RegisteredAgentType registered_agent_types = 7;
 }
diff --git a/golem-registry-service/src/repo/model/deployment.rs b/golem-registry-service/src/repo/model/deployment.rs
@@ -599,6 +599,7 @@ impl TryFrom<DeploymentCompiledMcpRecord> for CompiledMcp {
             domain: Domain(value.domain),
             agent_type_implementers: mcp_data.implementers,
             security_scheme: None,
+            registered_agent_types: Vec::new(),
         })
     }
 }
diff --git a/golem-registry-service/src/services/deployment/deployment_context.rs b/golem-registry-service/src/services/deployment/deployment_context.rs
@@ -297,6 +297,7 @@ impl DeploymentContext {
                 domain: domain.clone(),
                 agent_type_implementers,
                 security_scheme,
+                registered_agent_types: Vec::new(),
             };
             all_compiled_mcps.push(compiled_mcp);
         }
diff --git a/golem-registry-service/src/services/deployment/mcp.rs b/golem-registry-service/src/services/deployment/mcp.rs
@@ -2,6 +2,7 @@ use crate::repo::deployment::DeploymentRepo;
 use crate::repo::model::deployment::DeployRepoError;
 use crate::services::security_scheme::SecuritySchemeService;
 use golem_common::base_model::domain_registration::Domain;
+use golem_common::model::agent::RegisteredAgentType;
 use golem_common::{SafeDisplay, error_forwarding};
 use golem_service_base::custom_api::SecuritySchemeDetails;
 use golem_service_base::mcp::CompiledMcp;
@@ -56,6 +57,7 @@ impl DeployedMcpService {
         match optional_deployment {
             Some(deployment) => {
                 let security_scheme_id = deployment.mcp_data.value().security_scheme_id;
+                let environment_id = deployment.environment_id;
                 let mut compiled_mcp = CompiledMcp::try_from(deployment)?;
 
                 if let Some(scheme_id) = security_scheme_id {
@@ -76,6 +78,50 @@ impl DeployedMcpService {
                     });
                 }
 
+                let mut registered_agent_types = Vec::new();
+                for agent_type_name in compiled_mcp.agent_type_implementers.keys() {
+                    match self
+                        .deployment_repo
+                        .get_deployed_agent_type(environment_id, &agent_type_name.0)
+                        .await
+                    {
+                        Ok(Some(record)) => {
+                            match golem_common::model::agent::DeployedRegisteredAgentType::try_from(
+                                record,
+                            ) {
+                                Ok(deployed) => {
+                                    registered_agent_types
+                                        .push(RegisteredAgentType::from(deployed));
+                                }
+                                Err(e) => {
+                                    tracing::warn!(
+                                        "Failed to convert agent type {} for domain {}: {}",
+                                        agent_type_name.0,
+                                        domain.0,
+                                        e
+                                    );
+                                }
+                            }
+                        }
+                        Ok(None) => {
+                            tracing::warn!(
+                                "Agent type {} not found for domain {}",
+                                agent_type_name.0,
+                                domain.0,
+                            );
+                        }
+                        Err(e) => {
+                            tracing::warn!(
+                                "Failed to fetch agent type {} for domain {}: {}",
+                                agent_type_name.0,
+                                domain.0,
+                                e
+                            );
+                        }
+                    }
+                }
+                compiled_mcp.registered_agent_types = registered_agent_types;
+
                 Ok(compiled_mcp)
             }
             None => Err(DeployedMcpError::NoActiveMcpForDomain(domain.clone())),
diff --git a/golem-service-base/src/mcp/mod.rs b/golem-service-base/src/mcp/mod.rs
@@ -1,24 +1,40 @@
+// Copyright 2024-2026 Golem Cloud
+//
+// Licensed under the Golem Source License v1.1 (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.
+
 mod protobuf;
 
 use golem_common::base_model::account::AccountId;
 use golem_common::base_model::deployment::DeploymentRevision;
 use golem_common::base_model::domain_registration::Domain;
 use golem_common::base_model::environment::EnvironmentId;
-use golem_common::model::agent::AgentTypeName;
+use golem_common::model::agent::{AgentTypeName, RegisteredAgentType};
 use golem_common::model::component::{ComponentId, ComponentRevision};
 use std::collections::HashMap;
 
 use crate::custom_api::SecuritySchemeDetails;
 
 pub type AgentTypeImplementers = HashMap<AgentTypeName, (ComponentId, ComponentRevision)>;
 
+#[derive(Clone)]
 pub struct CompiledMcp {
     pub account_id: AccountId,
     pub environment_id: EnvironmentId,
     pub deployment_revision: DeploymentRevision,
     pub domain: Domain,
     pub agent_type_implementers: AgentTypeImplementers,
     pub security_scheme: Option<SecuritySchemeDetails>,
+    pub registered_agent_types: Vec<RegisteredAgentType>,
 }
 
 impl CompiledMcp {
diff --git a/golem-service-base/src/mcp/protobuf.rs b/golem-service-base/src/mcp/protobuf.rs
@@ -1,6 +1,6 @@
 use crate::mcp::{AgentTypeImplementers, CompiledMcp};
 use golem_common::base_model::domain_registration::Domain;
-use golem_common::model::agent::AgentTypeName;
+use golem_common::model::agent::{AgentTypeName, RegisteredAgentType};
 
 impl From<CompiledMcp> for golem_api_grpc::proto::golem::mcp::CompiledMcp {
     fn from(value: CompiledMcp) -> Self {
@@ -23,6 +23,11 @@ impl From<CompiledMcp> for golem_api_grpc::proto::golem::mcp::CompiledMcp {
                 })
                 .collect(),
             security_scheme: value.security_scheme.map(|s| s.into()),
+            registered_agent_types: value
+                .registered_agent_types
+                .into_iter()
+                .map(|rat| rat.into())
+                .collect(),
         }
     }
 }
@@ -50,6 +55,12 @@ impl TryFrom<golem_api_grpc::proto::golem::mcp::CompiledMcp> for CompiledMcp {
             })
             .collect::<Result<_, String>>()?;
 
+        let registered_agent_types: Vec<RegisteredAgentType> = value
+            .registered_agent_types
+            .into_iter()
+            .map(|rat| rat.try_into())
+            .collect::<Result<_, String>>()?;
+
         Ok(Self {
             account_id: value
                 .account_id
@@ -72,6 +83,7 @@ impl TryFrom<golem_api_grpc::proto::golem::mcp::CompiledMcp> for CompiledMcp {
                 .map(|s| s.try_into())
                 .transpose()
                 .map_err(|e| format!("Invalid security_scheme: {}", e))?,
+            registered_agent_types,
         })
     }
 }
diff --git a/golem-worker-service/src/mcp/agent_mcp_server.rs b/golem-worker-service/src/mcp/agent_mcp_server.rs
@@ -114,86 +114,69 @@ pub async fn get_agent_capabilities(
     let account_id = compiled_mcp.account_id;
     let environment_id = compiled_mcp.environment_id;
 
-    let agent_types = compiled_mcp.agent_types();
-
     tracing::info!(
-        "Found {} agent types for domain {}: {:?}",
-        agent_types.len(),
+        "Found {} registered agent types for domain {}: {:?}",
+        compiled_mcp.registered_agent_types.len(),
         domain.0,
-        agent_types
+        compiled_mcp
+            .registered_agent_types
             .iter()
-            .map(|at| at.0.clone())
+            .map(|rat| rat.agent_type.type_name.0.clone())
             .collect::<Vec<_>>()
     );
 
-    for agent_type_name in &agent_types {
-        match mcp_definition_lookup
-            .resolve_agent_type(domain, agent_type_name)
-            .await
-        {
-            Ok(registered_agent_type) => {
-                tracing::debug!(
-                    "Resolved agent type {} for domain {}: implemented by component {}, methods: {:?}",
-                    agent_type_name.0,
-                    domain.0,
-                    registered_agent_type.implemented_by.component_id.0,
-                    registered_agent_type
-                        .agent_type
-                        .methods
-                        .iter()
-                        .map(|m| m.name.clone())
-                        .collect::<Vec<_>>()
-                );
+    for registered_agent_type in &compiled_mcp.registered_agent_types {
+        tracing::debug!(
+            "Processing agent type {} for domain {}: implemented by component {}, methods: {:?}",
+            registered_agent_type.agent_type.type_name.0,
+            domain.0,
+            registered_agent_type.implemented_by.component_id.0,
+            registered_agent_type
+                .agent_type
+                .methods
+                .iter()
+                .map(|m| m.name.clone())
+                .collect::<Vec<_>>()
+        );
+
+        let agent_type = &registered_agent_type.agent_type;
+        let component_id = registered_agent_type.implemented_by.component_id;
+
+        if let Some(prompt_hint) = &agent_type.constructor.prompt_hint {
+            prompts.push(AgentMcpPrompt::from_constructor_hint(
+                &agent_type.type_name,
+                &agent_type.description,
+                prompt_hint,
+            ));
+        }
 
-                let agent_type = &registered_agent_type.agent_type;
+        for method in &agent_type.methods {
+            if let Some(prompt_hint) = &method.prompt_hint {
+                prompts.push(AgentMcpPrompt::from_method_hint(
+                    &agent_type.type_name,
+                    method,
+                    &agent_type.constructor,
+                    prompt_hint,
+                ));
+            }
 
-                let component_id = registered_agent_type.implemented_by.component_id;
+            let agent_method_mcp = McpAgentCapability::from_agent_method(
+                &account_id,
+                &environment_id,
+                &agent_type.type_name,
+                method,
+                &agent_type.constructor,
+                component_id,
+            );
 
-                if let Some(prompt_hint) = &agent_type.constructor.prompt_hint {
-                    prompts.push(AgentMcpPrompt::from_constructor_hint(
-                        &agent_type.type_name,
-                        &agent_type.description,
-                        prompt_hint,
-                    ));
+            match agent_method_mcp {
+                McpAgentCapability::Tool(agent_mcp_tool) => {
+                    tools.push(*agent_mcp_tool);
                 }
-
-                for method in &agent_type.methods {
-                    if let Some(prompt_hint) = &method.prompt_hint {
-                        prompts.push(AgentMcpPrompt::from_method_hint(
-                            &agent_type.type_name,
-                            method,
-                            &agent_type.constructor,
-                            prompt_hint,
-                        ));
-                    }
-
-                    let agent_method_mcp = McpAgentCapability::from_agent_method(
-                        &account_id,
-                        &environment_id,
-                        &agent_type.type_name,
-                        method,
-                        &agent_type.constructor,
-                        component_id,
-                    );
-
-                    match agent_method_mcp {
-                        McpAgentCapability::Tool(agent_mcp_tool) => {
-                            tools.push(*agent_mcp_tool);
-                        }
-                        McpAgentCapability::Resource(agent_mcp_resource) => {
-                            resources.push(*agent_mcp_resource);
-                        }
-                    }
+                McpAgentCapability::Resource(agent_mcp_resource) => {
+                    resources.push(*agent_mcp_resource);
                 }
             }
-            Err(e) => {
-                tracing::error!(
-                    "Failed to resolve agent type {} for domain {}: {}",
-                    agent_type_name.0,
-                    domain.0,
-                    e
-                );
-            }
         }
     }
 
diff --git a/golem-worker-service/src/mcp/mcp_capabilities_lookup.rs b/golem-worker-service/src/mcp/mcp_capabilities_lookup.rs

Original file line number	Diff line number	Diff line change
`@@ -14,4 +14,5 @@ message CompiledMcp {`
`14`	`14`	`string domain = 4;`
`15`	`15`	`map<string, golem.registry.RegisteredAgentTypeImplementer> agent_type_implementers = 5;`
`16`	`16`	`golem.customapi.SecuritySchemeDetails security_scheme = 6;`
	`17`	`+ repeated golem.registry.RegisteredAgentType registered_agent_types = 7;`
`17`	`18`	`}`
Original file line number	Diff line number	Diff line change
`@@ -599,6 +599,7 @@ impl TryFrom<DeploymentCompiledMcpRecord> for CompiledMcp {`
`599`	`599`	`domain: Domain(value.domain),`
`600`	`600`	`agent_type_implementers: mcp_data.implementers,`
`601`	`601`	`security_scheme: None,`
	`602`	`+ registered_agent_types: Vec::new(),`
`602`	`603`	`})`
`603`	`604`	`}`
`604`	`605`	`}`
Original file line number	Diff line number	Diff line change
`@@ -297,6 +297,7 @@ impl DeploymentContext {`
`297`	`297`	`domain: domain.clone(),`
`298`	`298`	`agent_type_implementers,`
`299`	`299`	`security_scheme,`
	`300`	`+ registered_agent_types: Vec::new(),`
`300`	`301`	`};`
`301`	`302`	`all_compiled_mcps.push(compiled_mcp);`
`302`	`303`	`}`