From 56ede303ed97b874530fa1df3a303d2fc5012efc Mon Sep 17 00:00:00 2001 From: Ruben Fiszel Date: Mon, 28 Apr 2025 19:16:04 +0200 Subject: [PATCH] fix: improve mcp mode api --- .../20250428170426_mcp_mode_log.down.sql | 1 + .../20250428170426_mcp_mode_log.up.sql | 2 + backend/src/main.rs | 8 +- backend/src/monitor.rs | 118 ++++++++++++------ backend/windmill-api/src/agent_workers_ee.rs | 4 +- backend/windmill-api/src/lib.rs | 89 ++++++------- 6 files changed, 126 insertions(+), 96 deletions(-) create mode 100644 backend/migrations/20250428170426_mcp_mode_log.down.sql create mode 100644 backend/migrations/20250428170426_mcp_mode_log.up.sql diff --git a/backend/migrations/20250428170426_mcp_mode_log.down.sql b/backend/migrations/20250428170426_mcp_mode_log.down.sql new file mode 100644 index 0000000000..d2f607c5b8 --- /dev/null +++ b/backend/migrations/20250428170426_mcp_mode_log.down.sql @@ -0,0 +1 @@ +-- Add down migration script here diff --git a/backend/migrations/20250428170426_mcp_mode_log.up.sql b/backend/migrations/20250428170426_mcp_mode_log.up.sql new file mode 100644 index 0000000000..038bb8234a --- /dev/null +++ b/backend/migrations/20250428170426_mcp_mode_log.up.sql @@ -0,0 +1,2 @@ +-- Add up migration script here +ALTER TYPE log_mode ADD VALUE 'mcp'; \ No newline at end of file diff --git a/backend/src/main.rs b/backend/src/main.rs index b2d77f352b..7226d4a3af 100644 --- a/backend/src/main.rs +++ b/backend/src/main.rs @@ -318,7 +318,7 @@ async fn windmill_main() -> anyhow::Result<()> { .ok() .and_then(|x| x.parse::().ok()) .unwrap_or(false) - && (mode == Mode::Server || mode == Mode::Standalone || mode == Mode::MCP); + && (mode == Mode::Server || mode == Mode::Standalone); let indexer_mode = mode == Mode::Indexer; let mcp_mode = mode == Mode::MCP; @@ -446,7 +446,7 @@ Windmill Community Edition {GIT_VERSION} if !valid_key && !server_mode { tracing::error!("Invalid license key, workers require a valid license key"); } - if server_mode { + if server_mode || mcp_mode { if let Some(db) = conn.as_sql() { // only force renewal if invalid but not empty (= expired) let renewed_now = maybe_renew_license_key_on_start( @@ -466,10 +466,10 @@ Windmill Community Edition {GIT_VERSION} } } - if server_mode || worker_mode || indexer_mode { + if server_mode || worker_mode || indexer_mode || mcp_mode { let port_var = std::env::var("PORT").ok().and_then(|x| x.parse().ok()); - let port = if server_mode || indexer_mode { + let port = if server_mode || indexer_mode || mcp_mode { port_var.unwrap_or(DEFAULT_PORT as u16) } else { port_var.unwrap_or(0) diff --git a/backend/src/monitor.rs b/backend/src/monitor.rs index e8bb9be822..7366cb4f53 100644 --- a/backend/src/monitor.rs +++ b/backend/src/monitor.rs @@ -34,8 +34,12 @@ use windmill_common::ee::{jobs_waiting_alerts, worker_groups_alerts}; #[cfg(feature = "oauth2")] use windmill_common::global_settings::OAUTH_SETTING; use windmill_common::{ - utils::empty_string_as_none, - agent_workers::DECODED_AGENT_TOKEN, auth::create_token_for_owner, ee::CriticalErrorChannel, error, flow_status::{FlowStatus, FlowStatusModule}, global_settings::{ + agent_workers::DECODED_AGENT_TOKEN, + auth::create_token_for_owner, + ee::CriticalErrorChannel, + error, + flow_status::{FlowStatus, FlowStatusModule}, + global_settings::{ BASE_URL_SETTING, BUNFIG_INSTALL_SCOPES_SETTING, CRITICAL_ALERT_MUTE_UI_SETTING, CRITICAL_ERROR_CHANNELS_SETTING, DEFAULT_TAGS_PER_WORKSPACE_SETTING, DEFAULT_TAGS_WORKSPACES_SETTING, EXPOSE_DEBUG_METRICS_SETTING, EXPOSE_METRICS_SETTING, @@ -45,13 +49,32 @@ use windmill_common::{ NUGET_CONFIG_SETTING, OTEL_SETTING, PIP_INDEX_URL_SETTING, REQUEST_SIZE_LIMIT_SETTING, REQUIRE_PREEXISTING_USER_FOR_OAUTH_SETTING, RETENTION_PERIOD_SECS_SETTING, SAML_METADATA_SETTING, SCIM_TOKEN_SETTING, TIMEOUT_WAIT_RESULT_SETTING, - }, indexer::load_indexer_config, jobs::QueuedJob, jwt::JWT_SECRET, oauth2::REQUIRE_PREEXISTING_USER_FOR_OAUTH, server::load_smtp_config, tracing_init::JSON_FMT, users::truncate_token, utils::{now_from_db, rd_string, report_critical_error, Mode}, worker::{ - load_worker_config, reload_custom_tags_setting, store_pull_query, store_suspended_pull_query, update_min_version, Connection, DEFAULT_TAGS_PER_WORKSPACE, DEFAULT_TAGS_WORKSPACES, INDEXER_CONFIG, SCRIPT_TOKEN_EXPIRY, SMTP_CONFIG, TMP_DIR, WORKER_CONFIG, WORKER_GROUP - }, KillpillSender, BASE_URL, CRITICAL_ALERT_MUTE_UI_ENABLED, CRITICAL_ERROR_CHANNELS, DB, DEFAULT_HUB_BASE_URL, HUB_BASE_URL, JOB_RETENTION_SECS, METRICS_DEBUG_ENABLED, METRICS_ENABLED, MONITOR_LOGS_ON_OBJECT_STORE, OTEL_LOGS_ENABLED, OTEL_METRICS_ENABLED, OTEL_TRACING_ENABLED, SERVICE_LOG_RETENTION_SECS, + }, + indexer::load_indexer_config, + jobs::QueuedJob, + jwt::JWT_SECRET, + oauth2::REQUIRE_PREEXISTING_USER_FOR_OAUTH, + server::load_smtp_config, + tracing_init::JSON_FMT, + users::truncate_token, + utils::empty_string_as_none, + utils::{now_from_db, rd_string, report_critical_error, Mode}, + worker::{ + load_worker_config, reload_custom_tags_setting, store_pull_query, + store_suspended_pull_query, update_min_version, Connection, DEFAULT_TAGS_PER_WORKSPACE, + DEFAULT_TAGS_WORKSPACES, INDEXER_CONFIG, SCRIPT_TOKEN_EXPIRY, SMTP_CONFIG, TMP_DIR, + WORKER_CONFIG, WORKER_GROUP, + }, + KillpillSender, BASE_URL, CRITICAL_ALERT_MUTE_UI_ENABLED, CRITICAL_ERROR_CHANNELS, DB, + DEFAULT_HUB_BASE_URL, HUB_BASE_URL, JOB_RETENTION_SECS, METRICS_DEBUG_ENABLED, METRICS_ENABLED, + MONITOR_LOGS_ON_OBJECT_STORE, OTEL_LOGS_ENABLED, OTEL_METRICS_ENABLED, OTEL_TRACING_ENABLED, + SERVICE_LOG_RETENTION_SECS, }; use windmill_queue::{cancel_job, MiniPulledJob, SameWorkerPayload}; use windmill_worker::{ - handle_job_error, AuthedClient, JobCompletedSender, SameWorkerSender, BUNFIG_INSTALL_SCOPES, INSTANCE_PYTHON_VERSION, JOB_DEFAULT_TIMEOUT, KEEP_JOB_DIR, MAVEN_REPOS, NO_DEFAULT_MAVEN, NPM_CONFIG_REGISTRY, NUGET_CONFIG, PIP_EXTRA_INDEX_URL, PIP_INDEX_URL + handle_job_error, AuthedClient, JobCompletedSender, SameWorkerSender, BUNFIG_INSTALL_SCOPES, + INSTANCE_PYTHON_VERSION, JOB_DEFAULT_TIMEOUT, KEEP_JOB_DIR, MAVEN_REPOS, NO_DEFAULT_MAVEN, + NPM_CONFIG_REGISTRY, NUGET_CONFIG, PIP_EXTRA_INDEX_URL, PIP_INDEX_URL, }; #[cfg(feature = "parquet")] @@ -135,7 +158,6 @@ pub async fn initial_load( } } - if let Err(e) = load_metrics_enabled(conn).await { tracing::error!("Error loading expose metrics: {e:#}"); } @@ -172,12 +194,14 @@ pub async fn initial_load( } Connection::Http(_) => { // TODO: reload worker config from http - WORKER_CONFIG.write().await.worker_tags = DECODED_AGENT_TOKEN.as_ref().map(|x| x.tags.clone()).unwrap_or_default(); + WORKER_CONFIG.write().await.worker_tags = DECODED_AGENT_TOKEN + .as_ref() + .map(|x| x.tags.clone()) + .unwrap_or_default(); } } } - if let Err(e) = reload_hub_base_url_setting(conn, server_mode).await { tracing::error!("Error reloading hub base url: {:?}", e) } @@ -190,7 +214,6 @@ pub async fn initial_load( if let Err(e) = reload_custom_tags_setting(db).await { tracing::error!("Error reloading custom tags: {:?}", e) } - } #[cfg(feature = "parquet")] @@ -225,7 +248,8 @@ pub async fn initial_load( } pub async fn load_metrics_enabled(conn: &Connection) -> error::Result<()> { - let metrics_enabled = load_value_from_global_settings_with_conn(conn, EXPOSE_METRICS_SETTING, true).await; + let metrics_enabled = + load_value_from_global_settings_with_conn(conn, EXPOSE_METRICS_SETTING, true).await; match metrics_enabled { Ok(Some(serde_json::Value::Bool(t))) => METRICS_ENABLED.store(t, Ordering::Relaxed), _ => (), @@ -345,13 +369,13 @@ pub async fn reload_critical_alert_mute_ui_setting(conn: &Connection) -> error:: load_value_from_global_settings_with_conn(conn, CRITICAL_ALERT_MUTE_UI_SETTING, true).await { CRITICAL_ALERT_MUTE_UI_ENABLED.store(t, Ordering::Relaxed); - } Ok(()) } pub async fn load_metrics_debug_enabled(conn: &Connection) -> error::Result<()> { - let metrics_enabled = load_value_from_global_settings_with_conn(conn, EXPOSE_DEBUG_METRICS_SETTING, true).await; + let metrics_enabled = + load_value_from_global_settings_with_conn(conn, EXPOSE_DEBUG_METRICS_SETTING, true).await; match metrics_enabled { Ok(Some(serde_json::Value::Bool(t))) => { METRICS_DEBUG_ENABLED.store(t, Ordering::Relaxed); @@ -575,7 +599,9 @@ async fn send_log_file_to_object_store( }; let exists = LAST_LOG_FILE_SENT.lock().map(|last_log_file_sent| { - last_log_file_sent.map(|last_log_file_sent| last_log_file_sent >= ts).unwrap_or(false) + last_log_file_sent + .map(|last_log_file_sent| last_log_file_sent >= ts) + .unwrap_or(false) }); if exists.unwrap_or(false) { @@ -1002,11 +1028,21 @@ pub async fn reload_nuget_config_setting(conn: &Connection) { .await; } pub async fn reload_maven_repos_setting(conn: &Connection) { - reload_option_setting_with_tracing(conn, windmill_common::global_settings::MAVEN_REPOS_SETTING, "MAVEN_REPOS", MAVEN_REPOS.clone()) - .await; + reload_option_setting_with_tracing( + conn, + windmill_common::global_settings::MAVEN_REPOS_SETTING, + "MAVEN_REPOS", + MAVEN_REPOS.clone(), + ) + .await; } pub async fn reload_no_default_maven_setting(conn: &Connection) { - let value = load_value_from_global_settings_with_conn(conn, windmill_common::global_settings::NO_DEFAULT_MAVEN_SETTING, true).await; + let value = load_value_from_global_settings_with_conn( + conn, + windmill_common::global_settings::NO_DEFAULT_MAVEN_SETTING, + true, + ) + .await; match value { Ok(Some(serde_json::Value::Bool(t))) => NO_DEFAULT_MAVEN.store(t, Ordering::Relaxed), Err(e) => { @@ -1175,7 +1211,6 @@ pub async fn load_value_from_global_settings( Ok(r) } - pub async fn load_value_from_global_settings_with_conn( conn: &Connection, setting_name: &str, @@ -1185,14 +1220,18 @@ pub async fn load_value_from_global_settings_with_conn( Connection::Sql(db) => Ok(load_value_from_global_settings(db, setting_name).await?), Connection::Http(client) => { if load_from_http { - client.get::>(&format!("/api/agent_workers/get_global_setting/{}", setting_name)).await - .map_err(|e| anyhow::anyhow!("Error loading setting {}: {}", setting_name, e)) + client + .get::>(&format!( + "/api/agent_workers/get_global_setting/{}", + setting_name + )) + .await + .map_err(|e| anyhow::anyhow!("Error loading setting {}: {}", setting_name, e)) } else { Ok(None) } } } - } pub async fn reload_option_setting( @@ -1314,12 +1353,12 @@ pub async fn monitor_db( let zombie_jobs_f = async { if server_mode && !initial_load && !*DISABLE_ZOMBIE_JOBS_MONITORING { if let Some(db) = conn.as_sql() { - handle_zombie_jobs(db, base_internal_url, "server").await; - match handle_zombie_flows(db).await { - Err(err) => { - tracing::error!("Error handling zombie flows: {:?}", err); - }, - _ => {} + handle_zombie_jobs(db, base_internal_url, "server").await; + match handle_zombie_flows(db).await { + Err(err) => { + tracing::error!("Error handling zombie flows: {:?}", err); + } + _ => {} } } } @@ -1327,7 +1366,7 @@ pub async fn monitor_db( let expired_items_f = async { if server_mode && !initial_load { if let Some(db) = conn.as_sql() { - delete_expired_items(&db).await; + delete_expired_items(&db).await; } } }; @@ -1496,11 +1535,7 @@ pub async fn reload_indexer_config(db: &Pool) { } } -pub async fn reload_worker_config( - db: &DB, - tx: KillpillSender, - kill_if_change: bool, -) { +pub async fn reload_worker_config(db: &DB, tx: KillpillSender, kill_if_change: bool) { let config = load_worker_config(db, tx.clone()).await; if let Err(e) = config { tracing::error!("Error reloading worker config: {:?}", e) @@ -1543,7 +1578,8 @@ pub async fn reload_worker_config( } pub async fn load_base_url(conn: &Connection) -> error::Result { - let q_base_url = load_value_from_global_settings_with_conn(conn, BASE_URL_SETTING, false).await?; + let q_base_url = + load_value_from_global_settings_with_conn(conn, BASE_URL_SETTING, false).await?; let std_base_url = std::env::var("BASE_URL") .ok() @@ -1574,10 +1610,9 @@ pub async fn load_base_url(conn: &Connection) -> error::Result { } pub async fn reload_base_url_setting(conn: &Connection) -> error::Result<()> { - #[cfg(feature = "oauth2")] let oauths = if let Some(db) = conn.as_sql() { - let q_oauth = load_value_from_global_settings (db, OAUTH_SETTING).await?; + let q_oauth = load_value_from_global_settings(db, OAUTH_SETTING).await?; if let Some(q) = q_oauth { if let Ok(v) = serde_json::from_value::< @@ -1863,7 +1898,8 @@ async fn handle_zombie_jobs(db: &Pool, base_internal_url: &str, worker mpsc::channel::(1); let same_worker_tx_never_used = SameWorkerSender(same_worker_tx_never_used, Arc::new(AtomicU16::new(0))); - let (send_result_never_used, _send_result_rx_never_used) = JobCompletedSender::new_never_used(); + let (send_result_never_used, _send_result_rx_never_used) = + JobCompletedSender::new_never_used(); let label = if job.permissioned_as != format!("u/{}", job.created_by) && job.permissioned_as != job.created_by @@ -2073,8 +2109,12 @@ async fn cancel_zombie_flow_job( Ok(()) } -pub async fn reload_hub_base_url_setting(conn: &Connection, server_mode: bool) -> error::Result<()> { - let hub_base_url = load_value_from_global_settings_with_conn(conn, HUB_BASE_URL_SETTING, true).await?; +pub async fn reload_hub_base_url_setting( + conn: &Connection, + server_mode: bool, +) -> error::Result<()> { + let hub_base_url = + load_value_from_global_settings_with_conn(conn, HUB_BASE_URL_SETTING, true).await?; let base_url = if let Some(q) = hub_base_url { if let Ok(v) = serde_json::from_value::(q.clone()) { diff --git a/backend/windmill-api/src/agent_workers_ee.rs b/backend/windmill-api/src/agent_workers_ee.rs index af31e61229..f3cd1f2474 100644 --- a/backend/windmill-api/src/agent_workers_ee.rs +++ b/backend/windmill-api/src/agent_workers_ee.rs @@ -22,7 +22,7 @@ pub fn workspaced_service( ) -> ( Router, Option>, - windmill_worker::JobCompletedSender, + Option, ) { use windmill_common::worker::Connection; use windmill_worker::JobCompletedSender; @@ -32,7 +32,7 @@ pub fn workspaced_service( let router = Router::new(); - (router, None, job_completed_tx) + (router, None, Some(job_completed_tx)) } #[derive(Clone, Debug, Deserialize, Serialize)] diff --git a/backend/windmill-api/src/lib.rs b/backend/windmill-api/src/lib.rs index 4245d0d736..b7de918407 100644 --- a/backend/windmill-api/src/lib.rs +++ b/backend/windmill-api/src/lib.rs @@ -411,7 +411,7 @@ pub async fn run_server( Router::new() }; - if !*CLOUD_HOSTED && server_mode { + if !*CLOUD_HOSTED && server_mode && !mcp_mode { #[cfg(feature = "websocket")] { let ws_killpill_rx = killpill_rx.resubscribe(); @@ -465,29 +465,35 @@ pub async fn run_server( .unwrap_or("localhost".to_string()); // Setup MCP server - #[cfg(feature = "mcp")] - let (mcp_sse_server, mcp_router) = setup_mcp_server(addr, "/api/mcp/w/:workspace_id")?; - #[cfg(feature = "mcp")] - let mcp_main_ct = mcp_sse_server.config.ct.clone(); // Token to signal shutdown *to* MCP - #[cfg(feature = "mcp")] - let mcp_service_ct = mcp_sse_server.with_service(McpRunner::new); // Token to wait for MCP *service* shutdown + #[allow(unused_variables)] + let (mcp_router, mcp_main_ct, mcp_service_ct) = { + #[cfg(feature = "mcp")] + if server_mode || mcp_mode { + let (mcp_sse_server, mcp_router) = setup_mcp_server(addr, "/api/mcp/w/:workspace_id")?; + #[cfg(feature = "mcp")] + let mcp_main_ct = mcp_sse_server.config.ct.clone(); // Token to signal shutdown *to* MCP + #[cfg(feature = "mcp")] + let mcp_service_ct = mcp_sse_server.with_service(McpRunner::new); // Token to wait for MCP *service* shutdown + (mcp_router, Some(mcp_main_ct), Some(mcp_service_ct)) + } else { + (Router::new(), None, None) + } + + #[cfg(not(feature = "mcp"))] + (Router::new(), None::<()>, None::<()>) + }; #[cfg(feature = "agent_worker_server")] let (agent_workers_router, agent_workers_bg_processor, agent_workers_killpill_tx) = - agent_workers_ee::workspaced_service(db.clone(), _base_internal_url.clone()); + if server_mode { + agent_workers_ee::workspaced_service(db.clone(), _base_internal_url.clone()) + } else { + (Router::new(), None, None) + }; #[cfg(feature = "agent_worker_server")] let agent_cache = Arc::new(AgentCache::new()); - // used on mcp mode only - #[cfg(feature = "mcp")] - let mcp_app = Router::new() - .nest("/api/mcp/w/:workspace_id", mcp_router.clone()) - .layer(from_extractor::()) - .layer(middleware_stack.clone()); - #[cfg(not(feature = "mcp"))] - let mcp_app = Router::new(); - // build our application with a route let app = Router::new() .nest( @@ -619,16 +625,7 @@ pub async fn run_server( .layer(from_extractor::()) .layer(cors.clone()), ) - .nest("/mcp/w/:workspace_id", { - #[cfg(feature = "mcp")] - { - mcp_router - } - #[cfg(not(feature = "mcp"))] - { - Router::new() - } - }) + .nest("/mcp/w/:workspace_id", mcp_router) .layer(from_extractor::()) .nest( "/w/:workspace_id/jobs_u", @@ -739,23 +736,7 @@ pub async fn run_server( ) }; - let mcp_app = if disable_response_logs { - mcp_app - } else { - mcp_app.layer( - TraceLayer::new_for_http() - .on_response(MyOnResponse {}) - .make_span_with(MyMakeSpan {}) - .on_request(()) - .on_failure(MyOnFailure {}), - ) - }; - - let server = if mcp_mode { - axum::serve(listener, mcp_app.into_make_service()) - } else { - axum::serve(listener, app.into_make_service()) - }; + let server = axum::serve(listener, app.into_make_service()); tracing::info!( instance = %*INSTANCE_NAME, @@ -771,18 +752,24 @@ pub async fn run_server( let server = server.with_graceful_shutdown(async move { killpill_rx.recv().await.ok(); #[cfg(feature = "agent_worker_server")] - if let Err(e) = agent_workers_killpill_tx.kill().await { - tracing::error!("Error killing agent workers: {e:#}"); + if let Some(agent_workers_killpill_tx) = agent_workers_killpill_tx { + if let Err(e) = agent_workers_killpill_tx.kill().await { + tracing::error!("Error killing agent workers: {e:#}"); + } } tracing::info!("Graceful shutdown of server"); #[cfg(feature = "mcp")] { - tracing::info!("Received shutdown signal, cancelling MCP server..."); - mcp_main_ct.cancel(); - tracing::info!("Waiting for MCP service cancellation..."); - mcp_service_ct.cancelled().await; - tracing::info!("MCP service cancelled."); + if let Some(mcp_main_ct) = mcp_main_ct { + tracing::info!("Received shutdown signal, cancelling MCP server..."); + mcp_main_ct.cancel(); + } + if let Some(mcp_service_ct) = mcp_service_ct { + tracing::info!("Waiting for MCP service cancellation..."); + mcp_service_ct.cancelled().await; + tracing::info!("MCP service cancelled."); + } } });