kalugin66/windmill
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: improve mcp mode api

Browse Source
This commit is contained in:
Ruben Fiszel
2025-04-28 19:16:04 +02:00
parent d3e187b0b8
commit 56ede303ed
6 changed files with 126 additions and 96 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
backend/migrations/20250428170426_mcp_mode_log.down.sql Normal file
View File

@@ -0,0 +1 @@
-- Add down migration script here

2
backend/migrations/20250428170426_mcp_mode_log.up.sql Normal file
View File

@@ -0,0 +1,2 @@
-- Add up migration script here
ALTER TYPE log_mode ADD VALUE 'mcp';

8
backend/src/main.rs
View File

@@ -318,7 +318,7 @@ async fn windmill_main() -> anyhow::Result<()> {
.ok()
.and_then(|x| x.parse::<bool>().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)

118
backend/src/monitor.rs
View File

@@ -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::<Option<serde_json::Value>>(&format!("/api/agent_workers/get_global_setting/{}", setting_name)).await
.map_err(|e| anyhow::anyhow!("Error loading setting {}: {}", setting_name, e))
client
.get::<Option<serde_json::Value>>(&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<T: FromStr + DeserializeOwned>(
@@ -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<Postgres>) {
}
}
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<String> {
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<String> {
}
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<Postgres>, base_internal_url: &str, worker
mpsc::channel::<SameWorkerPayload>(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::<String>(q.clone()) {

4
backend/windmill-api/src/agent_workers_ee.rs
View File

@@ -22,7 +22,7 @@ pub fn workspaced_service(
) -> (
Router,
Option<tokio::task::JoinHandle<()>>,
windmill_worker::JobCompletedSender,
Option<windmill_worker::JobCompletedSender>,
) {
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)]

89
backend/windmill-api/src/lib.rs
View File

@@ -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::<OptAuthed>())
.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::<OptAuthed>())
.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::<OptAuthed>())
.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.");
}
}
});