diff --git a/app/src/ai/agent_sdk/driver.rs b/app/src/ai/agent_sdk/driver.rs index ccd67074a1b..8a57576d2a4 100644 --- a/app/src/ai/agent_sdk/driver.rs +++ b/app/src/ai/agent_sdk/driver.rs @@ -119,6 +119,7 @@ mod error_classification; pub(crate) mod git_credentials; pub(crate) mod harness; mod harness_output_monitor; +pub(crate) mod managed_mcp_refresh; pub(super) mod output; mod snapshot; pub(crate) mod terminal; @@ -139,6 +140,7 @@ async fn with_credential_refreshes( git_task_id: Option, ai_client: Arc, oidc_strategy: Option<(String, String, String)>, + managed_mcp_refresh: Option, foreground: &ModelSpawner, ) -> T where @@ -162,12 +164,26 @@ where } .fuse(); + let managed_mcp_refresh = async move { + match managed_mcp_refresh { + Some(params) => managed_mcp_refresh::refresh_loop(params, foreground).await, + None => future::pending::<()>().await, + } + } + .fuse(); + let run_future = run_future.fuse(); - futures::pin_mut!(run_future, git_refresh, bedrock_refresh); + futures::pin_mut!( + run_future, + git_refresh, + bedrock_refresh, + managed_mcp_refresh + ); futures::select! { result = run_future => result, _ = git_refresh => unreachable!("git credentials refresh loop resolved unexpectedly"), _ = bedrock_refresh => unreachable!("Bedrock credentials refresh loop resolved unexpectedly"), + _ = managed_mcp_refresh => unreachable!("managed MCP refresh loop resolved unexpectedly"), } } @@ -870,6 +886,9 @@ struct ResolvedMcpSpecs { /// `createManagedMcpClientConfig`: the `uid` to re-mint the short-lived /// proxy config with, keyed by installation UUID. managed_uids: HashMap, + /// Expiry of the proxy config minted for each managed installation, + /// keyed by installation UUID; drives the proactive re-mint schedule. + managed_expirations: HashMap>, } impl From for AgentDriverError { @@ -1720,6 +1739,11 @@ impl AgentDriver { resolved .managed_uids .insert(installation.uuid(), uuid.to_string()); + if let Some(expires_at) = client_config.expires_at { + resolved + .managed_expirations + .insert(installation.uuid(), expires_at.utc()); + } } resolved.ephemeral_installations.extend(installations); } @@ -1764,6 +1788,11 @@ impl AgentDriver { resolved .managed_uids .insert(installation.uuid(), id.clone()); + if let Some(expires_at) = client_config.expires_at { + resolved + .managed_expirations + .insert(installation.uuid(), expires_at.utc()); + } } resolved.ephemeral_installations.extend(installations); } @@ -2288,36 +2317,52 @@ impl AgentDriver { .await .map_err(|err| format!("{err:#}"))?; - let server_name = stale.templatable_mcp_server().name.clone(); - let fresh = Self::installations_from_managed_client_config_json( + Self::rebuild_managed_installation( + &stale, &client_config.mcp_config_json, task_id, &managed_uid, + &secrets, ) - .map_err(|err| err.to_string())? - .into_iter() - .find(|candidate| candidate.templatable_mcp_server().name == server_name) - .ok_or_else(|| { - format!("server '{server_name}' missing from re-minted managed MCP config") - })?; - - // Preserve the stale instance's identities: the installation - // UUID keys all manager state (and is random for local runs, - // so re-parsing yields a different one), and the template - // UUID keys the log file. - let mut fresh_server = fresh.templatable_mcp_server().clone(); - fresh_server.uuid = stale.template_uuid(); - let mut rebuilt = TemplatableMCPServerInstallation::new( - stale.uuid(), - fresh_server, - fresh.variable_values().clone(), - ); - rebuilt.apply_secrets(&secrets); - Ok(rebuilt) }) }) } + /// Rebuilds a managed installation from a freshly minted client config, + /// preserving the stale instance's identities: the installation UUID + /// keys all manager state (and is random for local runs, so re-parsing + /// yields a different one), and the template UUID keys the log file. + pub(crate) fn rebuild_managed_installation( + stale: &TemplatableMCPServerInstallation, + mcp_config_json: &str, + task_id: Option, + managed_uid: &str, + secrets: &HashMap, + ) -> Result { + let server_name = stale.templatable_mcp_server().name.clone(); + let fresh = Self::installations_from_managed_client_config_json( + mcp_config_json, + task_id, + managed_uid, + ) + .map_err(|err| err.to_string())? + .into_iter() + .find(|candidate| candidate.templatable_mcp_server().name == server_name) + .ok_or_else(|| { + format!("server '{server_name}' missing from re-minted managed MCP config") + })?; + + let mut fresh_server = fresh.templatable_mcp_server().clone(); + fresh_server.uuid = stale.template_uuid(); + let mut rebuilt = TemplatableMCPServerInstallation::new( + stale.uuid(), + fresh_server, + fresh.variable_values().clone(), + ); + rebuilt.apply_secrets(secrets); + Ok(rebuilt) + } + /// Subscribe to [`FileBasedMCPManagerEvent::CloudEnvMcpScanComplete`] /// paths and return a receiver that fires with auto-start-requested server UUIDs once every repo /// reports in. Must be called **before** `prepare_environment` so no events are missed. @@ -2828,6 +2873,7 @@ impl AgentDriver { task_id_for_refresh, ai_client_for_refresh, oidc_strategy_for_refresh, + managed_mcp_refresh_params, ) = async { let (setup_events, environment_snapshot_reporter) = foreground .spawn(|me, ctx| { @@ -2894,6 +2940,10 @@ impl AgentDriver { .await?; // For the Oz harness only: set up MCP servers, model overrides, and profile information. + let mut managed_mcp_refresh_params: Option = + None; + let mut managed_refresh_entries: Vec = + Vec::new(); if matches!(&task.harness, HarnessKind::Oz) { let mcp_specs = task.mcp_specs.clone(); let managed_mcp_client = foreground @@ -2908,6 +2958,19 @@ impl AgentDriver { let existing_uuids = resolved_mcp_specs.local_uuids; let mut ephemeral_installations = resolved_mcp_specs.ephemeral_installations; let managed_uids = resolved_mcp_specs.managed_uids; + let managed_expirations = resolved_mcp_specs.managed_expirations; + managed_refresh_entries = managed_uids + .iter() + .filter_map(|(installation_uuid, managed_uid)| { + managed_expirations.get(installation_uuid).map(|expires_at| { + managed_mcp_refresh::ManagedRefreshEntry { + installation_uuid: *installation_uuid, + managed_uid: managed_uid.clone(), + expires_at: *expires_at, + } + }) + }) + .collect(); // Attach the built-in Factory MCP server. Interactive // clients attach built-ins via @@ -3002,6 +3065,26 @@ impl AgentDriver { }) .await; Self::handle_mcp_startup_result(mcp_startup_result, &foreground).await?; + + // Proactively re-mint managed proxy configs before they expire. + if FeatureFlag::McpSelfHeal.is_enabled() && !managed_refresh_entries.is_empty() { + let (managed_mcp_client, task_id, secrets) = foreground + .spawn(|me, ctx| { + ( + ServerApiProvider::as_ref(ctx).get_managed_mcp_client(), + me.task_id, + me.secrets.clone(), + ) + }) + .await?; + managed_mcp_refresh_params = + Some(managed_mcp_refresh::ManagedMcpRefreshParams { + entries: std::mem::take(&mut managed_refresh_entries), + managed_mcp_client, + task_id, + secrets, + }); + } let profile = task.profile.clone(); setup_events .record_result(SetupStep::AgentProfileConfiguration, async { @@ -3259,6 +3342,7 @@ impl AgentDriver { task_id_for_refresh, ai_client_for_refresh, oidc_strategy_for_refresh, + managed_mcp_refresh_params, )) } .instrument(setup_span) @@ -3284,6 +3368,7 @@ impl AgentDriver { task_id_for_refresh, ai_client_for_refresh, oidc_strategy_for_refresh, + managed_mcp_refresh_params, &foreground, ) .await?; @@ -3337,6 +3422,7 @@ impl AgentDriver { task_id_for_refresh, ai_client_for_refresh, oidc_strategy_for_refresh, + managed_mcp_refresh_params, &foreground, ) .await diff --git a/app/src/ai/agent_sdk/driver/managed_mcp_refresh.rs b/app/src/ai/agent_sdk/driver/managed_mcp_refresh.rs new file mode 100644 index 00000000000..1949c5ab252 --- /dev/null +++ b/app/src/ai/agent_sdk/driver/managed_mcp_refresh.rs @@ -0,0 +1,207 @@ +//! Proactive re-minting of managed MCP proxy configs during agent runs. +//! +//! Managed (backend-proxied) MCP servers authenticate with a short-lived +//! proxy session token minted once at run startup. External sessions expire +//! after a few hours, after which every tool call fails until a new token is +//! minted. This loop re-mints each managed server's config shortly before it +//! expires and respawns the server with the fresh token, so long runs never +//! hit the expiry cliff. The reactive path (re-mint on a 401 during +//! reconnect) remains the backstop when this loop can't help. +//! +//! Like the git/Bedrock credential loops, [`refresh_loop`] never resolves on +//! its own — it is raced against the harness future via `futures::select!` +//! in `with_credential_refreshes` and dropped when the run finishes. + +use std::collections::HashMap; +use std::sync::Arc; +use std::time::Duration; + +use chrono::{DateTime, Utc}; +use uuid::Uuid; +use warp_managed_secrets::ManagedSecretValue; +use warpui::{ModelSpawner, SingletonEntity as _}; + +use super::AgentDriver; +use crate::ai::agent_sdk::retry::{is_transient_graphql_or_http_error, with_bounded_retry_using}; +use crate::ai::ambient_agents::AmbientAgentTaskId; +use crate::ai::mcp::TemplatableMCPServerManager; +use crate::ai::mcp::parsing::resolve_json; +use crate::server::server_api::managed_mcp::ManagedMcpClient; + +/// Re-mint this long before a token's expiry (mirrors warp-server's own +/// refresh lead time for downstream OAuth tokens). +const REFRESH_LEAD: Duration = Duration::from_secs(5 * 60); +/// Never spin faster than this, even for tokens that are about to expire. +const MIN_DELAY: Duration = Duration::from_secs(60); +/// Attempt budget for a proactive re-mint; off the tool-call latency path, +/// so it gets the same patience as run-startup resolution. +const REMINT_MAX_ATTEMPTS: usize = 6; + +/// One managed server's re-mint schedule entry. +pub(crate) struct ManagedRefreshEntry { + pub installation_uuid: Uuid, + pub managed_uid: String, + pub expires_at: DateTime, +} + +/// Everything the proactive re-mint loop needs, captured at run startup. +pub(crate) struct ManagedMcpRefreshParams { + pub entries: Vec, + pub managed_mcp_client: Arc, + pub task_id: Option, + pub secrets: Arc>, +} + +enum RefreshOutcome { + /// Token re-minted (and the server respawned if the config changed); + /// schedule the next refresh at the new expiry. + Rescheduled(DateTime), + /// Stop scheduling this server; the reason is logged by the caller. + Done(String), +} + +/// Never-resolving loop that re-mints managed proxy configs before expiry. +pub(crate) async fn refresh_loop( + params: ManagedMcpRefreshParams, + foreground: &ModelSpawner, +) { + let ManagedMcpRefreshParams { + mut entries, + managed_mcp_client, + task_id, + secrets, + } = params; + + loop { + let Some(next_expiry) = entries.iter().map(|entry| entry.expires_at).min() else { + // Nothing left to schedule; park until the run ends. + return futures::future::pending::<()>().await; + }; + warpui::r#async::Timer::after(refresh_delay(next_expiry, Utc::now())).await; + + let now = Utc::now(); + let mut index = 0; + while index < entries.len() { + if !is_due(entries[index].expires_at, now) { + index += 1; + continue; + } + let entry = &mut entries[index]; + match refresh_entry(entry, &managed_mcp_client, task_id, &secrets, foreground).await { + RefreshOutcome::Rescheduled(new_expiry) => { + log::info!( + "Re-minted managed MCP config '{}'; next refresh before {new_expiry}", + entry.managed_uid + ); + entry.expires_at = new_expiry; + index += 1; + } + RefreshOutcome::Done(reason) => { + log::info!( + "Stopping proactive re-mint of managed MCP config '{}': {reason}", + entry.managed_uid + ); + entries.swap_remove(index); + } + } + } + } +} + +/// How long to sleep before the next re-mint pass: wake [`REFRESH_LEAD`] +/// before the earliest expiry, but never spin faster than [`MIN_DELAY`]. +fn refresh_delay(next_expiry: DateTime, now: DateTime) -> Duration { + (next_expiry - now) + .to_std() + .unwrap_or_default() + .saturating_sub(REFRESH_LEAD) + .max(MIN_DELAY) +} + +/// Whether a token is inside its refresh window. +fn is_due(expires_at: DateTime, now: DateTime) -> bool { + expires_at - now <= chrono::Duration::from_std(REFRESH_LEAD).unwrap_or_default() +} + +async fn refresh_entry( + entry: &ManagedRefreshEntry, + managed_mcp_client: &Arc, + task_id: Option, + secrets: &Arc>, + foreground: &ModelSpawner, +) -> RefreshOutcome { + let installation_uuid = entry.installation_uuid; + let managed_uid = entry.managed_uid.clone(); + + // The retained config is the current source of truth for this server + // (a reactive re-mint may have already replaced the startup config). + let stale = match foreground + .spawn(move |_, ctx| { + TemplatableMCPServerManager::as_ref(ctx).retained_installation(installation_uuid) + }) + .await + { + Ok(Some(stale)) => stale, + Ok(None) => return RefreshOutcome::Done("server no longer active".to_string()), + Err(_) => return RefreshOutcome::Done("driver shutting down".to_string()), + }; + + let client_config = match with_bounded_retry_using( + &format!("proactively re-mint managed MCP config '{managed_uid}'"), + REMINT_MAX_ATTEMPTS, + is_transient_graphql_or_http_error, + || managed_mcp_client.create_managed_mcp_client_config(managed_uid.clone()), + ) + .await + { + Ok(client_config) => client_config, + Err(err) => { + // The reactive reconnect path remains as the backstop. + return RefreshOutcome::Done(format!("re-mint failed: {err:#}")); + } + }; + + let Some(new_expiry) = client_config.expires_at.map(|time| time.utc()) else { + return RefreshOutcome::Done("re-minted config has no expiry".to_string()); + }; + if new_expiry <= entry.expires_at { + // Runtime-kind tokens are anchored to the task's start, so a re-mint + // cannot extend the deadline — the sandbox ends then anyway. + return RefreshOutcome::Done( + "token deadline is fixed (runtime-anchored); nothing to extend".to_string(), + ); + } + + let fresh = match AgentDriver::rebuild_managed_installation( + &stale, + &client_config.mcp_config_json, + task_id, + &managed_uid, + secrets, + ) { + Ok(fresh) => fresh, + Err(err) => return RefreshOutcome::Done(format!("rebuild failed: {err}")), + }; + + // Respawn only when the rendered config actually changed; respawning + // drops the live MCP session, and an interrupted tool call relies on + // the reconnect retry to recover. + if resolve_json(&fresh) != resolve_json(&stale) { + let respawn = foreground + .spawn(move |_, ctx| { + TemplatableMCPServerManager::handle(ctx).update(ctx, move |manager, ctx| { + manager.respawn_with_installation(fresh, ctx); + }); + }) + .await; + if respawn.is_err() { + return RefreshOutcome::Done("driver shutting down".to_string()); + } + } + + RefreshOutcome::Rescheduled(new_expiry) +} + +#[cfg(test)] +#[path = "managed_mcp_refresh_tests.rs"] +mod tests; diff --git a/app/src/ai/agent_sdk/driver/managed_mcp_refresh_tests.rs b/app/src/ai/agent_sdk/driver/managed_mcp_refresh_tests.rs new file mode 100644 index 00000000000..4f9107c5582 --- /dev/null +++ b/app/src/ai/agent_sdk/driver/managed_mcp_refresh_tests.rs @@ -0,0 +1,39 @@ +//! Tests for the proactive managed-MCP re-mint schedule math. + +use chrono::{TimeDelta, Utc}; + +use super::{MIN_DELAY, REFRESH_LEAD, is_due, refresh_delay}; + +#[test] +fn wakes_one_lead_before_the_earliest_expiry() { + let now = Utc::now(); + let expiry = now + TimeDelta::hours(3); + let delay = refresh_delay(expiry, now); + let expected = std::time::Duration::from_secs(3 * 60 * 60) - REFRESH_LEAD; + // Allow a little slack for the sub-second remainder of `now`. + assert!( + delay >= expected - std::time::Duration::from_secs(1) && delay <= expected, + "unexpected delay: {delay:?}" + ); +} + +#[test] +fn never_spins_faster_than_the_minimum_delay() { + let now = Utc::now(); + // Expiry inside the lead window, at it, and in the past all clamp. + for offset in [ + TimeDelta::minutes(4), + TimeDelta::zero(), + TimeDelta::minutes(-10), + ] { + assert_eq!(refresh_delay(now + offset, now), MIN_DELAY); + } +} + +#[test] +fn due_exactly_within_the_lead_window() { + let now = Utc::now(); + assert!(is_due(now + TimeDelta::minutes(4), now)); + assert!(is_due(now - TimeDelta::minutes(1), now)); + assert!(!is_due(now + TimeDelta::minutes(6), now)); +} diff --git a/app/src/ai/agent_sdk/driver_tests.rs b/app/src/ai/agent_sdk/driver_tests.rs index 96c2197e47f..c90dfb5d543 100644 --- a/app/src/ai/agent_sdk/driver_tests.rs +++ b/app/src/ai/agent_sdk/driver_tests.rs @@ -282,6 +282,35 @@ fn managed_resolution_records_managed_uids_for_remint() { } } +#[test] +fn managed_resolution_records_expirations_for_proactive_refresh() { + let uuid = uuid::Uuid::parse_str("550e8400-e29b-41d4-a716-446655440000").unwrap(); + let config_json = r#"{"mcpServers":{"linear":{"url":"https://app.warp.dev/mcp/proxy/abc","headers":{"Authorization":"Bearer tok"}}}}"#; + let expires_at = chrono::Utc::now() + chrono::TimeDelta::hours(3); + let mut mock = MockManagedMcpClient::new(); + mock.expect_create_managed_mcp_client_config() + .times(1) + .returning(move |_| { + let mut output = managed_client_config_output(config_json); + output.expires_at = Some(warp_graphql::scalars::Time::new(expires_at)); + Ok(output) + }); + + let resolved = block_on(AgentDriver::resolve_mcp_specs_with_local_uuids( + &[MCPSpec::Uuid(uuid)], + &HashSet::new(), + Arc::new(mock), + None, + )) + .unwrap(); + + let installation_uuid = resolved.ephemeral_installations[0].uuid(); + assert_eq!( + resolved.managed_expirations.get(&installation_uuid), + Some(&expires_at) + ); +} + #[test] fn managed_refresher_preserves_identity_and_swaps_token() { let old_json = r#"{"mcpServers":{"linear":{"url":"https://app.warp.dev/mcp/proxy/abc","headers":{"Authorization":"Bearer old-token"}}}}"#; diff --git a/app/src/ai/mcp/templatable_manager/native.rs b/app/src/ai/mcp/templatable_manager/native.rs index e843342bdc3..07990b707ce 100644 --- a/app/src/ai/mcp/templatable_manager/native.rs +++ b/app/src/ai/mcp/templatable_manager/native.rs @@ -2169,6 +2169,59 @@ impl TemplatableMCPServerManager { } } + /// Returns the installation a server was last spawned with, if retained. + pub(crate) fn retained_installation( + &self, + installation_uuid: Uuid, + ) -> Option { + self.spawn_configs + .get(&installation_uuid) + .map(|config| config.installation.clone()) + } + + /// Respawns a running server with a pre-refreshed installation (e.g. a + /// managed server whose proxy token was proactively re-minted). Unlike + /// `reconnect_server` this does not invoke the retained refresher — the + /// caller already holds the fresh config — and it defers to any reactive + /// reconnect already in flight instead of double-spawning. + pub(crate) fn respawn_with_installation( + &mut self, + installation: TemplatableMCPServerInstallation, + ctx: &mut ModelContext, + ) { + let installation_uuid = installation.uuid(); + if self.pending_reconnections.contains_key(&installation_uuid) { + log::debug!( + "Skipping proactive respawn of {installation_uuid}: a reconnect is in flight" + ); + return; + } + let Some(provenance) = self + .spawn_configs + .get(&installation_uuid) + .map(|config| config.provenance.clone()) + else { + log::debug!("Skipping proactive respawn of {installation_uuid}: no retained config"); + return; + }; + + log::info!("Respawning MCP server {installation_uuid} with a refreshed config"); + if let Some(server_info) = self.active_servers.remove(&installation_uuid) { + ctx.spawn(server_info.shutdown(), |_, _, _| {}); + } + if let Some(spawned_info) = self.spawned_servers.remove(&installation_uuid) { + spawned_info.abort_handle.abort(); + } + if let Some(logger) = self.server_loggers.remove(&installation_uuid) { + logger.close(); + } + self.pending_oauth_csrf + .retain(|_, v| *v != installation_uuid); + self.authorization_urls.remove(&installation_uuid); + + self.spawn_server_impl(installation, provenance, SpawnMode::Reconnect, ctx); + } + /// Records a failed reconnect attempt and blocks further attempts for a /// jittered, exponentially growing interval (500ms doubling to a ~32s /// cap, mirroring `retry_strategies`). Cleared on the next successful