diff --git a/backend/Cargo.lock b/backend/Cargo.lock index b7b69cb93f..8d1c8a70e1 100644 --- a/backend/Cargo.lock +++ b/backend/Cargo.lock @@ -10852,9 +10852,9 @@ dependencies = [ [[package]] name = "rmcp" -version = "0.6.4" +version = "0.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "41ab0892f4938752b34ae47cb53910b1b0921e55e77ddb6e44df666cab17939f" +checksum = "6f35acda8f89fca5fd8c96cae3c6d5b4c38ea0072df4c8030915f3b5ff469c1c" dependencies = [ "base64 0.22.1", "bytes", @@ -10866,6 +10866,7 @@ dependencies = [ "paste", "pin-project-lite", "rand 0.9.0", + "reqwest 0.12.24", "rmcp-macros", "schemars 1.0.4", "serde", @@ -10882,9 +10883,9 @@ dependencies = [ [[package]] name = "rmcp-macros" -version = "0.6.4" +version = "0.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1827cd98dab34cade0513243c6fe0351f0f0b2c9d6825460bcf45b42804bdda0" +checksum = "c9f1d5220aaa23b79c3d02e18f7a554403b3ccea544bbb6c69d6bcb3e854a274" dependencies = [ "darling 0.21.3", "proc-macro2", @@ -15399,6 +15400,7 @@ dependencies = [ "reqwest 0.12.24", "reqwest-middleware", "reqwest-retry", + "rmcp", "semver 1.0.27", "serde", "serde_json", @@ -15811,6 +15813,7 @@ dependencies = [ "regex", "reqwest 0.12.24", "reqwest-middleware", + "rmcp", "rust_decimal", "serde", "serde_json", diff --git a/backend/windmill-api/Cargo.toml b/backend/windmill-api/Cargo.toml index 76d05a7188..34cb7adfe9 100644 --- a/backend/windmill-api/Cargo.toml +++ b/backend/windmill-api/Cargo.toml @@ -40,7 +40,7 @@ mcp = ["dep:rmcp"] python = [] [dependencies] -rmcp = { version = "0.6.4", features=["transport-streamable-http-server", "transport-streamable-http-server-session", "transport-worker"], optional = true } +rmcp = { version = "0.8.1", features=["transport-streamable-http-server", "transport-streamable-http-server-session", "transport-worker"], optional = true } windmill-queue.workspace = true windmill-common = { workspace = true, default-features = false } windmill-audit.workspace = true diff --git a/backend/windmill-api/openapi.yaml b/backend/windmill-api/openapi.yaml index 591f8992a0..504d175150 100644 --- a/backend/windmill-api/openapi.yaml +++ b/backend/windmill-api/openapi.yaml @@ -4297,6 +4297,35 @@ paths: - path - value + /w/{workspace}/resources/mcp_tools/{path}: + get: + summary: get MCP tools from resource + operationId: getMcpTools + tags: + - resource + parameters: + - $ref: "#/components/parameters/WorkspaceId" + - $ref: "#/components/parameters/Path" + responses: + "200": + description: list of MCP tools + content: + application/json: + schema: + type: array + items: + type: object + properties: + name: + type: string + description: + type: string + parameters: + type: object + required: + - name + - parameters + /w/{workspace}/resources/list_names/{name}: get: summary: list resource names diff --git a/backend/windmill-api/src/resources.rs b/backend/windmill-api/src/resources.rs index 192caa5deb..7e6e6cb6f1 100644 --- a/backend/windmill-api/src/resources.rs +++ b/backend/windmill-api/src/resources.rs @@ -73,6 +73,7 @@ pub fn workspaced_service() -> Router { get(file_resource_ext_to_resource_type), ) .route("/type/create", post(create_resource_type)) + .route("/mcp_tools/*path", get(get_mcp_tools)) } pub fn public_service() -> Router { @@ -1392,6 +1393,64 @@ where Ok(resource) } +/// Get list of tools from an MCP resource +async fn get_mcp_tools( + authed: ApiAuthed, + Extension(db): Extension, + Extension(user_db): Extension, + Path((w_id, path)): Path<(String, StripPath)>, +) -> JsonResult> { + let path = path.to_path(); + check_scopes(&authed, || format!("resources:read:{}", path))?; + + let mut tx = user_db.begin(&authed).await?; + + // Fetch the MCP resource from database + let resource_value_o = sqlx::query_scalar!( + "SELECT value as \"value: sqlx::types::Json>\" FROM resource WHERE path = $1 AND workspace_id = $2", + &path, + &w_id + ) + .fetch_optional(&mut *tx) + .await?; + + tx.commit().await?; + + if resource_value_o.is_none() { + explain_resource_perm_error(&path, &w_id, &db, &authed).await?; + } + + let resource_value = not_found_if_none(resource_value_o, "Resource", path)? + .ok_or_else(|| Error::BadRequest(format!("Empty resource value for {}", path)))?; + + // Parse MCP resource + let mcp_resource = + serde_json::from_str::(resource_value.0.get()) + .map_err(|e| Error::BadRequest(format!("Failed to parse MCP resource: {}", e)))?; + + // Create MCP client connection + let client = windmill_common::mcp_client::McpClient::from_resource(mcp_resource, &db, &w_id) + .await + .map_err(|e| Error::ExecutionErr(format!("Failed to connect to MCP server: {}", e)))?; + + // Get raw MCP tools and convert to JSON + let tools: Vec = client + .available_tools() + .iter() + .map(|tool| { + serde_json::to_value(tool) + .map_err(|e| Error::ExecutionErr(format!("Failed to serialize MCP tool: {}", e))) + }) + .collect::>>()?; + + // Gracefully shutdown the client + if let Err(e) = client.shutdown().await { + tracing::warn!("Failed to shutdown MCP client: {}", e); + } + + Ok(Json(tools)) +} + #[derive(Deserialize, Serialize)] struct GitRepositoryResource { url: String, diff --git a/backend/windmill-common/Cargo.toml b/backend/windmill-common/Cargo.toml index c24b0d503c..8317940079 100644 --- a/backend/windmill-common/Cargo.toml +++ b/backend/windmill-common/Cargo.toml @@ -97,6 +97,7 @@ tempfile.workspace = true systemstat.workspace = true size.workspace = true globset.workspace = true +rmcp = { version = "0.8.1", features = ["client", "transport-streamable-http-client", "transport-streamable-http-client-reqwest"] } opentelemetry-semantic-conventions = { workspace = true, optional = true } opentelemetry-otlp = { workspace = true, optional = true } diff --git a/backend/windmill-common/src/flow_status.rs b/backend/windmill-common/src/flow_status.rs index f3dade1c0e..44da5f6639 100644 --- a/backend/windmill-common/src/flow_status.rs +++ b/backend/windmill-common/src/flow_status.rs @@ -187,7 +187,18 @@ struct UntaggedFlowStatusModule { #[derive(Serialize, Deserialize, Debug, Clone)] #[serde(tag = "type", rename_all = "snake_case")] pub enum AgentAction { - ToolCall { job_id: uuid::Uuid, function_name: String, module_id: String }, + ToolCall { + job_id: uuid::Uuid, + function_name: String, + module_id: String, + }, + McpToolCall { + call_id: uuid::Uuid, + function_name: String, + resource_path: String, + #[serde(skip_serializing_if = "Option::is_none")] + arguments: Option, + }, Message {}, } diff --git a/backend/windmill-common/src/flows.rs b/backend/windmill-common/src/flows.rs index 810368f901..9d554db7e1 100644 --- a/backend/windmill-common/src/flows.rs +++ b/backend/windmill-common/src/flows.rs @@ -176,9 +176,19 @@ impl FlowValue { | Flow { .. } | FlowScript { .. } | Identity) => cb(&s, &module.id)?, - ForloopFlow { modules, .. } - | WhileloopFlow { modules, .. } - | AIAgent { tools: modules, .. } => Self::traverse_leafs(&modules, cb)?, + ForloopFlow { modules, .. } | WhileloopFlow { modules, .. } => { + Self::traverse_leafs(&modules, cb)? + } + AIAgent { tools, .. } => { + for tool in tools { + match &tool.value { + ToolValue::FlowModule(module_value) => cb(module_value, &tool.id)?, + ToolValue::Mcp(_) => { + // MCP tools don't have a FlowModuleValue to traverse + } + } + } + } BranchOne { branches, .. } | BranchAll { branches, .. } => { for branch in branches { Self::traverse_leafs(&branch.modules, cb)?; @@ -601,6 +611,96 @@ pub struct Branch { pub parallel: bool, } +// Tool types for AI Agent +#[derive(Serialize, Debug, Clone, Deserialize)] +pub struct AgentTool { + pub id: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub summary: Option, + pub value: ToolValue, +} + +// Convert FlowModule -> AgentTool +impl From for AgentTool { + fn from(flow_module: FlowModule) -> Self { + let module_value = serde_json::from_str::(flow_module.value.get()) + .unwrap_or(FlowModuleValue::Identity); + + AgentTool { + id: flow_module.id, + summary: flow_module.summary, + value: ToolValue::FlowModule(module_value), + } + } +} + +// Convert AgentTool -> FlowModule (only for FlowModule type tools) +impl From<&AgentTool> for Option { + fn from(tool: &AgentTool) -> Self { + match &tool.value { + ToolValue::FlowModule(module_value) => Some(FlowModule { + id: tool.id.clone(), + value: to_raw_value(module_value), + summary: tool.summary.clone(), + ..Default::default() + }), + ToolValue::Mcp(_) => None, // MCP tools can't be converted to FlowModule + } + } +} + +#[derive(Serialize, Debug, Clone)] +#[serde(tag = "tool_type", rename_all = "lowercase")] +pub enum ToolValue { + FlowModule(FlowModuleValue), + Mcp(McpToolValue), +} + +// Custom deserializer for backward compatibility with old flows +impl<'de> Deserialize<'de> for ToolValue { + fn deserialize(deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + use serde::de::Error; + + let content = serde_json::Value::deserialize(deserializer)?; + + // First, try to deserialize as the new tagged format (with tool_type field) + #[derive(Deserialize)] + #[serde(tag = "tool_type", rename_all = "lowercase")] + enum TaggedToolValue { + FlowModule(FlowModuleValue), + Mcp(McpToolValue), + } + + if let Ok(tagged) = TaggedToolValue::deserialize(&content) { + return Ok(match tagged { + TaggedToolValue::FlowModule(v) => ToolValue::FlowModule(v), + TaggedToolValue::Mcp(v) => ToolValue::Mcp(v), + }); + } + + // Fall back to legacy format (direct FlowModuleValue without tool_type) + FlowModuleValue::deserialize(&content) + .map(ToolValue::FlowModule) + .map_err(|_| { + D::Error::custom( + "expected ToolValue with tool_type field or legacy FlowModuleValue", + ) + }) + } +} + +#[derive(Serialize, Deserialize, Debug, Clone)] +pub struct McpToolValue { + pub resource_path: String, + #[serde(default)] + pub include_tools: Vec, + #[serde(default)] + pub exclude_tools: Vec, +} + #[derive(Serialize, Debug, Clone)] #[serde( tag = "type", @@ -725,7 +825,7 @@ pub enum FlowModuleValue { // AI agent node AIAgent { input_transforms: HashMap, - tools: Vec, + tools: Vec, }, } @@ -762,7 +862,7 @@ struct UntaggedFlowModuleValue { default_node: Option, modules_node: Option, assets: Option>, - tools: Option>, + tools: Option>, pass_flow_input_directly: Option, } diff --git a/backend/windmill-common/src/lib.rs b/backend/windmill-common/src/lib.rs index 4ff4e18723..7d75f36a91 100644 --- a/backend/windmill-common/src/lib.rs +++ b/backend/windmill-common/src/lib.rs @@ -54,8 +54,12 @@ pub mod job_s3_helpers_ee; #[cfg(feature = "parquet")] pub mod job_s3_helpers_oss; +#[cfg(feature = "private")] +pub mod git_sync_ee; +pub mod git_sync_oss; pub mod jobs; pub mod jwt; +pub mod mcp_client; pub mod more_serde; pub mod oauth2; #[cfg(all(feature = "enterprise", feature = "openidconnect", feature = "private"))] @@ -87,9 +91,6 @@ pub mod variables; pub mod worker; pub mod worker_group_job_stats; pub mod workspaces; -#[cfg(feature = "private")] -pub mod git_sync_ee; -pub mod git_sync_oss; pub const DEFAULT_MAX_CONNECTIONS_SERVER: u32 = 50; pub const DEFAULT_MAX_CONNECTIONS_WORKER: u32 = 5; diff --git a/backend/windmill-common/src/mcp_client.rs b/backend/windmill-common/src/mcp_client.rs new file mode 100644 index 0000000000..e12aebe5bd --- /dev/null +++ b/backend/windmill-common/src/mcp_client.rs @@ -0,0 +1,228 @@ +use crate::variables::get_secret_value_as_admin; +use crate::DB; +use anyhow::{Context, Result}; +use reqwest::header::{HeaderMap, HeaderName, HeaderValue}; +use serde_json::{json, Value}; +use std::str::FromStr; + +use rmcp::model::Tool as McpTool; +use rmcp::{ + model::{ + CallToolRequestParam, ClientCapabilities, ClientInfo, Implementation, + InitializeRequestParam, + }, + service::RunningService, + transport::{ + streamable_http_client::StreamableHttpClientTransportConfig, StreamableHttpClientTransport, + }, + RoleClient, ServiceExt, +}; + +use std::collections::HashMap; + +use serde::{Deserialize, Serialize}; + +/// MCP server resource configuration +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct McpResource { + /// Name of the MCP resource (used for prefixing tools) + pub name: String, + /// HTTP URL for the MCP server endpoint + pub url: String, + /// Optional token for authentication + #[serde(skip_serializing_if = "Option::is_none")] + pub token: Option, + /// Optional headers + #[serde(skip_serializing_if = "Option::is_none")] + pub headers: Option>, +} + +/// Metadata for tracking MCP tool sources +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct McpToolSource { + /// Name of the MCP resource this tool comes from + pub name: String, + /// Original tool name in the MCP server + pub tool_name: String, + /// Path of the MCP resource + pub resource_path: String, +} + +/// MCP client for communicating with external MCP servers +pub struct McpClient { + /// The underlying rmcp client + client: RunningService, + /// Cached list of available tools from the server + available_tools: Vec, +} + +impl McpClient { + /// Create a new MCP client from a resource configuration + pub async fn from_resource(resource: McpResource, db: &DB, w_id: &str) -> Result { + // Build custom reqwest client with headers if provided + let mut headers = HeaderMap::new(); + if let Some(token_path) = &resource.token { + if !token_path.trim().is_empty() { + let value = + get_secret_value_as_admin(db, w_id, token_path.trim_start_matches("$var:")) + .await?; + headers.insert( + HeaderName::from_static("authorization"), + HeaderValue::from_str(format!("Bearer {}", value).as_str())?, + ); + } + } + if let Some(resource_headers) = &resource.headers { + for (key, value) in resource_headers { + match (HeaderName::from_str(key), HeaderValue::from_str(value)) { + (Ok(name), Ok(value)) => { + headers.insert(name, value); + } + _ => { + tracing::warn!("Invalid header: {}={}", key, value); + } + } + } + } + + let reqwest_client = reqwest::Client::builder() + .default_headers(headers) + .build() + .context("Failed to build HTTP client")?; + + // Create the HTTP transport with custom client + let config = StreamableHttpClientTransportConfig::with_uri(resource.url.as_str()); + let transport = StreamableHttpClientTransport::with_client(reqwest_client, config); + + // Set up client info + let client_info = ClientInfo { + protocol_version: Default::default(), + capabilities: ClientCapabilities::default(), + client_info: Implementation { + name: "windmill-ai-agent".to_string(), + title: Some("Windmill AI Agent".to_string()), + version: env!("CARGO_PKG_VERSION").to_string(), + website_url: None, + icons: None, + }, + }; + + // Initialize the connection + let client = client_info + .serve(transport) + .await + .context("Failed to connect to MCP server")?; + + // Immediately fetch available tools + let available_tools = client + .list_tools(Default::default()) + .await + .context("Failed to list tools from MCP server")? + .tools; + + Ok(Self { client, available_tools }) + } + + /// Get the list of available tools from the MCP server + pub fn available_tools(&self) -> &[McpTool] { + &self.available_tools + } + + /// Call a tool on the MCP server, with openai-style arguments + pub async fn call_tool(&self, name: &str, arguments: &str) -> Result { + // Convert OpenAI-style arguments to MCP format + let mcp_args = + Self::openai_args_to_mcp_args(arguments).context("Failed to parse tool arguments")?; + + let result = self + .client + .call_tool(CallToolRequestParam { name: name.to_string().into(), arguments: mcp_args }) + .await + .context(format!("Failed to call MCP tool: {}", name))?; + + // Convert the result to a JSON value + // MCP tools return ToolResult which contains content array + let result_json = + serde_json::to_value(&result).context("Failed to serialize MCP tool result")?; + + Ok(result_json) + } + + /// Close the connection + pub async fn shutdown(self) -> Result<()> { + self.client.cancel().await?; + Ok(()) + } + + /// Fix array schemas to ensure they have the required 'items' property + /// OpenAI requires all array types to have an 'items' field. MCP servers may + /// return schemas without this field, so we add a default. + pub fn fix_array_schemas(schema: &mut Value) { + if let Value::Object(obj) = schema { + // Check if this is an array type + if let Some(type_val) = obj.get("type") { + let is_array = match type_val { + Value::String(s) => s == "array", + Value::Array(arr) => arr.iter().any(|v| v.as_str() == Some("array")), + _ => false, + }; + + // If it's an array and missing 'items', add a default + if is_array && !obj.contains_key("items") { + obj.insert("items".to_string(), json!({})); + } + } + + // Recursively fix nested schemas + if let Some(Value::Object(props)) = obj.get_mut("properties") { + for value in props.values_mut() { + Self::fix_array_schemas(value); + } + } + + // Fix items if present (for nested arrays) + if let Some(items) = obj.get_mut("items") { + Self::fix_array_schemas(items); + } + + // Fix oneOf, anyOf, allOf schemas + for key in &["oneOf", "anyOf", "allOf"] { + if let Some(Value::Array(schemas)) = obj.get_mut(*key) { + for schema in schemas { + Self::fix_array_schemas(schema); + } + } + } + + // Fix additionalProperties if it's a schema + if let Some(additional) = obj.get_mut("additionalProperties") { + if additional.is_object() { + Self::fix_array_schemas(additional); + } + } + } + } + + /// Convert OpenAI-style tool call arguments to MCP format + /// OpenAI sends arguments as a JSON string, MCP expects a Map + fn openai_args_to_mcp_args( + args_str: &str, + ) -> Result>> { + if args_str.trim().is_empty() { + return Ok(None); + } + + let args_value: serde_json::Value = + serde_json::from_str(args_str).context("Failed to parse tool call arguments")?; + + match args_value { + serde_json::Value::Object(map) => Ok(Some(map)), + serde_json::Value::Null => Ok(None), + _ => Ok(Some( + vec![("value".to_string(), args_value)] + .into_iter() + .collect(), + )), + } + } +} diff --git a/backend/windmill-worker/Cargo.toml b/backend/windmill-worker/Cargo.toml index d249a78319..f36d2c5d58 100644 --- a/backend/windmill-worker/Cargo.toml +++ b/backend/windmill-worker/Cargo.toml @@ -57,6 +57,7 @@ windmill-parser-sql.workspace = true windmill-parser-graphql.workspace = true windmill-parser-php = { workspace = true, optional = true } windmill-git-sync.workspace = true +rmcp = { version = "0.8.1", features = ["client", "transport-streamable-http-client", "transport-streamable-http-client-reqwest"] } flume.workspace = true sqlx.workspace = true uuid.workspace = true diff --git a/backend/windmill-worker/src/ai/mod.rs b/backend/windmill-worker/src/ai/mod.rs index 4d11f100a5..4ad67b9e6e 100644 --- a/backend/windmill-worker/src/ai/mod.rs +++ b/backend/windmill-worker/src/ai/mod.rs @@ -5,4 +5,6 @@ pub mod image_handler; pub mod providers; pub mod query_builder; pub mod sse; +pub mod tools; pub mod types; +pub mod utils; diff --git a/backend/windmill-worker/src/ai/tools.rs b/backend/windmill-worker/src/ai/tools.rs new file mode 100644 index 0000000000..1f52565182 --- /dev/null +++ b/backend/windmill-worker/src/ai/tools.rs @@ -0,0 +1,697 @@ +use crate::ai::providers::openai::OpenAIToolCall; +use crate::ai::query_builder::StreamEventProcessor; +use crate::ai::types::*; +use crate::ai::utils::{ + add_message_to_conversation, execute_mcp_tool, get_flow_chat_settings, get_step_name_from_flow, + update_flow_status_module_with_actions, update_flow_status_module_with_actions_success, + FlowChatSettings, +}; +use crate::common::{error_to_value, OccupancyMetrics}; +use crate::result_processor::handle_non_flow_job_error; +use crate::worker_flow::{raw_script_to_payload, script_to_payload, JobPayloadWithTag}; +use crate::{ + create_job_dir, handle_queued_job, JobCompletedReceiver, JobCompletedSender, SendResult, + SendResultPayload, +}; +use anyhow::Context; +use serde_json::value::RawValue; +use std::{collections::HashMap, sync::Arc}; +use uuid::Uuid; +use windmill_common::jobs::JobPayload; +use windmill_common::mcp_client::{McpClient, McpToolSource}; +use windmill_common::{ + client::AuthedClient, + db::DB, + error::{to_anyhow, Error}, + flow_conversations::MessageType, + flow_status::AgentAction, + flows::FlowModuleValue, + worker::Connection, +}; +use windmill_queue::{ + get_mini_pulled_job, push, JobCompleted, MiniPulledJob, PushArgs, PushIsolationLevel, +}; + +/// Context for tool execution containing all required references and state +pub struct ToolExecutionContext<'a> { + // Database & connections + pub db: &'a DB, + pub conn: &'a Connection, + + // Job context + pub job: &'a MiniPulledJob, + pub parent_job: &'a Uuid, + pub summary: &'a Option<&'a str>, + + // Execution parameters + pub client: &'a AuthedClient, + pub worker_dir: &'a str, + pub base_internal_url: &'a str, + pub worker_name: &'a str, + pub hostname: &'a str, + + // Runtime state + pub occupancy_metrics: &'a mut OccupancyMetrics, + pub job_completed_tx: &'a JobCompletedSender, + pub killpill_rx: &'a mut tokio::sync::broadcast::Receiver<()>, + + // Optional streaming & chat + pub stream_event_processor: Option<&'a StreamEventProcessor>, + pub chat_settings: &'a mut Option, +} + +/// Execute all tool calls from an AI response +pub async fn execute_tool_calls( + mut ctx: ToolExecutionContext<'_>, + tool_calls: &[OpenAIToolCall], + tools: &[Tool], + mcp_clients: &HashMap>, + actions: &mut Vec, + final_events_str: &mut String, + structured_output_tool_name: &Option, +) -> Result<(Vec, Option, bool), Error> { + let mut messages = Vec::new(); + let mut used_structured_output_tool = false; + let mut final_content = None; + + for tool_call in tool_calls.iter() { + // Stream tool call progress + if let Some(stream_event_processor) = ctx.stream_event_processor { + let event = StreamingEvent::ToolExecution { + call_id: tool_call.id.clone(), + function_name: tool_call.function.name.clone(), + }; + stream_event_processor.send(event, final_events_str).await?; + } + + // Check if this is the structured output tool + if structured_output_tool_name + .as_ref() + .map_or(false, |name| tool_call.function.name == *name) + { + used_structured_output_tool = true; + messages.push(OpenAIMessage { + role: "tool".to_string(), + content: Some(OpenAIContent::Text( + "Successfully ran structured_output tool".to_string(), + )), + tool_call_id: Some(tool_call.id.clone()), + ..Default::default() + }); + messages.push(OpenAIMessage { + role: "assistant".to_string(), + content: Some(OpenAIContent::Text(tool_call.function.arguments.clone())), + agent_action: Some(AgentAction::Message {}), + ..Default::default() + }); + final_content = Some(OpenAIContent::Text(tool_call.function.arguments.clone())); + break; + } + + let tool = tools + .iter() + .find(|t| t.def.function.name == tool_call.function.name); + + if let Some(tool) = tool { + // Check if this is an MCP tool + if let Some(mcp_source) = &tool.mcp_source { + execute_mcp_tool_call( + &mut ctx, + tool_call, + mcp_clients, + mcp_source, + actions, + &mut messages, + final_events_str, + ) + .await?; + } else if tool.module.is_some() { + execute_windmill_tool( + &mut ctx, + tool_call, + tool, + actions, + &mut messages, + final_events_str, + ) + .await?; + } else { + return Err(Error::internal_err(format!( + "Tool type not supported: {}", + tool_call.function.name + ))); + } + } else { + return Err(Error::internal_err(format!( + "Tool not found: {}", + tool_call.function.name + ))); + } + } + + Ok((messages, final_content, used_structured_output_tool)) +} + +/// Execute an MCP tool call +async fn execute_mcp_tool_call( + ctx: &mut ToolExecutionContext<'_>, + tool_call: &OpenAIToolCall, + mcp_clients: &HashMap>, + mcp_source: &McpToolSource, + actions: &mut Vec, + messages: &mut Vec, + final_events_str: &mut String, +) -> Result<(), Error> { + let tool_result = + execute_mcp_tool(mcp_clients, mcp_source, &tool_call.function.arguments).await; + + let call_id = ulid::Ulid::new().into(); + let resource_path = &mcp_source.resource_path; + let tool_name = &tool_call.function.name; + let arguments = serde_json::from_str(&tool_call.function.arguments).ok(); + + actions.push(AgentAction::McpToolCall { + call_id, + function_name: tool_name.clone(), + resource_path: resource_path.clone(), + arguments: arguments.clone(), + }); + + match tool_result { + Ok(result) => { + let result_str = + serde_json::to_string_pretty(&result).unwrap_or_else(|_| result.to_string()); + + messages.push(OpenAIMessage { + role: "tool".to_string(), + content: Some(OpenAIContent::Text(result_str.clone())), + tool_call_id: Some(tool_call.id.clone()), + agent_action: Some(AgentAction::McpToolCall { + call_id, + function_name: tool_name.clone(), + resource_path: resource_path.clone(), + arguments: arguments.clone(), + }), + ..Default::default() + }); + + // Stream tool result + if let Some(stream_event_processor) = ctx.stream_event_processor { + let event = StreamingEvent::ToolResult { + call_id: tool_call.id.clone(), + function_name: tool_call.function.name.clone(), + result: result.to_string(), + success: true, + }; + stream_event_processor.send(event, final_events_str).await?; + } + + // Add tool message to conversation if chat_input_enabled + let content = format!("Used {} tool", tool_call.function.name); + add_tool_message_to_chat(ctx, None, &content, true).await; + } + Err(e) => { + let error_msg = format!("MCP tool error: {}", e); + tracing::error!("{}", error_msg); + + messages.push(OpenAIMessage { + role: "tool".to_string(), + content: Some(OpenAIContent::Text(error_msg.clone())), + tool_call_id: Some(tool_call.id.clone()), + agent_action: Some(AgentAction::McpToolCall { + call_id, + function_name: tool_name.clone(), + resource_path: resource_path.clone(), + arguments: arguments.clone(), + }), + ..Default::default() + }); + + // Stream tool error + if let Some(stream_event_processor) = ctx.stream_event_processor { + let event = StreamingEvent::ToolResult { + call_id: tool_call.id.clone(), + function_name: tool_name.clone(), + result: error_msg.clone(), + success: false, + }; + stream_event_processor.send(event, final_events_str).await?; + } + + // Add tool message to conversation if chat_input_enabled + add_tool_message_to_chat(ctx, None, &error_msg, false).await; + } + } + + Ok(()) +} + +/// Execute a Windmill tool (script or flow) +async fn execute_windmill_tool( + ctx: &mut ToolExecutionContext<'_>, + tool_call: &OpenAIToolCall, + tool: &Tool, + actions: &mut Vec, + messages: &mut Vec, + final_events_str: &mut String, +) -> Result<(), Error> { + // Regular Windmill tools must have a module + let tool_module = tool.module.as_ref().ok_or_else(|| { + Error::internal_err(format!( + "Tool {} has no module (MCP tools should be handled above)", + tool_call.function.name + )) + })?; + + let job_id = ulid::Ulid::new().into(); + actions.push(AgentAction::ToolCall { + job_id, + function_name: tool_call.function.name.clone(), + module_id: tool_module.id.clone(), + }); + + update_flow_status_module_with_actions(ctx.db, ctx.parent_job, actions).await?; + + let raw_tool_call_args = if tool_call.function.arguments.is_empty() { + "{}".to_string() + } else { + tool_call.function.arguments.clone() + }; + + let tool_call_args = serde_json::from_str::>>( + &raw_tool_call_args, + ) + .with_context(|| { + format!( + "Failed to parse tool call arguments for tool call {}: {}", + tool_call.function.name, tool_call.function.arguments + ) + })?; + + let job_payload = match tool_module.get_value()? { + FlowModuleValue::Script { path: script_path, hash: script_hash, tag_override, .. } => { + script_to_payload( + script_hash, + script_path, + ctx.db, + ctx.job, + tool_module, + tag_override, + tool_module.apply_preprocessor, + ) + .await? + } + FlowModuleValue::RawScript { + path, + content, + language, + lock, + tag, + custom_concurrency_key, + concurrent_limit, + concurrency_time_window_s, + .. + } => { + let path = path + .unwrap_or_else(|| format!("{}/tools/{}", ctx.job.runnable_path(), tool_module.id)); + + raw_script_to_payload( + path, + content, + language, + lock, + custom_concurrency_key, + concurrent_limit, + concurrency_time_window_s, + tool_module, + tag, + tool_module.delete_after_use.unwrap_or(false), + ) + } + FlowModuleValue::FlowScript { + id, + language, + custom_concurrency_key, + concurrent_limit, + concurrency_time_window_s, + tag, + .. + } => { + let path = format!("{}/tools/{}", ctx.job.runnable_path(), tool_module.id); + + let payload = JobPayloadWithTag { + payload: JobPayload::FlowScript { + id, + language, + custom_concurrency_key: custom_concurrency_key.clone(), + concurrent_limit, + concurrency_time_window_s, + cache_ttl: tool_module.cache_ttl.map(|x| x as i32), + dedicated_worker: None, + path, + }, + tag: tag.clone(), + delete_after_use: tool_module.delete_after_use.unwrap_or(false), + timeout: None, + on_behalf_of: None, + }; + payload + } + _ => { + return Err(Error::internal_err(format!( + "Unsupported tool: {}", + tool_call.function.name + ))); + } + }; + + let mut tx = ctx.db.begin().await?; + + let job_perms = + windmill_common::auth::get_job_perms(&mut *tx, &ctx.job.id, &ctx.job.workspace_id) + .await? + .map(|x| x.into()); + + let (email, permissioned_as) = if let Some(on_behalf_of) = job_payload.on_behalf_of.as_ref() { + (&on_behalf_of.email, on_behalf_of.permissioned_as.clone()) + } else { + ( + &ctx.job.permissioned_as_email, + ctx.job.permissioned_as.to_owned(), + ) + }; + + let job_priority = tool_module.priority.or(ctx.job.priority); + + let tx = PushIsolationLevel::Transaction(tx); + let (uuid, tx) = push( + ctx.db, + tx, + &ctx.job.workspace_id, + job_payload.payload, + PushArgs { args: &tool_call_args, extra: None }, + &ctx.job.created_by, + email, + permissioned_as, + Some(&format!("job-span-{}", ctx.job.id)), + None, + ctx.job.schedule_path(), + Some(ctx.job.id), + None, + None, + Some(job_id), + false, + false, + None, + ctx.job.visible_to_owner, + Some(ctx.job.tag.clone()), + job_payload.timeout, + None, + job_priority, + job_perms.as_ref(), + true, + None, + None, + ) + .await?; + + tx.commit().await?; + + let tool_job = get_mini_pulled_job(ctx.db, &uuid).await?; + + let Some(tool_job) = tool_job else { + return Err(Error::internal_err("Tool job not found".to_string())); + }; + + let tool_job = Arc::new(tool_job); + + let (inner_job_completed_tx, inner_job_completed_rx) = JobCompletedSender::new(ctx.conn, 1); + + let inner_job_completed_rx = inner_job_completed_rx.expect( + "inner_job_completed_tx should be set as agent jobs are not supported on agent workers", + ); + + // Spawn handle_queued_job on separate task to prevent tokio stack overflow + // Clone everything needed for the spawned task + let tool_job_spawn = tool_job.clone(); + let conn_spawn = ctx.conn.clone(); + let client_spawn = ctx.client.clone(); + let hostname_spawn = ctx.hostname.to_string(); + let worker_name_spawn = ctx.worker_name.to_string(); + let worker_dir_spawn = ctx.worker_dir.to_string(); + let base_internal_url_spawn = ctx.base_internal_url.to_string(); + let inner_job_completed_tx_spawn = inner_job_completed_tx.clone(); + let mut occupancy_metrics_spawn = ctx.occupancy_metrics.clone(); + let mut killpill_rx_spawn = ctx.killpill_rx.resubscribe(); + + // Spawn on separate tokio task with fresh stack + let join_handle = tokio::task::spawn(async move { + #[cfg(feature = "benchmark")] + let mut bench_spawn = windmill_common::bench::BenchmarkIter::new(); + + let job_dir = create_job_dir(&worker_dir_spawn, tool_job_spawn.id).await; + + let result = handle_queued_job( + tool_job_spawn, + None, + None, + None, + None, + &conn_spawn, + &client_spawn, + &hostname_spawn, + &worker_name_spawn, + &worker_dir_spawn, + &job_dir, + None, + &base_internal_url_spawn, + inner_job_completed_tx_spawn, + &mut occupancy_metrics_spawn, + &mut killpill_rx_spawn, + None, + #[cfg(feature = "benchmark")] + &mut bench_spawn, + ) + .await; + + // Return both result and updated metrics + (result, occupancy_metrics_spawn) + }); + + // Await the spawned task + let (handle_result, updated_occupancy) = join_handle + .await + .map_err(|e| Error::internal_err(format!("Tool execution task failed: {}", e)))?; + + // Merge occupancy metrics back + ctx.occupancy_metrics.total_duration_of_running_jobs = + updated_occupancy.total_duration_of_running_jobs; + + // Continue with match on handle_result + match handle_result { + Err(err) => { + handle_tool_execution_error( + ctx, + tool_call, + tool_module, + &tool_job, + job_id, + err, + messages, + final_events_str, + ) + .await?; + } + Ok(success) => { + handle_tool_execution_success( + ctx, + tool_call, + tool_module, + job_id, + success, + inner_job_completed_rx, + messages, + final_events_str, + ) + .await?; + } + } + + Ok(()) +} + +/// Handle tool execution error +async fn handle_tool_execution_error( + ctx: &mut ToolExecutionContext<'_>, + tool_call: &OpenAIToolCall, + tool_module: &windmill_common::flows::FlowModule, + tool_job: &MiniPulledJob, + job_id: Uuid, + err: Error, + messages: &mut Vec, + final_events_str: &mut String, +) -> Result<(), Error> { + let err_string = format!("{}: {}", err.name(), err.to_string()); + let err_json = error_to_value(&err); + let _ = handle_non_flow_job_error( + ctx.db, + tool_job, + 0, + None, + err_string.clone(), + err_json, + ctx.worker_name, + ) + .await; + + let error_message = format!("Error running tool: {}", err_string); + messages.push(OpenAIMessage { + role: "tool".to_string(), + content: Some(OpenAIContent::Text(error_message.clone())), + tool_call_id: Some(tool_call.id.clone()), + agent_action: Some(AgentAction::ToolCall { + job_id, + function_name: tool_call.function.name.clone(), + module_id: tool_module.id.clone(), + }), + ..Default::default() + }); + + // Stream tool result (error case) + if let Some(stream_event_processor) = ctx.stream_event_processor { + let tool_result_event = StreamingEvent::ToolResult { + call_id: tool_call.id.clone(), + function_name: tool_call.function.name.clone(), + result: error_message.clone(), + success: false, + }; + stream_event_processor + .send(tool_result_event, final_events_str) + .await?; + } + + update_flow_status_module_with_actions_success(ctx.db, ctx.parent_job, false).await?; + + // Add tool message to conversation if chat_input_enabled (error case) + add_tool_message_to_chat(ctx, Some(job_id), &error_message, false).await; + + Ok(()) +} + +/// Handle tool execution success +async fn handle_tool_execution_success( + ctx: &mut ToolExecutionContext<'_>, + tool_call: &OpenAIToolCall, + tool_module: &windmill_common::flows::FlowModule, + job_id: Uuid, + success: bool, + inner_job_completed_rx: JobCompletedReceiver, + messages: &mut Vec, + final_events_str: &mut String, +) -> Result<(), Error> { + let send_result = inner_job_completed_rx.bounded_rx.try_recv().ok(); + + let result = if let Some(SendResult { + result: SendResultPayload::JobCompleted(JobCompleted { result, .. }), + .. + }) = send_result.as_ref() + { + ctx.job_completed_tx + .send(send_result.as_ref().unwrap().result.clone(), true) + .await + .map_err(to_anyhow)?; + result + } else { + if let Some(send_result) = send_result { + ctx.job_completed_tx + .send(send_result.result, true) + .await + .map_err(to_anyhow)?; + } + return Err(Error::internal_err( + "Tool job completed but no result".to_string(), + )); + }; + + messages.push(OpenAIMessage { + role: "tool".to_string(), + content: Some(OpenAIContent::Text(result.get().to_string())), + tool_call_id: Some(tool_call.id.clone()), + agent_action: Some(AgentAction::ToolCall { + job_id, + function_name: tool_call.function.name.clone(), + module_id: tool_module.id.clone(), + }), + ..Default::default() + }); + + // Stream tool result (success case) + if let Some(stream_event_processor) = ctx.stream_event_processor { + let tool_result_event = StreamingEvent::ToolResult { + call_id: tool_call.id.clone(), + function_name: tool_call.function.name.clone(), + result: result.get().to_string(), + success: true, + }; + stream_event_processor + .send(tool_result_event, final_events_str) + .await?; + } + + update_flow_status_module_with_actions_success(ctx.db, ctx.parent_job, success).await?; + + // Add tool message to conversation if chat_input_enabled + let content = if success { + format!("Used {} tool", tool_call.function.name) + } else { + format!("Error executing {}", tool_call.function.name) + }; + + add_tool_message_to_chat(ctx, Some(job_id), &content, success).await; + + Ok(()) +} + +/// Add tool message to conversation if chat is enabled +async fn add_tool_message_to_chat( + ctx: &mut ToolExecutionContext<'_>, + tool_job_id: Option, + content: &str, + success: bool, +) { + if ctx.chat_settings.is_none() { + *ctx.chat_settings = Some(get_flow_chat_settings(ctx.db, ctx.job).await); + } + + let chat_enabled = ctx + .chat_settings + .as_ref() + .map(|s| s.chat_input_enabled) + .unwrap_or(false); + + if chat_enabled { + if let Some(mid) = ctx.chat_settings.as_ref().and_then(|s| s.memory_id) { + let db_clone = ctx.db.clone(); + let step_name = + get_step_name_from_flow(ctx.summary.as_deref(), ctx.job.flow_step_id.as_deref()); + let content = content.to_string(); + + // Spawn task because we do not need to wait for the result + tokio::spawn(async move { + if let Err(e) = add_message_to_conversation( + &db_clone, + &mid, + tool_job_id, + &content, + MessageType::Tool, + &step_name, + success, + ) + .await + { + tracing::warn!("Failed to add tool message to conversation {}: {}", mid, e); + } + }); + } + } +} diff --git a/backend/windmill-worker/src/ai/types.rs b/backend/windmill-worker/src/ai/types.rs index aa440ce8f0..4a2de30926 100644 --- a/backend/windmill-worker/src/ai/types.rs +++ b/backend/windmill-worker/src/ai/types.rs @@ -1,3 +1,5 @@ +use crate::ai::providers::openai::OpenAIToolCall; +use windmill_common::mcp_client::McpToolSource; use serde::{Deserialize, Serialize}; use serde_json::value::RawValue; use std::collections::HashMap; @@ -7,10 +9,6 @@ use windmill_common::{ }; use windmill_parser::Typ; -use crate::ai::providers::openai::OpenAIToolCall; - -// Shared types used across multiple providers - #[derive(Serialize, Deserialize, Clone, Debug)] #[serde(tag = "type", rename_all = "snake_case")] pub enum ContentPart { @@ -88,9 +86,11 @@ pub struct ToolDef { pub function: ToolDefFunction, } +#[derive(Serialize, Clone, Debug)] pub struct Tool { - pub module: FlowModule, + pub module: Option, pub def: ToolDef, + pub mcp_source: Option, } #[derive(Deserialize, Serialize, Debug, Clone, PartialEq)] diff --git a/backend/windmill-worker/src/ai/utils.rs b/backend/windmill-worker/src/ai/utils.rs new file mode 100644 index 0000000000..70c022577a --- /dev/null +++ b/backend/windmill-worker/src/ai/utils.rs @@ -0,0 +1,465 @@ +use crate::ai::types::{ToolDef, ToolDefFunction}; +use anyhow::Context; +use serde_json::value::RawValue; +use std::{collections::HashMap, sync::Arc}; +use uuid::Uuid; +use windmill_common::mcp_client::{McpClient, McpResource, McpToolSource}; +use windmill_common::{ + ai_providers::AIProvider, + db::DB, + error::Error, + flow_conversations::{add_message_to_conversation_tx, MessageType}, + flow_status::AgentAction, + flows::Step, + jobs::JobKind, + scripts::{ScriptHash, ScriptLang}, + worker::to_raw_value, +}; +use windmill_queue::{flow_status::get_step_of_flow_status, MiniPulledJob}; + +use crate::{ai::types::*, parse_sig_of_lang}; + +pub fn parse_raw_script_schema( + content: &str, + language: &ScriptLang, +) -> Result, Error> { + let main_arg_signature = parse_sig_of_lang(content, Some(&language), None)?.unwrap(); // safe to unwrap as langauge is some + + let schema = OpenAPISchema { + r#type: Some(SchemaType::default()), + properties: Some( + main_arg_signature + .args + .iter() + .map(|arg| { + let name = arg.name.clone(); + let typ = OpenAPISchema::from_typ(&arg.typ); + (name, Box::new(typ)) + }) + .collect(), + ), + required: Some( + main_arg_signature + .args + .iter() + .map(|arg| arg.name.clone()) + .collect(), + ), + ..Default::default() + }; + + Ok(to_raw_value(&schema)) +} + +pub struct FlowJobRunnableIdAndRawFlow { + pub runnable_id: Option, + pub raw_flow: Option>>, + pub kind: JobKind, +} + +pub async fn get_flow_job_runnable_and_raw_flow( + db: &DB, + job_id: &uuid::Uuid, +) -> windmill_common::error::Result { + let job = sqlx::query_as!( + FlowJobRunnableIdAndRawFlow, + "SELECT runnable_id as \"runnable_id: ScriptHash\", raw_flow as \"raw_flow: _\", kind as \"kind: _\" FROM v2_job WHERE id = $1", + job_id + ) + .fetch_one(db) + .await?; + Ok(job) +} + +#[derive(Debug, Clone, Default)] +pub struct FlowChatSettings { + pub memory_id: Option, + pub chat_input_enabled: bool, +} + +/// Get chat settings (memory_id and chat_input_enabled) from root flow's flow_status +pub async fn get_flow_chat_settings(db: &DB, job: &MiniPulledJob) -> FlowChatSettings { + let root_job_id = job + .root_job + .or(job.flow_innermost_root_job) + .or(job.parent_job); + + let Some(root_job_id) = root_job_id else { + return FlowChatSettings::default(); + }; + + match sqlx::query!( + "SELECT + (flow_status->>'memory_id')::uuid as memory_id, + (flow_status->>'chat_input_enabled')::boolean as chat_input_enabled + FROM v2_job_status + WHERE id = $1", + root_job_id + ) + .fetch_optional(db) + .await + { + Ok(Some(row)) => FlowChatSettings { + memory_id: row.memory_id, + chat_input_enabled: row.chat_input_enabled.unwrap_or(false), + }, + Ok(None) => FlowChatSettings::default(), + Err(e) => { + tracing::warn!( + "Failed to get chat settings from flow status for job {}: {}", + job.id, + e + ); + FlowChatSettings::default() + } + } +} + +// Add message to conversation +pub async fn add_message_to_conversation( + db: &DB, + conversation_id: &Uuid, + job_id: Option, + message_content: &str, + message_type: MessageType, + step_name: &Option, + success: bool, +) -> Result<(), Error> { + let mut tx = db.begin().await?; + add_message_to_conversation_tx( + &mut tx, + *conversation_id, + job_id, + &message_content, + message_type, + step_name.as_deref(), + success, + ) + .await?; + tx.commit().await?; + Ok(()) +} + +/// Find a unique tool name for structured output tool to avoid collisions with user-provided tools +pub fn find_unique_tool_name(base_name: &str, existing_tools: Option<&[ToolDef]>) -> String { + let Some(tools) = existing_tools else { + return base_name.to_string(); + }; + + if !tools.iter().any(|t| t.function.name == base_name) { + return base_name.to_string(); + } + + for i in 1..100 { + let candidate = format!("{}_{}", base_name, i); + if !tools.iter().any(|t| t.function.name == candidate) { + return candidate; + } + } + + // Fallback with process id if somehow we can't find a unique name + format!("{}_{}_fallback", base_name, std::process::id()) +} + +pub async fn update_flow_status_module_with_actions( + db: &DB, + parent_job: &Uuid, + actions: &[AgentAction], +) -> Result<(), Error> { + let step = get_step_of_flow_status(db, parent_job.to_owned()).await?; + match step { + Step::Step { idx: step, .. } => { + sqlx::query!( + r#" + UPDATE v2_job_status SET + flow_status = jsonb_set( + flow_status, + array['modules', $3::TEXT, 'agent_actions'], + $2 + ) + WHERE id = $1 + "#, + parent_job, + sqlx::types::Json(actions) as _, + step as i32 + ) + .execute(db) + .await?; + } + _ => {} + } + Ok(()) +} + +pub async fn update_flow_status_module_with_actions_success( + db: &DB, + parent_job: &Uuid, + action_success: bool, +) -> Result<(), Error> { + let step = get_step_of_flow_status(db, parent_job.to_owned()).await?; + match step { + Step::Step { idx: step, .. } => { + // Append the new bool to the existing array, or create a new array if it doesn't exist + sqlx::query!( + r#" + UPDATE v2_job_status SET + flow_status = jsonb_set( + flow_status, + array['modules', $2::TEXT, 'agent_actions_success'], + COALESCE( + flow_status->'modules'->$2->'agent_actions_success', + to_jsonb(ARRAY[]::bool[]) + ) || to_jsonb(ARRAY[$3::bool]) + ) + WHERE id = $1 + "#, + parent_job, + step as i32, + action_success + ) + .execute(db) + .await?; + } + _ => {} + } + Ok(()) +} + +/// Get step name from the flow module (summary if exists, else id) +pub fn get_step_name_from_flow( + summary: Option<&str>, + flow_step_id: Option<&str>, +) -> Option { + let flow_step_id = flow_step_id?; + Some( + summary + .map(|s| s.to_string()) + .unwrap_or_else(|| format!("AI Agent Step {}", flow_step_id)), + ) +} + +/// Check if the provider is Anthropic (either direct or through OpenRouter) +pub fn is_anthropic_provider(provider: &ProviderWithResource) -> bool { + let provider_is_anthropic = provider.kind.is_anthropic(); + let is_openrouter_anthropic = + provider.kind == AIProvider::OpenRouter && provider.model.starts_with("anthropic/"); + provider_is_anthropic || is_openrouter_anthropic +} + +/// Cleanup MCP clients by gracefully shutting down connections +pub async fn cleanup_mcp_clients(mcp_clients: HashMap>) { + if mcp_clients.is_empty() { + return; + } + + tracing::debug!("Cleaning up {} MCP client(s)", mcp_clients.len()); + + for (resource_name, client) in mcp_clients { + // Try to unwrap the Arc to get the McpClient + match Arc::try_unwrap(client) { + Ok(client) => { + tracing::debug!("Shutting down MCP client for {}", resource_name); + if let Err(e) = client.shutdown().await { + tracing::warn!("Failed to shutdown MCP client for {}: {}", resource_name, e); + } + } + Err(arc) => { + // Other references still exist (shouldn't happen in normal flow) + tracing::warn!( + "MCP client for {} still has {} references, dropping without graceful shutdown", + resource_name, + Arc::strong_count(&arc) + ); + } + } + } +} + +/// Convert raw MCP tools to Windmill Tool format with source tracking +fn convert_mcp_tools_to_windmill_tools( + mcp_tools: &[rmcp::model::Tool], + resource_name: &str, + resource_path: &str, +) -> Result, Error> { + mcp_tools + .iter() + .map(|mcp_tool| { + let tool_name = format!("mcp_{}_{}", resource_name, mcp_tool.name); + + let mut schema_value = serde_json::to_value(&*mcp_tool.input_schema) + .context("Failed to convert MCP schema to JSON value")?; + McpClient::fix_array_schemas(&mut schema_value); + let parameters = to_raw_value(&schema_value); + + // Build the description from title and description + let description = if let Some(title) = &mcp_tool.title { + if let Some(desc) = &mcp_tool.description { + Some(format!("{}: {}", title, desc)) + } else { + Some(title.to_string()) + } + } else { + mcp_tool.description.as_ref().map(|d| d.to_string()) + }; + + let tool_def_function = + ToolDefFunction { name: tool_name.clone(), description, parameters }; + + let tool_def = ToolDef { r#type: "function".to_string(), function: tool_def_function }; + + Ok(Tool { + def: tool_def, + module: None, + mcp_source: Some(McpToolSource { + name: resource_name.to_string(), + tool_name: mcp_tool.name.to_string(), + resource_path: resource_path.to_string(), + }), + }) + }) + .collect() +} + +/// Configuration for loading tools from an MCP server resource +#[derive(Debug, Clone)] +pub struct McpResourceConfig { + pub resource_path: String, + pub include_tools: Option>, + pub exclude_tools: Option>, +} + +/// Apply include/exclude filters to a list of tools +/// Priority: include_tools > exclude_tools > all +/// - If include_tools is Some and non-empty: whitelist approach (keep only listed tools) +/// - Else if exclude_tools is Some and non-empty: blacklist approach (remove listed tools) +/// - Otherwise: no filtering (keep all tools) +fn apply_tool_filters( + tools: Vec, + include_tools: &Option>, + exclude_tools: &Option>, +) -> Vec { + // If include_tools is specified and non-empty, use whitelist approach + if let Some(include_list) = include_tools { + if !include_list.is_empty() { + return tools + .into_iter() + .filter(|tool| { + tool.mcp_source + .as_ref() + .map(|src| include_list.contains(&src.tool_name)) + .unwrap_or(false) + }) + .collect(); + } + } + + // If exclude_tools is specified and non-empty, use blacklist approach + if let Some(exclude_list) = exclude_tools { + if !exclude_list.is_empty() { + return tools + .into_iter() + .filter(|tool| { + tool.mcp_source + .as_ref() + .map(|src| !exclude_list.contains(&src.tool_name)) + .unwrap_or(true) + }) + .collect(); + } + } + + // No filtering - return all tools + tools +} + +/// Load tools from MCP servers and return both the clients and tools +/// Returns a map of resource name -> client, and a vector of tools +pub async fn load_mcp_tools( + db: &DB, + workspace_id: &str, + mcp_configs: Vec, +) -> Result<(HashMap>, Vec), Error> { + let mut all_mcp_tools = Vec::new(); + let mut mcp_clients = HashMap::new(); + + for config in mcp_configs { + tracing::debug!("Loading MCP tools from resource: {}", config.resource_path); + + let path = config.resource_path.trim_start_matches("$res:"); + let mcp_resource = { + // Fetch the resource from database + let resource= sqlx::query_scalar!( + "SELECT value as \"value: sqlx::types::Json>\" FROM resource WHERE path = $1 AND workspace_id = $2", + &path, + &workspace_id + ) + .fetch_optional(db) + .await? + .ok_or_else(|| Error::NotFound(format!("Could not find the resource {}, update the resource path in the workspace settings", config.resource_path)))? + .ok_or_else(|| Error::BadRequest(format!("Empty resource value for {}", config.resource_path)))?; + + serde_json::from_str::(resource.0.get()) + .context("Failed to parse MCP resource")? + }; + + let resource_name = mcp_resource.name.clone(); + + // Create new MCP client for this execution + tracing::debug!("Creating fresh MCP client for {}", resource_name); + let client = McpClient::from_resource(mcp_resource, db, workspace_id) + .await + .context("Failed to create MCP client")?; + + // Get raw MCP tools from client + let raw_mcp_tools = client.available_tools(); + + // Convert to Windmill Tool format + let converted_tools = + convert_mcp_tools_to_windmill_tools(raw_mcp_tools, &resource_name, &path)?; + + // Apply include/exclude filters + let filtered_tools = apply_tool_filters( + converted_tools, + &config.include_tools, + &config.exclude_tools, + ); + + tracing::info!( + "Loaded {} tools from MCP server '{}' (filtered from {} available tools)", + filtered_tools.len(), + resource_name, + raw_mcp_tools.len() + ); + + all_mcp_tools.extend(filtered_tools); + + // Store client for later use and cleanup + let mcp_client = Arc::new(client); + mcp_clients.insert(resource_name, mcp_client); + } + + Ok((mcp_clients, all_mcp_tools)) +} + +/// Execute an MCP tool by routing the call to the appropriate MCP client +pub async fn execute_mcp_tool( + mcp_clients: &HashMap>, + mcp_source: &McpToolSource, + arguments_str: &str, +) -> Result { + // Get the MCP client from the provided map + let mcp_client = mcp_clients.get(&mcp_source.name).ok_or_else(|| { + Error::internal_err(format!( + "MCP client not found for resource: {}", + mcp_source.name + )) + })?; + + // Call the MCP tool + let result = mcp_client + .call_tool(&mcp_source.tool_name, arguments_str) + .await + .context("MCP tool call failed")?; + + Ok(result) +} diff --git a/backend/windmill-worker/src/ai_executor.rs b/backend/windmill-worker/src/ai_executor.rs index bbd40979fe..6b6596f9cb 100644 --- a/backend/windmill-worker/src/ai_executor.rs +++ b/backend/windmill-worker/src/ai_executor.rs @@ -1,30 +1,34 @@ +use crate::ai::tools::{execute_tool_calls, ToolExecutionContext}; +use crate::ai::utils::{ + add_message_to_conversation, cleanup_mcp_clients, find_unique_tool_name, + get_flow_chat_settings, get_flow_job_runnable_and_raw_flow, get_step_name_from_flow, + is_anthropic_provider, load_mcp_tools, parse_raw_script_schema, + update_flow_status_module_with_actions, update_flow_status_module_with_actions_success, + FlowChatSettings, +}; use crate::memory_oss::{read_from_memory, write_to_memory}; -use anyhow::Context; use async_recursion::async_recursion; use regex::Regex; use serde_json::value::RawValue; use std::{collections::HashMap, sync::Arc}; -use ulid; use uuid::Uuid; +use windmill_common::mcp_client::McpClient; use windmill_common::{ - ai_providers::{AIProvider, AZURE_API_VERSION}, + ai_providers::AZURE_API_VERSION, cache, client::AuthedClient, db::DB, - error::{self, to_anyhow, Error}, - flow_conversations::{add_message_to_conversation_tx, MessageType}, + error::{self, Error}, + flow_conversations::MessageType, flow_status::AgentAction, - flows::{FlowModuleValue, FlowValue, Step}, + flows::{FlowModule, FlowModuleValue, ToolValue}, get_latest_hash_for_path, jobs::JobKind, - scripts::{get_full_hub_script_by_path, ScriptHash, ScriptLang}, + scripts::get_full_hub_script_by_path, utils::{StripPath, HTTP_CLIENT}, worker::{to_raw_value, Connection}, }; -use windmill_queue::{ - flow_status::get_step_of_flow_status, get_mini_pulled_job, push, CanceledBy, JobCompleted, - MiniPulledJob, PushArgs, PushIsolationLevel, -}; +use windmill_queue::{CanceledBy, MiniPulledJob}; use crate::{ ai::{ @@ -34,15 +38,9 @@ use crate::{ }, types::*, }, - common::{ - build_args_map, error_to_value, resolve_job_timeout, OccupancyMetrics, StreamNotifier, - }, - create_job_dir, + common::{build_args_map, resolve_job_timeout, OccupancyMetrics, StreamNotifier}, handle_child::run_future_with_polling_update_job_poller, - handle_queued_job, parse_sig_of_lang, - result_processor::handle_non_flow_job_error, - worker_flow::{raw_script_to_payload, script_to_payload}, - JobCompletedSender, SendResult, SendResultPayload, + JobCompletedSender, }; lazy_static::lazy_static! { @@ -51,124 +49,6 @@ lazy_static::lazy_static! { const MAX_AGENT_ITERATIONS: usize = 10; -fn parse_raw_script_schema(content: &str, language: &ScriptLang) -> Result, Error> { - let main_arg_signature = parse_sig_of_lang(content, Some(&language), None)?.unwrap(); // safe to unwrap as langauge is some - - let schema = OpenAPISchema { - r#type: Some(SchemaType::default()), - properties: Some( - main_arg_signature - .args - .iter() - .map(|arg| { - let name = arg.name.clone(); - let typ = OpenAPISchema::from_typ(&arg.typ); - (name, Box::new(typ)) - }) - .collect(), - ), - required: Some( - main_arg_signature - .args - .iter() - .map(|arg| arg.name.clone()) - .collect(), - ), - ..Default::default() - }; - - Ok(to_raw_value(&schema)) -} - -pub struct FlowJobRunnableIdAndRawFlow { - pub runnable_id: Option, - pub raw_flow: Option>>, - pub kind: JobKind, -} - -pub async fn get_flow_job_runnable_and_raw_flow( - db: &DB, - job_id: &uuid::Uuid, -) -> windmill_common::error::Result { - let job = sqlx::query_as!( - FlowJobRunnableIdAndRawFlow, - "SELECT runnable_id as \"runnable_id: ScriptHash\", raw_flow as \"raw_flow: _\", kind as \"kind: _\" FROM v2_job WHERE id = $1", - job_id - ) - .fetch_one(db) - .await?; - Ok(job) -} - -#[derive(Debug, Clone, Default)] -struct FlowChatSettings { - memory_id: Option, - chat_input_enabled: bool, -} - -/// Get chat settings (memory_id and chat_input_enabled) from root flow's flow_status -async fn get_flow_chat_settings(db: &DB, job: &MiniPulledJob) -> FlowChatSettings { - let root_job_id = job - .root_job - .or(job.flow_innermost_root_job) - .or(job.parent_job); - - let Some(root_job_id) = root_job_id else { - return FlowChatSettings::default(); - }; - - match sqlx::query!( - "SELECT - (flow_status->>'memory_id')::uuid as memory_id, - (flow_status->>'chat_input_enabled')::boolean as chat_input_enabled - FROM v2_job_status - WHERE id = $1", - root_job_id - ) - .fetch_optional(db) - .await - { - Ok(Some(row)) => FlowChatSettings { - memory_id: row.memory_id, - chat_input_enabled: row.chat_input_enabled.unwrap_or(false), - }, - Ok(None) => FlowChatSettings::default(), - Err(e) => { - tracing::warn!( - "Failed to get chat settings from flow status for job {}: {}", - job.id, - e - ); - FlowChatSettings::default() - } - } -} - -// Add message to conversation -async fn add_message_to_conversation( - db: &DB, - conversation_id: &Uuid, - job_id: &Uuid, - message_content: &str, - message_type: MessageType, - step_name: &Option, - success: bool, -) -> Result<(), Error> { - let mut tx = db.begin().await?; - add_message_to_conversation_tx( - &mut tx, - *conversation_id, - Some(*job_id), - &message_content, - message_type, - step_name.as_deref(), - success, - ) - .await?; - tx.commit().await?; - Ok(()) -} - pub async fn handle_ai_agent_job( // connection conn: &Connection, @@ -191,7 +71,6 @@ pub async fn handle_ai_agent_job( has_stream: &mut bool, ) -> Result, Error> { let args = build_args_map(job, client, conn).await?; - let args = serde_json::from_str::(&serde_json::to_string(&args)?)?; let Some(flow_step_id) = &job.flow_step_id else { @@ -225,6 +104,7 @@ pub async fn handle_ai_agent_job( let value = flow_data.value(); let module = value.modules.iter().find(|m| m.id == *flow_step_id); + let summary = module.as_ref().and_then(|m| m.summary.clone()); let Some(module) = module else { return Err(Error::internal_err( @@ -238,7 +118,38 @@ pub async fn handle_ai_agent_job( )); }; - let tools = futures::future::try_join_all(tools.into_iter().map(|mut t| { + // Separate Windmill tools from MCP tools and extract MCP resource configs + let mut windmill_modules: Vec = Vec::new(); + let mut mcp_configs: Vec = Vec::new(); + + for tool in tools { + match &tool.value { + ToolValue::Mcp(mcp_config) => { + // This is an MCP tool - extract config + tracing::debug!( + "MCP server module: path={}, include={:?}, exclude={:?}", + mcp_config.resource_path, + mcp_config.include_tools, + mcp_config.exclude_tools + ); + mcp_configs.push(crate::ai::utils::McpResourceConfig { + resource_path: mcp_config.resource_path.clone(), + include_tools: Some(mcp_config.include_tools.clone()), + exclude_tools: Some(mcp_config.exclude_tools.clone()), + }); + } + ToolValue::FlowModule(_) => { + // Regular Windmill flow module (script, flow, etc.) - convert to FlowModule + tracing::debug!("Windmill module: {:?}", tool.id); + if let Some(flow_module) = Option::::from(&tool) { + windmill_modules.push(flow_module); + } + } + } + } + + // Process Windmill flow modules into Tool definitions + let tools = futures::future::try_join_all(windmill_modules.into_iter().map(|mut t| { let conn = conn; let db = db; let job = job; @@ -279,9 +190,14 @@ pub async fn handle_ai_agent_job( .await?; Ok(Some(hub_script.schema)) } else { - let hash = get_latest_hash_for_path(db, &job.workspace_id, path, true) - .await? - .0; + let hash = get_latest_hash_for_path( + db, + &job.workspace_id, + path.as_str(), + true, + ) + .await? + .0; // update module definition to use a fixed hash so all tool calls match the same schema t.value = to_raw_value(&FlowModuleValue::Script { hash: Some(hash), @@ -333,12 +249,23 @@ pub async fn handle_ai_agent_job( }), }, }, - module: t, + module: Some(t), + mcp_source: None, }) } })) .await?; + // Load MCP tools if configured + let mut tools = tools; + let mcp_clients = if !mcp_configs.is_empty() { + let (clients, mcp_tools) = load_mcp_tools(db, &job.workspace_id, mcp_configs).await?; + tools.extend(mcp_tools); + clients + } else { + HashMap::new() + }; + let mut inner_occupancy_metrics = occupancy_metrics.clone(); let stream_notifier = StreamNotifier::new(conn, job); @@ -354,7 +281,8 @@ pub async fn handle_ai_agent_job( parent_job, &args, &tools, - value, + &mcp_clients, + summary.as_deref(), client, &mut inner_occupancy_metrics, job_completed_tx, @@ -380,114 +308,12 @@ pub async fn handle_ai_agent_job( ) .await?; + // Cleanup MCP clients + cleanup_mcp_clients(mcp_clients).await; + Ok(result) } -/// Find a unique tool name for structured output tool to avoid collisions with user-provided tools -fn find_unique_tool_name(base_name: &str, existing_tools: Option<&[ToolDef]>) -> String { - let Some(tools) = existing_tools else { - return base_name.to_string(); - }; - - if !tools.iter().any(|t| t.function.name == base_name) { - return base_name.to_string(); - } - - for i in 1..100 { - let candidate = format!("{}_{}", base_name, i); - if !tools.iter().any(|t| t.function.name == candidate) { - return candidate; - } - } - - // Fallback with process id if somehow we can't find a unique name - format!("{}_{}_fallback", base_name, std::process::id()) -} - -async fn update_flow_status_module_with_actions( - db: &DB, - parent_job: &Uuid, - actions: &[AgentAction], -) -> Result<(), Error> { - let step = get_step_of_flow_status(db, parent_job.to_owned()).await?; - match step { - Step::Step { idx: step, .. } => { - sqlx::query!( - r#" - UPDATE v2_job_status SET - flow_status = jsonb_set( - flow_status, - array['modules', $3::TEXT, 'agent_actions'], - $2 - ) - WHERE id = $1 - "#, - parent_job, - sqlx::types::Json(actions) as _, - step as i32 - ) - .execute(db) - .await?; - } - _ => {} - } - Ok(()) -} - -async fn update_flow_status_module_with_actions_success( - db: &DB, - parent_job: &Uuid, - action_success: bool, -) -> Result<(), Error> { - let step = get_step_of_flow_status(db, parent_job.to_owned()).await?; - match step { - Step::Step { idx: step, .. } => { - // Append the new bool to the existing array, or create a new array if it doesn't exist - sqlx::query!( - r#" - UPDATE v2_job_status SET - flow_status = jsonb_set( - flow_status, - array['modules', $2::TEXT, 'agent_actions_success'], - COALESCE( - flow_status->'modules'->$2->'agent_actions_success', - to_jsonb(ARRAY[]::bool[]) - ) || to_jsonb(ARRAY[$3::bool]) - ) - WHERE id = $1 - "#, - parent_job, - step as i32, - action_success - ) - .execute(db) - .await?; - } - _ => {} - } - Ok(()) -} - -/// Get step name from the flow module (summary if exists, else id) -fn get_step_name_from_flow(flow_value: &FlowValue, flow_step_id: Option<&str>) -> Option { - let flow_step_id = flow_step_id?; - let module = flow_value.modules.iter().find(|m| m.id == flow_step_id)?; - Some( - module - .summary - .clone() - .unwrap_or_else(|| format!("AI Agent Step {}", module.id)), - ) -} - -/// Check if the provider is Anthropic (either direct or through OpenRouter) -fn is_anthropic_provider(provider: &ProviderWithResource) -> bool { - let provider_is_anthropic = provider.kind.is_anthropic(); - let is_openrouter_anthropic = - provider.kind == AIProvider::OpenRouter && provider.model.starts_with("anthropic/"); - provider_is_anthropic || is_openrouter_anthropic -} - #[async_recursion] pub async fn run_agent( // connection @@ -499,7 +325,8 @@ pub async fn run_agent( parent_job: &Uuid, args: &AIAgentArgs, tools: &[Tool], - flow_value: &FlowValue, + mcp_clients: &HashMap>, + summary: Option<&str>, // job execution context client: &AuthedClient, @@ -751,7 +578,7 @@ pub async fn run_agent( let db_clone = db.clone(); let message_content = response_content.clone(); let step_name = get_step_name_from_flow( - flow_value, + summary.as_deref(), job.flow_step_id.as_deref(), ); @@ -760,7 +587,7 @@ pub async fn run_agent( if let Err(e) = add_message_to_conversation( &db_clone, &mid, - &agent_job_id, + Some(agent_job_id), &message_content, MessageType::Assistant, &step_name, @@ -790,481 +617,42 @@ pub async fn run_agent( ..Default::default() }); - // Handle tool calls (keeping existing tool execution logic) - for tool_call in tool_calls.iter() { - // Stream tool call progress - if let Some(ref stream_event_processor) = stream_event_processor { - let event = StreamingEvent::ToolExecution { - call_id: tool_call.id.clone(), - function_name: tool_call.function.name.clone(), - }; - stream_event_processor - .send(event, &mut final_events_str) - .await?; - } + // Handle tool calls using extracted tools module + let tool_execution_ctx = ToolExecutionContext { + db, + conn, + job, + parent_job, + summary: &summary, + client, + worker_dir, + base_internal_url, + worker_name, + hostname, + occupancy_metrics, + job_completed_tx, + killpill_rx, + stream_event_processor: stream_event_processor.as_ref(), + chat_settings: &mut chat_settings, + }; - // Check if this is the structured output tool - if structured_output_tool_name - .as_ref() - .map_or(false, |name| tool_call.function.name == *name) - { - used_structured_output_tool = true; - messages.push(OpenAIMessage { - role: "tool".to_string(), - content: Some(OpenAIContent::Text( - "Successfully ran structured_output tool".to_string(), - )), - tool_call_id: Some(tool_call.id.clone()), - ..Default::default() - }); - messages.push(OpenAIMessage { - role: "assistant".to_string(), - content: Some(OpenAIContent::Text( - tool_call.function.arguments.clone(), - )), - agent_action: Some(AgentAction::Message {}), - ..Default::default() - }); - content = - Some(OpenAIContent::Text(tool_call.function.arguments.clone())); - break; - } + let (tool_messages, tool_content, tool_used_structured_output) = + execute_tool_calls( + tool_execution_ctx, + &tool_calls, + &tools, + mcp_clients, + &mut actions, + &mut final_events_str, + &structured_output_tool_name, + ) + .await?; - // Execute regular tool - let tool = tools - .iter() - .find(|t| t.def.function.name == tool_call.function.name); - if let Some(tool) = tool { - let job_id = ulid::Ulid::new().into(); - actions.push(AgentAction::ToolCall { - job_id, - function_name: tool_call.function.name.clone(), - module_id: tool.module.id.clone(), - }); - - update_flow_status_module_with_actions(db, parent_job, &actions) - .await?; - - let raw_tool_call_args = if tool_call.function.arguments.is_empty() - { - "{}".to_string() - } else { - tool_call.function.arguments.clone() - }; - let tool_call_args = - serde_json::from_str::>>( - &raw_tool_call_args, - ) - .with_context(|| { - format!( - "Failed to parse tool call arguments for tool call {}: {}", - tool_call.function.name, tool_call.function.arguments - ) - })?; - - let job_payload = match tool.module.get_value()? { - FlowModuleValue::Script { - path: script_path, - hash: script_hash, - tag_override, - .. - } => { - let payload = script_to_payload( - script_hash, - script_path, - db, - job, - &tool.module, - tag_override, - tool.module.apply_preprocessor, - ) - .await?; - payload - } - FlowModuleValue::RawScript { - path, - content, - language, - lock, - tag, - custom_concurrency_key, - concurrent_limit, - concurrency_time_window_s, - .. - } => { - let path = path.unwrap_or_else(|| { - format!( - "{}/tools/{}", - job.runnable_path(), - tool.module.id - ) - }); - - let payload = raw_script_to_payload( - path, - content, - language, - lock, - custom_concurrency_key, - concurrent_limit, - concurrency_time_window_s, - &tool.module, - tag, - tool.module.delete_after_use.unwrap_or(false), - ); - payload - } - _ => { - return Err(Error::internal_err(format!( - "Unsupported tool: {}", - tool_call.function.name - ))); - } - }; - - let mut tx = db.begin().await?; - - let job_perms = windmill_common::auth::get_job_perms( - &mut *tx, - &job.id, - &job.workspace_id, - ) - .await? - .map(|x| x.into()); - - let (email, permissioned_as) = - if let Some(on_behalf_of) = job_payload.on_behalf_of.as_ref() { - (&on_behalf_of.email, on_behalf_of.permissioned_as.clone()) - } else { - (&job.permissioned_as_email, job.permissioned_as.to_owned()) - }; - - let job_priority = tool.module.priority.or(job.priority); - - let tx = PushIsolationLevel::Transaction(tx); - let (uuid, tx) = push( - db, - tx, - &job.workspace_id, - job_payload.payload, - PushArgs { args: &tool_call_args, extra: None }, - &job.created_by, - email, - permissioned_as, - Some(&format!("job-span-{}", job.id)), - None, - job.schedule_path(), - Some(job.id), - None, - None, - Some(job_id), - false, - false, - None, - job.visible_to_owner, - Some(job.tag.clone()), - job_payload.timeout, - None, - job_priority, - job_perms.as_ref(), - true, - None, - None, - ) - .await?; - - tx.commit().await?; - - let tool_job = get_mini_pulled_job(db, &uuid).await?; - - let Some(tool_job) = tool_job else { - return Err(Error::internal_err( - "Tool job not found".to_string(), - )); - }; - - let tool_job = Arc::new(tool_job); - - let (inner_job_completed_tx, inner_job_completed_rx) = - JobCompletedSender::new(&conn, 1); - - let inner_job_completed_rx = inner_job_completed_rx.expect( - "inner_job_completed_tx should be set as agent jobs are not supported on agent workers", - ); - - // Spawn handle_queued_job on separate task to prevent tokio stack overflow - // Clone everything needed for the spawned task - let tool_job_spawn = tool_job.clone(); - let conn_spawn = conn.clone(); - let client_spawn = client.clone(); - let hostname_spawn = hostname.to_string(); - let worker_name_spawn = worker_name.to_string(); - let worker_dir_spawn = worker_dir.to_string(); - let base_internal_url_spawn = base_internal_url.to_string(); - let inner_job_completed_tx_spawn = inner_job_completed_tx.clone(); - let mut occupancy_metrics_spawn = occupancy_metrics.clone(); - let mut killpill_rx_spawn = killpill_rx.resubscribe(); - - // Spawn on separate tokio task with fresh stack - let join_handle = tokio::task::spawn(async move { - #[cfg(feature = "benchmark")] - let mut bench_spawn = - windmill_common::bench::BenchmarkIter::new(); - - let job_dir = - create_job_dir(&worker_dir_spawn, tool_job_spawn.id).await; - - let result = handle_queued_job( - tool_job_spawn, - None, - None, - None, - None, - &conn_spawn, - &client_spawn, - &hostname_spawn, - &worker_name_spawn, - &worker_dir_spawn, - &job_dir, - None, - &base_internal_url_spawn, - inner_job_completed_tx_spawn, - &mut occupancy_metrics_spawn, - &mut killpill_rx_spawn, - None, - #[cfg(feature = "benchmark")] - &mut bench_spawn, - ) - .await; - - // Return both result and updated metrics - (result, occupancy_metrics_spawn) - }); - - // Await the spawned task - let (handle_result, updated_occupancy) = - join_handle.await.map_err(|e| { - Error::internal_err(format!( - "Tool execution task failed: {}", - e - )) - })?; - - // Merge occupancy metrics back - occupancy_metrics.total_duration_of_running_jobs = - updated_occupancy.total_duration_of_running_jobs; - - // Continue with match on handle_result - match handle_result { - Err(err) => { - let err_string = - format!("{}: {}", err.name(), err.to_string()); - let err_json = error_to_value(&err); - let _ = handle_non_flow_job_error( - db, - &tool_job, - 0, - None, - err_string.clone(), - err_json, - worker_name, - ) - .await; - let error_message = - format!("Error running tool: {}", err_string); - messages.push(OpenAIMessage { - role: "tool".to_string(), - content: Some(OpenAIContent::Text( - error_message.clone(), - )), - tool_call_id: Some(tool_call.id.clone()), - agent_action: Some(AgentAction::ToolCall { - job_id, - function_name: tool_call.function.name.clone(), - module_id: tool.module.id.clone(), - }), - ..Default::default() - }); - // Stream tool result (error case) - if let Some(ref stream_event_processor) = - stream_event_processor - { - let tool_result_event = StreamingEvent::ToolResult { - call_id: tool_call.id.clone(), - function_name: tool_call.function.name.clone(), - result: error_message.clone(), - success: false, - }; - stream_event_processor - .send(tool_result_event, &mut final_events_str) - .await?; - } - - update_flow_status_module_with_actions_success( - db, parent_job, false, - ) - .await?; - - // Add tool message to conversation if chat_input_enabled (error case) - if chat_settings.is_none() { - chat_settings = - Some(get_flow_chat_settings(db, job).await); - } - let chat_enabled = chat_settings - .as_ref() - .map(|s| s.chat_input_enabled) - .unwrap_or(false); - - if chat_enabled { - if let Some(mid) = - chat_settings.as_ref().and_then(|s| s.memory_id) - { - let tool_job_id = job_id; - let db_clone = db.clone(); - let step_name = get_step_name_from_flow( - flow_value, - job.flow_step_id.as_deref(), - ); - - // Spawn task because we do not need to wait for the result - tokio::spawn(async move { - if let Err(e) = add_message_to_conversation( - &db_clone, - &mid, - &tool_job_id, - &error_message, - MessageType::Tool, - &step_name, - false, - ) - .await - { - tracing::warn!("Failed to add tool error message to conversation {}: {}", mid, e); - } - }); - } - } - } - Ok(success) => { - let send_result = - inner_job_completed_rx.bounded_rx.try_recv().ok(); - - let result = if let Some(SendResult { - result: - SendResultPayload::JobCompleted(JobCompleted { - result, - .. - }), - .. - }) = send_result.as_ref() - { - job_completed_tx - .send( - send_result.as_ref().unwrap().result.clone(), - true, - ) - .await - .map_err(to_anyhow)?; - result - } else { - if let Some(send_result) = send_result { - job_completed_tx - .send(send_result.result, true) - .await - .map_err(to_anyhow)?; - } - return Err(Error::internal_err( - "Tool job completed but no result".to_string(), - )); - }; - messages.push(OpenAIMessage { - role: "tool".to_string(), - content: Some(OpenAIContent::Text( - result.get().to_string(), - )), - tool_call_id: Some(tool_call.id.clone()), - agent_action: Some(AgentAction::ToolCall { - job_id, - function_name: tool_call.function.name.clone(), - module_id: tool.module.id.clone(), - }), - ..Default::default() - }); - - // Stream tool result (success case) - if let Some(ref stream_event_processor) = - stream_event_processor - { - let tool_result_event = StreamingEvent::ToolResult { - call_id: tool_call.id.clone(), - function_name: tool_call.function.name.clone(), - result: result.get().to_string(), - success: true, - }; - stream_event_processor - .send(tool_result_event, &mut final_events_str) - .await?; - } - - update_flow_status_module_with_actions_success( - db, parent_job, success, - ) - .await?; - - // Add tool message to conversation if chat_input_enabled - if chat_settings.is_none() { - chat_settings = - Some(get_flow_chat_settings(db, job).await); - } - let chat_enabled = chat_settings - .as_ref() - .map(|s| s.chat_input_enabled) - .unwrap_or(false); - - if chat_enabled { - if let Some(mid) = - chat_settings.as_ref().and_then(|s| s.memory_id) - { - let tool_job_id = job_id; - let db_clone = db.clone(); - let tool_name = tool_call.function.name.clone(); - let step_name = get_step_name_from_flow( - flow_value, - job.flow_step_id.as_deref(), - ); - let content = if success { - format!("Used {} tool", tool_name) - } else { - format!("Error executing {}", tool_name) - }; - - // Spawn task because we do not need to wait for the result - tokio::spawn(async move { - if let Err(e) = add_message_to_conversation( - &db_clone, - &mid, - &tool_job_id, - &content, - MessageType::Tool, - &step_name, - success, - ) - .await - { - tracing::warn!("Failed to add tool message to conversation {}: {}", mid, e); - } - }); - } - } - } - } - } else { - return Err(Error::internal_err(format!( - "Tool not found: {}", - tool_call.function.name - ))); - } + messages.extend(tool_messages); + if let Some(tc) = tool_content { + content = Some(tc); } + used_structured_output_tool = tool_used_structured_output; } ParsedResponse::Image { base64_data } => { // For image output with tools, we got an image response diff --git a/backend/windmill-worker/src/worker_flow.rs b/backend/windmill-worker/src/worker_flow.rs index 7a75068fcf..c566c99a07 100644 --- a/backend/windmill-worker/src/worker_flow.rs +++ b/backend/windmill-worker/src/worker_flow.rs @@ -3178,7 +3178,9 @@ async fn push_next_flow_job( tracing::debug!(id = %flow_job.id, root_id = %job_root, "pushed next flow job: {uuid}"); - if value_with_parallel.type_ == "forloopflow" && value_with_parallel.parallel.unwrap_or(false) { + if value_with_parallel.type_ == "forloopflow" + && value_with_parallel.parallel.unwrap_or(false) + { if let Some(parallelism_transform) = &value_with_parallel.parallelism { tracing::debug!(id = %flow_job.id, root_id = %job_root, "evaluating parallelism expression for forloopflow job {uuid}"); diff --git a/frontend/src/lib/components/AIAgentLogViewer.svelte b/frontend/src/lib/components/AIAgentLogViewer.svelte index c8fd14d6fc..1671fe7b77 100644 --- a/frontend/src/lib/components/AIAgentLogViewer.svelte +++ b/frontend/src/lib/components/AIAgentLogViewer.svelte @@ -11,6 +11,7 @@ import FlowLogViewerWrapper from './FlowLogViewerWrapper.svelte' import { z } from 'zod' import { onMount } from 'svelte' + import type { AgentTool } from './flows/agentToolUtils' type AgentActionWithContent = NonNullable[number] & { content?: unknown @@ -29,6 +30,13 @@ module_id: z.string(), function_name: z.string() }), + z.object({ + type: z.literal('mcp_tool_call'), + call_id: z.string(), + function_name: z.string(), + resource_path: z.string(), + arguments: z.record(z.unknown()).optional() + }), z.object({ type: z.literal('message') }) @@ -39,7 +47,7 @@ }) interface Props { - tools: FlowModule[] + tools: AgentTool[] agentJob: Partial & Pick & { type: 'CompletedJob' } workspaceId?: string | undefined storedToolCallJobs?: Record @@ -69,6 +77,13 @@ job_id: toolCall.job_id } onToolJobLoaded?.(job, idx) + } else if (toolCall.type === 'mcp_tool_call') { + fakeModuleStates[idx.toString()] = { + type: 'Success', + args: toolCall.arguments ?? {}, + logs: '', + result: toolCall.content + } } else { fakeModuleStates[idx.toString()] = { type: 'Success', @@ -104,7 +119,15 @@ module_id: m.agent_action.module_id, function_name: m.agent_action.function_name } - : undefined) as AgentActionWithContent | undefined + : m.agent_action?.type === 'mcp_tool_call' + ? { + type: 'mcp_tool_call', + content: m.content, + call_id: m.agent_action.call_id, + function_name: m.agent_action.function_name, + arguments: m.agent_action.arguments + } + : undefined) as AgentActionWithContent | undefined ) .filter((m) => m !== undefined) @@ -122,13 +145,22 @@ type: 'identity' as const } } + } else if (toolCall.type === 'mcp_tool_call') { + return { + id: idx.toString(), + value: { + type: 'identity' as const + }, + summary: toolCall.function_name, + arguments: toolCall.arguments + } } else { const module = tools.find((m) => m.summary === toolCall.function_name) return module - ? { + ? ({ ...module, id: idx.toString() - } + } as FlowModule) : undefined } }) diff --git a/frontend/src/lib/components/FlowJobResult.svelte b/frontend/src/lib/components/FlowJobResult.svelte index 0d90b025d0..ed35246c19 100644 --- a/frontend/src/lib/components/FlowJobResult.svelte +++ b/frontend/src/lib/components/FlowJobResult.svelte @@ -2,8 +2,9 @@ import { Loader2 } from 'lucide-svelte' import DisplayResult from './DisplayResult.svelte' import LogViewer from './LogViewer.svelte' - import type { CompletedJob, FlowModule, Job } from '$lib/gen' + import type { CompletedJob, Job } from '$lib/gen' import AiAgentLogViewer from './AIAgentLogViewer.svelte' + import type { AgentTool } from './flows/agentToolUtils' interface Props { waitingForExecutor?: boolean @@ -21,7 +22,7 @@ downloadLogs?: boolean tagLabel?: string | undefined aiAgentStatus?: { - tools: FlowModule[] + tools: AgentTool[] agentJob: Partial & Pick & { type: 'CompletedJob' } storedToolCallJobs?: Record onToolJobLoaded?: (job: Job, idx: number) => void diff --git a/frontend/src/lib/components/FlowLogViewer.svelte b/frontend/src/lib/components/FlowLogViewer.svelte index 6374f69930..53082b1387 100644 --- a/frontend/src/lib/components/FlowLogViewer.svelte +++ b/frontend/src/lib/components/FlowLogViewer.svelte @@ -693,12 +693,16 @@
- - {mode === 'aiagent' - ? module.summary - ? 'Tool call' - : 'Message' - : module.id} + + {#if mode === 'aiagent'} + {#if module.summary} + Tool call: {module.summary} + {:else} + Message + {/if} + {:else} + {module.id} + {/if} {#if mode === 'flow'} {#if module.value.type === 'forloopflow'} @@ -715,7 +719,7 @@ Step {/if} {/if} - {#if module.summary} + {#if module.summary && mode !== 'aiagent'} : {module.summary} {/if} {#if hasEmptySubflowValue} diff --git a/frontend/src/lib/components/FlowStatusViewerInner.svelte b/frontend/src/lib/components/FlowStatusViewerInner.svelte index 7bd1750246..c6f9a45efe 100644 --- a/frontend/src/lib/components/FlowStatusViewerInner.svelte +++ b/frontend/src/lib/components/FlowStatusViewerInner.svelte @@ -43,9 +43,11 @@ import { AI_TOOL_CALL_PREFIX, AI_TOOL_MESSAGE_PREFIX, + AI_MCP_TOOL_CALL_PREFIX, getToolCallId } from './graph/renderers/nodes/AIToolNode.svelte' import JobAssetsViewer from './assets/JobAssetsViewer.svelte' + import McpToolCallDetails from './McpToolCallDetails.svelte' let { flowState: flowStateStore, @@ -678,6 +680,12 @@ job_id: action.job_id, type: success != undefined ? (success ? 'Success' : 'Failure') : 'InProgress' }) + } else if (action.type == 'mcp_tool_call') { + const mcpToolCallId = AI_MCP_TOOL_CALL_PREFIX + '-' + mod.id + '-' + idx + const success = mod.agent_actions_success?.[idx] + setModuleState(mcpToolCallId, { + type: success != undefined ? (success ? 'Success' : 'Failure') : 'InProgress' + }) } else if (action.type == 'message') { const toolCallId = getToolCallId(idx, mod.id) setModuleState(toolCallId, { @@ -1806,6 +1814,27 @@
+ {:else if selectedNode?.startsWith(AI_MCP_TOOL_CALL_PREFIX)} + {@const [, agentModuleId, toolCallIndex] = selectedNode.split('-')} + {@const agentNode = localModuleStates?.[agentModuleId]} + {@const agentActions = agentNode?.agent_actions} + {@const mcpActionIndex = parseInt(toolCallIndex)} + {@const mcpAction = + agentActions && mcpActionIndex >= 0 && mcpActionIndex < agentActions.length + ? agentActions[mcpActionIndex] + : undefined} + {#if mcpAction?.type === 'mcp_tool_call' && agentNode?.result?.messages} + {@const message = agentNode.result.messages.find( + (m) => m.agent_action?.call_id === mcpAction.call_id + )} + + {/if} {:else if selectedNode} {@const node = localModuleStates[selectedNode]} {#if selectedNode == 'end'} diff --git a/frontend/src/lib/components/McpToolCallDetails.svelte b/frontend/src/lib/components/McpToolCallDetails.svelte new file mode 100644 index 0000000000..560556a5a9 --- /dev/null +++ b/frontend/src/lib/components/McpToolCallDetails.svelte @@ -0,0 +1,44 @@ + + +
+ +
+ {functionName} + + {type} + +
+ + +
+

Arguments

+ {#if args && typeof args === 'object' && Object.keys(args).length > 0} + + {:else} +

No arguments

+ {/if} +
+ + +
+

Result

+
+ +
+
+
diff --git a/frontend/src/lib/components/copilot/chat/flow/FlowAIChat.svelte b/frontend/src/lib/components/copilot/chat/flow/FlowAIChat.svelte index 5dd76d2ba0..dc48253240 100644 --- a/frontend/src/lib/components/copilot/chat/flow/FlowAIChat.svelte +++ b/frontend/src/lib/components/copilot/chat/flow/FlowAIChat.svelte @@ -17,6 +17,7 @@ import { aiChatManager } from '../AIChatManager.svelte' import { refreshStateStore } from '$lib/svelte5Utils.svelte' import DiffDrawer from '$lib/components/DiffDrawer.svelte' + import type { AgentTool } from '$lib/components/flows/agentToolUtils' let { flowModuleSchemaMap @@ -267,7 +268,7 @@ const indexToInsertAt = index + 1 - let newModules: FlowModule[] | undefined = undefined + let newModules: FlowModule[] | AgentTool[] | undefined = undefined switch (step.type) { case 'rawscript': { const inlineScript = { @@ -532,7 +533,7 @@ module.value.parallelism = undefined } else if (module.value.parallel || opts.parallel === true) { // Only set parallelism if parallel is enabled - const n = Math.max(1, Math.floor(Math.abs(opts.parallelism))); + const n = Math.max(1, Math.floor(Math.abs(opts.parallelism))) module.value.parallelism = { type: 'static', value: n diff --git a/frontend/src/lib/components/flows/agentToolUtils.ts b/frontend/src/lib/components/flows/agentToolUtils.ts new file mode 100644 index 0000000000..2c71ec2805 --- /dev/null +++ b/frontend/src/lib/components/flows/agentToolUtils.ts @@ -0,0 +1,58 @@ +import type { AiAgent, FlowModule, FlowModuleValue } from '$lib/gen' + +// Type aliases for better readability +export type AgentTool = AiAgent['tools'][number] +export type FlowModuleTool = AgentTool & { value: { tool_type: 'flowmodule' } & FlowModuleValue } +export type McpTool = AgentTool & { + value: { + tool_type: 'mcp' + resource_path: string + include_tools?: string[] + exclude_tools?: string[] + } +} + +/** + * Type guard to check if a tool is a FlowModule tool + */ +export function isFlowModuleTool(tool: AgentTool): tool is FlowModuleTool { + return tool.value.tool_type === undefined || tool.value.tool_type === 'flowmodule' +} + +/** + * Type guard to check if a tool is an MCP tool + */ +export function isMcpTool(tool: AgentTool): tool is McpTool { + return tool.value.tool_type === 'mcp' +} + +/** + * Create an MCP tool from resource path + */ +export function createMcpTool(id: string): McpTool { + return { + id, + summary: '', + value: { + tool_type: 'mcp', + resource_path: '', + include_tools: [], + exclude_tools: [] + } + } +} + +/** + * Convert a FlowModule back to an AgentTool + * Used when saving changes back to the AI Agent tools array + */ +export function flowModuleToAgentTool(flowModule: FlowModule): AgentTool { + return { + id: flowModule.id, + summary: flowModule.summary, + value: { + tool_type: 'flowmodule', + ...flowModule.value + } as FlowModuleTool['value'] + } +} diff --git a/frontend/src/lib/components/flows/content/AgentToolWrapper.svelte b/frontend/src/lib/components/flows/content/AgentToolWrapper.svelte new file mode 100644 index 0000000000..7933e06048 --- /dev/null +++ b/frontend/src/lib/components/flows/content/AgentToolWrapper.svelte @@ -0,0 +1,49 @@ + + +{#if isFlowModuleTool(tool)} + + +{:else if isMcpTool(tool)} + + +{/if} diff --git a/frontend/src/lib/components/flows/content/FlowInputsQuick.svelte b/frontend/src/lib/components/flows/content/FlowInputsQuick.svelte index 76b7beaa85..f88030f29a 100644 --- a/frontend/src/lib/components/flows/content/FlowInputsQuick.svelte +++ b/frontend/src/lib/components/flows/content/FlowInputsQuick.svelte @@ -172,7 +172,10 @@ ] let topLevelNodes: [string, string][] = $state([]) - function computeToplevelNodeChoices(funcDesc: string, preFilter: 'all' | 'workspace' | 'hub') { + function computeToplevelNodeChoices( + funcDesc: string, + preFilter: 'all' | 'workspace' | 'hub' + ) { if (funcDesc.length > 0 && preFilter == 'all' && kind == 'script') { topLevelNodes = allToplevelNodes.filter((node) => node[0].toLowerCase().startsWith(funcDesc.toLowerCase()) diff --git a/frontend/src/lib/components/flows/content/FlowModuleWrapper.svelte b/frontend/src/lib/components/flows/content/FlowModuleWrapper.svelte index d1f4a55f3c..95d1ce7cd8 100644 --- a/frontend/src/lib/components/flows/content/FlowModuleWrapper.svelte +++ b/frontend/src/lib/components/flows/content/FlowModuleWrapper.svelte @@ -21,6 +21,7 @@ import FlowWhileLoop from './FlowWhileLoop.svelte' import type { TriggerContext } from '$lib/components/triggers' import { formatCron } from '$lib/utils' + import AgentToolWrapper from './AgentToolWrapper.svelte' const { selectedId, flowStateStore } = getContext('FlowEditorContext') @@ -293,12 +294,16 @@ {/if} {/each} {:else if flowModule.value.type === 'aiagent'} - {#each flowModule.value.tools as _, index (index)} - t.id === $selectedId)} + {#if toolIndex !== -1} + - {/each} + {/if} {/if} diff --git a/frontend/src/lib/components/flows/content/McpToolEditor.svelte b/frontend/src/lib/components/flows/content/McpToolEditor.svelte new file mode 100644 index 0000000000..570b228693 --- /dev/null +++ b/frontend/src/lib/components/flows/content/McpToolEditor.svelte @@ -0,0 +1,194 @@ + + + + +
+ + + {#snippet children()} +

+ MCP clients allow AI agents to access and execute a list of tools made available by an MCP + server. +
+ Choose an MCP resource to make its tools available to the agent. +
+
+ Note: Only HTTP streamable MCP servers are supported. +

+ {/snippet} + + + +
+ +
+ + {#if tool.value.resource_path?.length > 0} + +
+ +
+ + +
+ {#snippet action()} + + {/snippet} +
+ {#if tools.error} +
+ {tools.error?.body?.message || + tools.error?.message || + 'Failed to load tools from MCP server'} +
+ {/if} +
+ {#if tools.status === 'loading'} +
Loading tools...
+ {:else if (tools.value ?? []).length === 0} +
+ {tools.error + ? 'Failed to load tools. Please check the resource path and try again.' + : 'No tools loaded yet. Click "Refresh Tools" to fetch tools from the MCP server.'} +
+ {:else} +
+ {#each tools.value ?? [] as tool} +
+ {tool.name} + {#if tool.description} + — {tool.description} + {/if} +
+ {/each} +
+ {/if} +
+
+
+ + + {#if tool.value.include_tools && tool.value.exclude_tools} +
+
+
+ +
+
+ +
+
+
+ {/if} + {/if} +
diff --git a/frontend/src/lib/components/flows/conversations/FlowChatInterface.svelte b/frontend/src/lib/components/flows/conversations/FlowChatInterface.svelte index 7b81e74604..2b1d776368 100644 --- a/frontend/src/lib/components/flows/conversations/FlowChatInterface.svelte +++ b/frontend/src/lib/components/flows/conversations/FlowChatInterface.svelte @@ -72,69 +72,67 @@ } -
-
- -
- {#if deploymentInProgress} - - {/if} - {#if manager.isLoadingMessages} -
- -
- {:else if manager.messages.length === 0} -
- -

Start a conversation

-

Send a message to run the flow and see the results

-
- {:else} -
- {#each manager.messages as message (message.id)} - - {/each} - {#if manager.isWaitingForResponse} -
- - Processing... -
- {/if} -
- {/if} -
+
+ +
+ {#if deploymentInProgress} + + {/if} + {#if manager.isLoadingMessages} +
+ +
+ {:else if manager.messages.length === 0} +
+ +

Start a conversation

+

Send a message to run the flow and see the results

+
+ {:else} +
+ {#each manager.messages as message (message.id)} + + {/each} + {#if manager.isWaitingForResponse} +
+ + Processing... +
+ {/if} +
+ {/if} +
- -
-
- -
-
+ +
+
+ +
+
diff --git a/frontend/src/lib/components/flows/conversations/FlowChatManager.svelte.ts b/frontend/src/lib/components/flows/conversations/FlowChatManager.svelte.ts index 6efb469ceb..7f98203995 100644 --- a/frontend/src/lib/components/flows/conversations/FlowChatManager.svelte.ts +++ b/frontend/src/lib/components/flows/conversations/FlowChatManager.svelte.ts @@ -380,9 +380,6 @@ class FlowChatManager { success } = this.parseStreamDeltas(data.new_result_stream) accumulatedContent += newContent - if (accumulatedContent.length > 0 || type === 'tool_result') { - this.isWaitingForResponse = false - } // Create tool message if type is tool_result if (type === 'tool_result') { diff --git a/frontend/src/lib/components/flows/conversations/FlowChatMessage.svelte b/frontend/src/lib/components/flows/conversations/FlowChatMessage.svelte index 9a143de1c0..ecb614321a 100644 --- a/frontend/src/lib/components/flows/conversations/FlowChatMessage.svelte +++ b/frontend/src/lib/components/flows/conversations/FlowChatMessage.svelte @@ -11,64 +11,59 @@ } let { message }: Props = $props() + + const messageClass = $derived.by(() => { + const base = 'max-w-[90%] min-w-0 rounded-lg w-fit' + if (message.message_type === 'user') { + return `${base} ml-auto bg-surface-secondary p-3` + } + return `${base} mr-auto bg-surface border ${message.success !== false ? 'border-gray-200 dark:border-gray-600' : '!border-red-500'}` + }) -
-
- {#if message.step_name} -
{message.step_name}
- {/if} +
+ {#if message.step_name} +
{message.step_name}
+ {/if} - {#if message.message_type === 'user'} -

{message.content}

- {:else if message.loading} -
- - Processing... -
- {:else if message.content} -
- {#if message.message_type === 'tool'} - {#if message.success !== false} - - {:else} - - {/if} + {#if message.message_type === 'user'} +

{message.content}

+ {:else if message.loading} +
+ + Processing... +
+ {:else if message.content} +
+ {#if message.message_type === 'tool'} + {#if message.success !== false} + + {:else} + {/if} -
- + -
+ } + ]} + />
- {:else} -

No result

- {/if} -
+
+ {:else} +

No result

+ {/if}
diff --git a/frontend/src/lib/components/flows/dfs.ts b/frontend/src/lib/components/flows/dfs.ts index 3e8709c206..d7627ae2ec 100644 --- a/frontend/src/lib/components/flows/dfs.ts +++ b/frontend/src/lib/components/flows/dfs.ts @@ -24,8 +24,8 @@ export function dfs( result = result.concat(dfs(branch, f, opts)) } } else if (module.value.type == 'aiagent' && !opts.skipToolNodes) { - result = result.concat(f(module, modules, [module.value.tools])) - result = result.concat(dfs(module.value.tools, f, opts)) + result = result.concat(f(module, modules, [module.value.tools as FlowModule[]])) + result = result.concat(dfs(module.value.tools as FlowModule[], f, opts)) } else { result.push(f(module, modules, [])) } diff --git a/frontend/src/lib/components/flows/map/FlowModuleSchemaMap.svelte b/frontend/src/lib/components/flows/map/FlowModuleSchemaMap.svelte index 4d6df0ef00..6fb0f8a10c 100644 --- a/frontend/src/lib/components/flows/map/FlowModuleSchemaMap.svelte +++ b/frontend/src/lib/components/flows/map/FlowModuleSchemaMap.svelte @@ -41,6 +41,7 @@ import { getStepHistoryLoaderContext } from '$lib/components/stepHistoryLoader.svelte' import { ModulesTestStates } from '$lib/components/modulesTest.svelte' import type { StateStore } from '$lib/utils' + import { type AgentTool, flowModuleToAgentTool, createMcpTool } from '../agentToolUtils' interface Props { sidebarSize?: number | undefined @@ -111,13 +112,14 @@ const { flowPropPickerConfig } = getContext('PropPickerContext') export async function insertNewModuleAtIndex( - modules: FlowModule[], + modules: FlowModule[] | AgentTool[], index: number, kind: InsertKind, wsScript?: { path: string; summary: string; hash: string | undefined }, wsFlow?: { path: string; summary: string }, - inlineScript?: InlineScript - ): Promise { + inlineScript?: InlineScript, + toolKind?: 'mcpTool' | 'flowmoduleTool' + ): Promise { push(history, flowStore.val) let module = emptyModule(flowStateStore.val, flowStore.val, kind == 'flow') let state = emptyFlowModuleState() @@ -167,8 +169,35 @@ } if (!modules) return [module] - modules.splice(index, 0, module) - return modules + + if (toolKind === 'mcpTool') { + // Create MCP AgentTool + const mcpTool = createMcpTool(module.id) + ;(modules as AgentTool[]).splice(index, 0, mcpTool) + return modules as AgentTool[] + } else if (toolKind === 'flowmoduleTool') { + // Create AgentTool from FlowModule + const agentTool = flowModuleToAgentTool(module) + ;(modules as AgentTool[]).splice(index, 0, agentTool) + return modules as AgentTool[] + } else { + // Standard FlowModule insertion (existing behavior) + modules.splice(index, 0, module) + return modules + } + } + + /** + * Helper function to remove an AgentTool by id from the tools array + * Tools are always leaf nodes, so we just need to delete their state directly + */ + function removeAgentToolById(tools: AgentTool[], id: string): AgentTool[] { + const index = tools.findIndex((tool) => tool.id == id) + if (index != -1) { + const [removed] = tools.splice(index, 1) + deleteFlowStateById(removed.id, flowStateStore) + } + return tools } export function removeAtId(modules: FlowModule[], id: string): FlowModule[] { @@ -194,7 +223,7 @@ }) mod.value.default = removeAtId(mod.value.default, id) } else if (mod.value.type == 'aiagent') { - mod.value.tools = removeAtId(mod.value.tools, id) + mod.value.tools = removeAgentToolById(mod.value.tools, id) } return mod }) @@ -489,6 +518,11 @@ } } else { const index = (detail.agentId ? targetModules?.length : detail.index) ?? 0 + const toolKind = detail.agentId + ? detail.kind === 'mcpTool' + ? 'mcpTool' + : 'flowmoduleTool' + : undefined await insertNewModuleAtIndex( targetModules, @@ -496,7 +530,8 @@ detail.kind, detail.script, detail.flow, - detail.inlineScript + detail.inlineScript, + toolKind ) const id = targetModules[index].id $selectedId = id diff --git a/frontend/src/lib/components/flows/map/InsertModuleInner.svelte b/frontend/src/lib/components/flows/map/InsertModuleInner.svelte index a02b632186..ff7fb95a6d 100644 --- a/frontend/src/lib/components/flows/map/InsertModuleInner.svelte +++ b/frontend/src/lib/components/flows/map/InsertModuleInner.svelte @@ -18,7 +18,7 @@ disableAi?: boolean kind?: 'script' | 'trigger' | 'preprocessor' | 'failure' allowTrigger?: boolean - scriptOnly?: boolean + toolMode?: boolean } let { @@ -27,7 +27,7 @@ disableAi = false, kind = 'script', allowTrigger = true, - scriptOnly = false + toolMode = false }: Props = $props() let customUi: undefined | FlowBuilderWhitelabelCustomUi = getContext('customUi') @@ -73,86 +73,96 @@
- {#if kind === 'script' && !scriptOnly} + {#if kind === 'script'}
{ + onSelect={() => { selectedKind = 'script' }} /> - {#if customUi?.triggers != false && allowTrigger} + {#if toolMode} { - selectedKind = 'trigger' - }} - /> - {/if} - { - selectedKind = 'approval' - }} - /> - {#if customUi?.flowNode != false} - { - selectedKind = 'flow' - }} - /> - {/if} - {#if stop} - { - selectedKind = 'script' - }} - /> - {/if} - - { - dispatch('close') - dispatch('new', { kind: 'forloop' }) - }} - /> - { - dispatch('close') - dispatch('new', { kind: 'whileloop' }) - }} - /> - { - dispatch('close') - dispatch('new', { kind: 'branchone' }) - }} - /> - { - dispatch('close') - dispatch('new', { kind: 'branchall' }) - }} - /> - {#if customUi?.aiAgent != false} - { + label="MCP" + onSelect={() => { + dispatch('pickMcpTool') dispatch('close') - dispatch('new', { kind: 'aiagent' }) }} /> + {:else} + {#if customUi?.triggers != false && allowTrigger} + { + selectedKind = 'trigger' + }} + /> + {/if} + { + selectedKind = 'approval' + }} + /> + {#if customUi?.flowNode != false} + { + selectedKind = 'flow' + }} + /> + {/if} + {#if stop} + { + selectedKind = 'script' + }} + /> + {/if} + + { + dispatch('close') + dispatch('new', { kind: 'forloop' }) + }} + /> + { + dispatch('close') + dispatch('new', { kind: 'whileloop' }) + }} + /> + { + dispatch('close') + dispatch('new', { kind: 'branchone' }) + }} + /> + { + dispatch('close') + dispatch('new', { kind: 'branchall' }) + }} + /> + {#if customUi?.aiAgent != false} + { + dispatch('close') + dispatch('new', { kind: 'aiagent' }) + }} + /> + {/if} {/if}
{/if} diff --git a/frontend/src/lib/components/flows/pickers/FlowToplevelNode.svelte b/frontend/src/lib/components/flows/pickers/FlowToplevelNode.svelte index 0248437963..47127a5771 100644 --- a/frontend/src/lib/components/flows/pickers/FlowToplevelNode.svelte +++ b/frontend/src/lib/components/flows/pickers/FlowToplevelNode.svelte @@ -19,4 +19,4 @@ - + diff --git a/frontend/src/lib/components/flows/pickers/TopLevelNode.svelte b/frontend/src/lib/components/flows/pickers/TopLevelNode.svelte index ab6746515d..9aff3abba9 100644 --- a/frontend/src/lib/components/flows/pickers/TopLevelNode.svelte +++ b/frontend/src/lib/components/flows/pickers/TopLevelNode.svelte @@ -6,65 +6,70 @@ ChevronRight, Code, GitBranch, + Plug, Repeat, Square, Zap } from 'lucide-svelte' - import { createEventDispatcher } from 'svelte' import { twMerge } from 'tailwind-merge' + import type { ComponentType } from 'svelte' - export let label: string - export let selected = false - export let returnIcon = false - const dispatch = createEventDispatcher() + interface Props { + label: string + selected?: boolean + returnIcon?: boolean + onSelect: () => void + class?: string + } + + let { label, selected, returnIcon, onSelect, class: className }: Props = $props() + + interface IconConfig { + icon: ComponentType + showChevron?: boolean + iconClass?: string + } + + const iconMap: Record = { + Action: { icon: Code, showChevron: true }, + Trigger: { icon: Zap, showChevron: true }, + 'Approval/Prompt': { icon: CheckCircle2, showChevron: true }, + Flow: { icon: BarsStaggered as unknown as ComponentType, showChevron: true }, + 'End Flow': { icon: Square }, + 'For loop': { icon: Repeat }, + 'While loop': { icon: Repeat }, + 'Branch to one': { icon: GitBranch }, + 'Branch to all': { icon: GitBranch }, + 'AI Agent': { icon: BotIcon, iconClass: 'text-violet-800 dark:text-violet-400' }, + MCP: { icon: Plug, showChevron: true } + } + + const config = $derived(iconMap[label]) +{#snippet iconWithText(icon: ComponentType, showChevron = false, iconClass = '')} + {@const Icon = icon} + + {label} + {#if showChevron} + + {/if} +{/snippet} +