From e11d268e30f8ce9a161328b4bd4a4b5ad21d980b Mon Sep 17 00:00:00 2001 From: Bryan Helmkamp Date: Wed, 5 Aug 2026 08:43:41 -0400 Subject: [PATCH 1/2] Bound doctor diagnostics within client timeout --- docs/public/api-reference/fabro-api.yaml | 9 ++ lib/apps/fabro-server/src/diagnostics.rs | 118 +++++++++++++++--- .../fabro-server/src/server/handler/system.rs | 59 ++++++++- 3 files changed, 168 insertions(+), 18 deletions(-) diff --git a/docs/public/api-reference/fabro-api.yaml b/docs/public/api-reference/fabro-api.yaml index 261badbbb..4382d19fa 100644 --- a/docs/public/api-reference/fabro-api.yaml +++ b/docs/public/api-reference/fabro-api.yaml @@ -594,6 +594,15 @@ paths: application/json: schema: $ref: "#/components/schemas/DiagnosticsReport" + "504": + description: Diagnostics operation timed out + headers: + x-request-id: + $ref: "#/components/headers/XRequestId" + content: + application/json: + schema: + $ref: "#/components/schemas/ErrorResponse" /api/v1/openapi.json: get: diff --git a/lib/apps/fabro-server/src/diagnostics.rs b/lib/apps/fabro-server/src/diagnostics.rs index 988f84404..86f625366 100644 --- a/lib/apps/fabro-server/src/diagnostics.rs +++ b/lib/apps/fabro-server/src/diagnostics.rs @@ -6,7 +6,7 @@ use base64::Engine as _; use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use fabro_auth::auth_issue_message; use fabro_llm::client::Client as LlmClient; -use fabro_llm::model_test::{ModelTestStatus, run_basic_model_probe}; +use fabro_llm::model_test::{ModelTestOutcome, ModelTestStatus, run_basic_model_probe}; use fabro_model::{Catalog, ProviderId}; use fabro_redact::redact_string; use fabro_sandbox::{DockerSandboxProvider, daytona}; @@ -23,6 +23,9 @@ use tokio::time::timeout; use crate::server::AppState; +const EXTERNAL_SERVICE_PROBE_TIMEOUT: Duration = Duration::from_secs(15); +const DOCKER_PROBE_TIMEOUT: Duration = Duration::from_secs(5); + fn http_client_or_check( name: &str, status: CheckStatus, @@ -254,11 +257,38 @@ async fn probe_single_provider( }; let model_id = model.id.clone(); - let outcome = run_basic_model_probe(model_id.as_str(), &provider, client).await; + let outcome = run_basic_model_probe(model_id.as_str(), provider.clone(), client); + provider_probe_with_timeout( + provider, + model_id.to_string(), + outcome, + EXTERNAL_SERVICE_PROBE_TIMEOUT, + ) + .await +} + +async fn provider_probe_with_timeout( + provider: ProviderId, + model_id: String, + probe: F, + probe_timeout: Duration, +) -> ProviderProbeResult +where + F: Future, +{ + let Ok(outcome) = timeout(probe_timeout, probe).await else { + return provider_probe_error( + provider, + Some(model_id), + probe_timeout_message(probe_timeout), + None, + ); + }; + match outcome.status { ModelTestStatus::Ok => ProviderProbeResult { provider, - model_id: Some(model_id.to_string()), + model_id: Some(model_id), status: ProviderProbeStatus::Ok, error_message: None, diagnostic_detail: None, @@ -267,16 +297,19 @@ async fn probe_single_provider( let raw = outcome .error_message .unwrap_or_else(|| "provider probe failed".to_string()); - provider_probe_error( - provider, - Some(model_id.to_string()), - redact_string(&raw), - None, - ) + provider_probe_error(provider, Some(model_id), redact_string(&raw), None) } } } +fn probe_timeout_message(probe_timeout: Duration) -> String { + if probe_timeout.subsec_nanos() == 0 { + format!("timeout ({}s)", probe_timeout.as_secs()) + } else { + format!("timeout ({}ms)", probe_timeout.as_millis()) + } +} + fn provider_probe_error( provider: ProviderId, model_id: Option, @@ -388,7 +421,7 @@ async fn check_github_app(state: &AppState) -> CheckResult { Err(result) => return result, }; let probe = timeout( - Duration::from_secs(15), + EXTERNAL_SERVICE_PROBE_TIMEOUT, http.get(format!("{}/user", fabro_github::github_api_base_url())) .header("Authorization", format!("Bearer {token}")) .header("Accept", "application/vnd.github+json") @@ -538,7 +571,7 @@ async fn check_github_app(state: &AppState) -> CheckResult { Err(result) => return result, }; let auth_result = timeout( - Duration::from_secs(15), + EXTERNAL_SERVICE_PROBE_TIMEOUT, fabro_github::get_authenticated_app(&http, &jwt, &fabro_github::github_api_base_url()), ) .await; @@ -581,7 +614,7 @@ async fn check_docker_sandbox(state: &AppState) -> CheckResult { .await .map_err(|err| err.display_with_causes()) }, - Duration::from_secs(5), + DOCKER_PROBE_TIMEOUT, ) .await } @@ -656,7 +689,29 @@ async fn check_cloud_sandbox(state: &AppState) -> CheckResult { }; }; - match state.check_daytona_api_key(api_key).await { + check_cloud_sandbox_with_probe( + || state.check_daytona_api_key(api_key), + EXTERNAL_SERVICE_PROBE_TIMEOUT, + ) + .await +} + +async fn check_cloud_sandbox_with_probe(probe: F, probe_timeout: Duration) -> CheckResult +where + F: FnOnce() -> Fut, + Fut: Future>, +{ + let Ok(probe) = timeout(probe_timeout, probe()).await else { + return CheckResult { + name: "Cloud Sandbox".to_string(), + status: CheckStatus::Error, + summary: probe_timeout_message(probe_timeout), + details: vec![CheckDetail::new("Daytona probe timed out".to_string())], + remediation: Some("Verify DAYTONA_API_KEY value and Daytona reachability".to_string()), + }; + }; + + match probe { Ok(check) if check.ok() => CheckResult { name: "Cloud Sandbox".to_string(), status: CheckStatus::Pass, @@ -750,7 +805,7 @@ async fn check_brave_search(state: &AppState) -> CheckResult { Err(result) => return result, }; - let probe = timeout(Duration::from_secs(15), async move { + let probe = timeout(EXTERNAL_SERVICE_PROBE_TIMEOUT, async move { http.get("https://api.search.brave.com/res/v1/web/search?q=test&count=1") .header("X-Subscription-Token", api_key) .send() @@ -1035,6 +1090,27 @@ mod tests { ); } + #[tokio::test] + async fn provider_probe_reports_provider_specific_timeout() { + assert_eq!( + probe_timeout_message(EXTERNAL_SERVICE_PROBE_TIMEOUT), + "timeout (15s)" + ); + + let result = provider_probe_with_timeout( + ProviderId::new("modal"), + "modal/test-model".to_string(), + std::future::pending::(), + Duration::from_millis(1), + ) + .await; + + assert_eq!(result.provider, ProviderId::new("modal")); + assert_eq!(result.model_id.as_deref(), Some("modal/test-model")); + assert_eq!(result.status, ProviderProbeStatus::Error); + assert_eq!(result.error_message.as_deref(), Some("timeout (1ms)")); + } + #[test] fn docker_sandbox_probe_passes_when_daemon_responds() { let result = docker_sandbox_probe_check(Ok(())); @@ -1154,6 +1230,20 @@ enabled = false ); } + #[tokio::test] + async fn check_cloud_sandbox_reports_timeout() { + let result = check_cloud_sandbox_with_probe( + std::future::pending::>, + Duration::from_millis(1), + ) + .await; + + assert_eq!(result.name, "Cloud Sandbox"); + assert_eq!(result.status, CheckStatus::Error); + assert_eq!(result.summary, "timeout (1ms)"); + assert_eq!(result.details[0].text, "Daytona probe timed out"); + } + #[tokio::test] async fn check_brave_search_ignores_env_backed_api_key() { let state = TestAppStateBuilder::new() diff --git a/lib/apps/fabro-server/src/server/handler/system.rs b/lib/apps/fabro-server/src/server/handler/system.rs index d26e445a6..a38fe33f5 100644 --- a/lib/apps/fabro-server/src/server/handler/system.rs +++ b/lib/apps/fabro-server/src/server/handler/system.rs @@ -1,5 +1,7 @@ use std::collections::BTreeMap; +use std::future::Future; use std::sync::Arc; +use std::time::Duration; use chrono::Utc; use fabro_slack::config::{ @@ -9,6 +11,7 @@ use fabro_slack::config::{ use fabro_static::EnvVars; use fabro_types::settings::server::GithubIntegrationSettings; use fabro_vault::Vault; +use tokio::time::timeout; use super::super::{ AggregateBilling, AggregateBillingTotals, ApiError, AppState, BilledTokenCounts, @@ -21,6 +24,8 @@ use super::super::{ resource_sampler, spawn_blocking, system_sandbox_provider, to_i64, }; +const SERVER_DIAGNOSTICS_TIMEOUT: Duration = Duration::from_secs(25); + pub(super) fn routes() -> Router> { Router::new() .route("/repos/github/{owner}/{name}", get(get_github_repo)) @@ -683,11 +688,34 @@ async fn get_github_repo( } async fn run_diagnostics(_auth: RequiredUser, State(state): State>) -> Response { - ( - StatusCode::OK, - Json(diagnostics::run_all(state.as_ref()).await), + diagnostics_response_with_timeout( + Box::pin(diagnostics::run_all(state.as_ref())), + SERVER_DIAGNOSTICS_TIMEOUT, ) - .into_response() + .await +} + +async fn diagnostics_response_with_timeout( + diagnostics: F, + operation_timeout: Duration, +) -> Response +where + F: Future, +{ + let Ok(report) = timeout(operation_timeout, diagnostics).await else { + tracing::warn!( + timeout_secs = operation_timeout.as_secs(), + "server diagnostics timed out" + ); + return ApiError::with_code( + StatusCode::GATEWAY_TIMEOUT, + "Server diagnostics timed out.", + "diagnostics_timeout", + ) + .into_response(); + }; + + (StatusCode::OK, Json(report)).into_response() } pub(in crate::server) async fn openapi_spec() -> Response { @@ -737,3 +765,26 @@ async fn get_aggregate_billing( }; (StatusCode::OK, Json(response)).into_response() } + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn diagnostics_response_returns_gateway_timeout_before_client_deadline() { + let response = diagnostics_response_with_timeout( + std::future::pending::(), + Duration::from_millis(1), + ) + .await; + + assert_eq!(response.status(), StatusCode::GATEWAY_TIMEOUT); + let body = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .expect("diagnostics timeout response body should be readable"); + let body: serde_json::Value = serde_json::from_slice(&body) + .expect("diagnostics timeout response should contain JSON"); + assert_eq!(body["errors"][0]["code"], "diagnostics_timeout"); + assert_eq!(body["errors"][0]["detail"], "Server diagnostics timed out."); + } +} From b7e3b660ffa076e4f9897c0769897a3530eca089 Mon Sep 17 00:00:00 2001 From: Release Repro Date: Thu, 20 Aug 2026 20:34:38 -0400 Subject: [PATCH 2/2] Simplify diagnostics timeout handling --- lib/apps/fabro-server/src/diagnostics.rs | 130 ++++++------------ lib/apps/fabro-server/src/server.rs | 19 ++- .../fabro-server/src/server/handler/system.rs | 12 +- lib/components/fabro-llm/src/model_test.rs | 43 +++++- .../fabro-sandbox/src/daytona/mod.rs | 89 ++++++++++-- 5 files changed, 176 insertions(+), 117 deletions(-) diff --git a/lib/apps/fabro-server/src/diagnostics.rs b/lib/apps/fabro-server/src/diagnostics.rs index 86f625366..9373cc81e 100644 --- a/lib/apps/fabro-server/src/diagnostics.rs +++ b/lib/apps/fabro-server/src/diagnostics.rs @@ -6,7 +6,7 @@ use base64::Engine as _; use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use fabro_auth::auth_issue_message; use fabro_llm::client::Client as LlmClient; -use fabro_llm::model_test::{ModelTestOutcome, ModelTestStatus, run_basic_model_probe}; +use fabro_llm::model_test::{ModelTestStatus, run_basic_model_probe_with_timeout}; use fabro_model::{Catalog, ProviderId}; use fabro_redact::redact_string; use fabro_sandbox::{DockerSandboxProvider, daytona}; @@ -255,35 +255,15 @@ async fn probe_single_provider( None, ); }; - let model_id = model.id.clone(); + let model_id = model.id.to_string(); - let outcome = run_basic_model_probe(model_id.as_str(), provider.clone(), client); - provider_probe_with_timeout( - provider, - model_id.to_string(), - outcome, + let outcome = run_basic_model_probe_with_timeout( + &model_id, + &provider, + client, EXTERNAL_SERVICE_PROBE_TIMEOUT, ) - .await -} - -async fn provider_probe_with_timeout( - provider: ProviderId, - model_id: String, - probe: F, - probe_timeout: Duration, -) -> ProviderProbeResult -where - F: Future, -{ - let Ok(outcome) = timeout(probe_timeout, probe).await else { - return provider_probe_error( - provider, - Some(model_id), - probe_timeout_message(probe_timeout), - None, - ); - }; + .await; match outcome.status { ModelTestStatus::Ok => ProviderProbeResult { @@ -302,14 +282,6 @@ where } } -fn probe_timeout_message(probe_timeout: Duration) -> String { - if probe_timeout.subsec_nanos() == 0 { - format!("timeout ({}s)", probe_timeout.as_secs()) - } else { - format!("timeout ({}ms)", probe_timeout.as_millis()) - } -} - fn provider_probe_error( provider: ProviderId, model_id: Option, @@ -689,28 +661,13 @@ async fn check_cloud_sandbox(state: &AppState) -> CheckResult { }; }; - check_cloud_sandbox_with_probe( - || state.check_daytona_api_key(api_key), - EXTERNAL_SERVICE_PROBE_TIMEOUT, - ) - .await + let probe = state + .check_daytona_api_key_with_timeout(api_key, EXTERNAL_SERVICE_PROBE_TIMEOUT) + .await; + cloud_sandbox_probe_check(probe) } -async fn check_cloud_sandbox_with_probe(probe: F, probe_timeout: Duration) -> CheckResult -where - F: FnOnce() -> Fut, - Fut: Future>, -{ - let Ok(probe) = timeout(probe_timeout, probe()).await else { - return CheckResult { - name: "Cloud Sandbox".to_string(), - status: CheckStatus::Error, - summary: probe_timeout_message(probe_timeout), - details: vec![CheckDetail::new("Daytona probe timed out".to_string())], - remediation: Some("Verify DAYTONA_API_KEY value and Daytona reachability".to_string()), - }; - }; - +fn cloud_sandbox_probe_check(probe: anyhow::Result) -> CheckResult { match probe { Ok(check) if check.ok() => CheckResult { name: "Cloud Sandbox".to_string(), @@ -733,13 +690,29 @@ where daytona::required_perms_display() )), }, - Err(err) => CheckResult { - name: "Cloud Sandbox".to_string(), - status: CheckStatus::Error, - summary: "Daytona credential rejected".to_string(), - details: vec![CheckDetail::new(format!("{err:#}"))], - remediation: Some("Verify DAYTONA_API_KEY value and Daytona reachability".to_string()), - }, + Err(err) => { + if let Some(timeout) = err.downcast_ref::() { + return CheckResult { + name: "Cloud Sandbox".to_string(), + status: CheckStatus::Error, + summary: format!("timeout ({:?})", timeout.timeout()), + details: vec![CheckDetail::new("Daytona probe timed out".to_string())], + remediation: Some( + "Verify DAYTONA_API_KEY value and Daytona reachability".to_string(), + ), + }; + } + + CheckResult { + name: "Cloud Sandbox".to_string(), + status: CheckStatus::Error, + summary: "Daytona credential rejected".to_string(), + details: vec![CheckDetail::new(format!("{err:#}"))], + remediation: Some( + "Verify DAYTONA_API_KEY value and Daytona reachability".to_string(), + ), + } + } } } @@ -1090,27 +1063,6 @@ mod tests { ); } - #[tokio::test] - async fn provider_probe_reports_provider_specific_timeout() { - assert_eq!( - probe_timeout_message(EXTERNAL_SERVICE_PROBE_TIMEOUT), - "timeout (15s)" - ); - - let result = provider_probe_with_timeout( - ProviderId::new("modal"), - "modal/test-model".to_string(), - std::future::pending::(), - Duration::from_millis(1), - ) - .await; - - assert_eq!(result.provider, ProviderId::new("modal")); - assert_eq!(result.model_id.as_deref(), Some("modal/test-model")); - assert_eq!(result.status, ProviderProbeStatus::Error); - assert_eq!(result.error_message.as_deref(), Some("timeout (1ms)")); - } - #[test] fn docker_sandbox_probe_passes_when_daemon_responds() { let result = docker_sandbox_probe_check(Ok(())); @@ -1230,13 +1182,11 @@ enabled = false ); } - #[tokio::test] - async fn check_cloud_sandbox_reports_timeout() { - let result = check_cloud_sandbox_with_probe( - std::future::pending::>, - Duration::from_millis(1), - ) - .await; + #[test] + fn check_cloud_sandbox_reports_timeout() { + let result = cloud_sandbox_probe_check(Err(anyhow::Error::new( + daytona::DaytonaCredentialProbeTimeout::new(Duration::from_millis(1)), + ))); assert_eq!(result.name, "Cloud Sandbox"); assert_eq!(result.status, CheckStatus::Error); diff --git a/lib/apps/fabro-server/src/server.rs b/lib/apps/fabro-server/src/server.rs index 5445ad619..4976880f8 100644 --- a/lib/apps/fabro-server/src/server.rs +++ b/lib/apps/fabro-server/src/server.rs @@ -1455,6 +1455,15 @@ impl AppState { pub(crate) async fn check_daytona_api_key( &self, api_key: String, + ) -> anyhow::Result { + self.check_daytona_api_key_with_timeout(api_key, daytona::DAYTONA_CREDENTIAL_PROBE_TIMEOUT) + .await + } + + pub(crate) async fn check_daytona_api_key_with_timeout( + &self, + api_key: String, + probe_timeout: Duration, ) -> anyhow::Result { let base_url = self .config_env_lookup(EnvVars::DAYTONA_API_URL) @@ -1463,8 +1472,14 @@ impl AppState { let org_id = self.config_env_lookup(EnvVars::DAYTONA_ORGANIZATION_ID); let http_client = fabro_http::http_client().context("failed to build HTTP client")?; - daytona::check_daytona_api_key_with(&base_url, org_id.as_deref(), api_key, http_client) - .await + daytona::check_daytona_api_key_with_timeout( + &base_url, + org_id.as_deref(), + api_key, + http_client, + probe_timeout, + ) + .await } /// Borrow the persistent store so sibling modules can open run readers diff --git a/lib/apps/fabro-server/src/server/handler/system.rs b/lib/apps/fabro-server/src/server/handler/system.rs index a38fe33f5..d74f65944 100644 --- a/lib/apps/fabro-server/src/server/handler/system.rs +++ b/lib/apps/fabro-server/src/server/handler/system.rs @@ -778,12 +778,12 @@ mod tests { ) .await; - assert_eq!(response.status(), StatusCode::GATEWAY_TIMEOUT); - let body = axum::body::to_bytes(response.into_body(), usize::MAX) - .await - .expect("diagnostics timeout response body should be readable"); - let body: serde_json::Value = serde_json::from_slice(&body) - .expect("diagnostics timeout response should contain JSON"); + let body = fabro_test::expect_axum_json( + response, + StatusCode::GATEWAY_TIMEOUT, + "GET /api/v1/system/diagnostics timeout", + ) + .await; assert_eq!(body["errors"][0]["code"], "diagnostics_timeout"); assert_eq!(body["errors"][0]["detail"], "Server diagnostics timed out."); } diff --git a/lib/components/fabro-llm/src/model_test.rs b/lib/components/fabro-llm/src/model_test.rs index 6a985fb1b..de09f6fde 100644 --- a/lib/components/fabro-llm/src/model_test.rs +++ b/lib/components/fabro-llm/src/model_test.rs @@ -1,3 +1,4 @@ +use std::future::Future; use std::sync::Arc; use std::time::Duration; @@ -63,22 +64,38 @@ pub async fn run_basic_model_probe( model_id: &str, provider: impl ToString, client: Arc, +) -> ModelTestOutcome { + run_basic_model_probe_with_timeout( + model_id, + provider, + client, + Duration::from_secs(ModelTestMode::Basic.timeout_secs()), + ) + .await +} + +pub async fn run_basic_model_probe_with_timeout( + model_id: &str, + provider: impl ToString, + client: Arc, + probe_timeout: Duration, ) -> ModelTestOutcome { let params = GenerateParams::new(model_id, client) .provider(provider.to_string()) .prompt("Say OK") .max_tokens(16); - let result = time::timeout( - Duration::from_secs(ModelTestMode::Basic.timeout_secs()), - generate::generate(params), - ) - .await; + basic_model_probe_outcome(generate::generate(params), probe_timeout).await +} - match result { +async fn basic_model_probe_outcome(probe: F, probe_timeout: Duration) -> ModelTestOutcome +where + F: Future>, +{ + match time::timeout(probe_timeout, probe).await { Ok(Ok(_)) => ModelTestOutcome::ok(), Ok(Err(err)) => ModelTestOutcome::error(err.to_string()), - Err(_) => ModelTestOutcome::error("timeout (30s)"), + Err(_) => ModelTestOutcome::error(format!("timeout ({probe_timeout:?})")), } } @@ -244,6 +261,18 @@ mod tests { ); } + #[tokio::test] + async fn basic_model_probe_reports_configured_timeout() { + let outcome = basic_model_probe_outcome( + std::future::pending::>(), + Duration::from_millis(1), + ) + .await; + + assert_eq!(outcome.status, ModelTestStatus::Error); + assert_eq!(outcome.error_message.as_deref(), Some("timeout (1ms)")); + } + #[test] fn deep_test_omits_effort_for_reasoning_without_effort_controls() { let info = test_model_with(ModelFeatures { diff --git a/lib/components/fabro-sandbox/src/daytona/mod.rs b/lib/components/fabro-sandbox/src/daytona/mod.rs index 6371f23d2..44d014159 100644 --- a/lib/components/fabro-sandbox/src/daytona/mod.rs +++ b/lib/components/fabro-sandbox/src/daytona/mod.rs @@ -1,5 +1,6 @@ use std::collections::HashMap; use std::fmt::Write; +use std::future::Future; use std::path::Path; use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; @@ -62,7 +63,8 @@ pub const DEFAULT_DAYTONA_API_URL: &str = "https://app.daytona.io/api"; pub(crate) const DAYTONA_DASHBOARD_SANDBOXES_URL: &str = "https://app.daytona.io/dashboard/sandboxes"; const FABRO_SANDBOX_USER_AGENT: &str = concat!("fabro-sandbox/", env!("CARGO_PKG_VERSION")); -const DAYTONA_PROBE_TIMEOUT: Duration = Duration::from_secs(20); +pub const DAYTONA_CREDENTIAL_PROBE_TIMEOUT: Duration = Duration::from_secs(20); +const DAYTONA_BASH_SESSION_PROBE_TIMEOUT: Duration = Duration::from_secs(20); const DAYTONA_START_TIMEOUT: Duration = Duration::from_mins(1); /// Upper bound on explicit and Drop-triggered Daytona cleanup calls (session /// deletion, temporary stdin files) so a stalled REST call cannot block @@ -156,6 +158,24 @@ pub struct DaytonaKeyCheck { pub missing: Vec, } +#[derive(Debug, thiserror::Error)] +#[error("Daytona credential probe timed out after {timeout:?}")] +pub struct DaytonaCredentialProbeTimeout { + timeout: Duration, +} + +impl DaytonaCredentialProbeTimeout { + #[must_use] + pub const fn new(timeout: Duration) -> Self { + Self { timeout } + } + + #[must_use] + pub const fn timeout(&self) -> Duration { + self.timeout + } +} + impl DaytonaKeyCheck { pub fn ok(&self) -> bool { self.missing.is_empty() @@ -253,6 +273,23 @@ pub async fn check_daytona_api_key_with( org_id: Option<&str>, api_key: String, http_client: fabro_http::HttpClient, +) -> anyhow::Result { + check_daytona_api_key_with_timeout( + base_url, + org_id, + api_key, + http_client, + DAYTONA_CREDENTIAL_PROBE_TIMEOUT, + ) + .await +} + +pub async fn check_daytona_api_key_with_timeout( + base_url: &str, + org_id: Option<&str>, + api_key: String, + http_client: fabro_http::HttpClient, + probe_timeout: Duration, ) -> anyhow::Result { let work = async { let client = build_daytona_client_with( @@ -287,12 +324,21 @@ pub async fn check_daytona_api_key_with( }) }; - match time::timeout(DAYTONA_PROBE_TIMEOUT, work).await { + daytona_credential_probe_with_timeout(work, probe_timeout).await +} + +async fn daytona_credential_probe_with_timeout( + probe: F, + probe_timeout: Duration, +) -> anyhow::Result +where + F: Future>, +{ + match time::timeout(probe_timeout, probe).await { Ok(result) => result, - Err(_) => Err(anyhow::anyhow!( - "Daytona credential probe timed out after {}s", - DAYTONA_PROBE_TIMEOUT.as_secs() - )), + Err(_) => Err(anyhow::Error::new(DaytonaCredentialProbeTimeout::new( + probe_timeout, + ))), } } @@ -576,16 +622,16 @@ impl DaytonaSandbox { /// non-POSIX, and completion assertions all hold. /// /// Costs one session round trip plus a single status poll per sandbox - /// lifecycle transition. `DAYTONA_PROBE_TIMEOUT` is the outer backstop for - /// a stalled REST call; the inner [`BASH_PROBE_TIMEOUT_MS`] is the deadline - /// for the command itself. Session cleanup runs outside that deadline under - /// its own bounded timeout. + /// lifecycle transition. `DAYTONA_BASH_SESSION_PROBE_TIMEOUT` is the outer + /// backstop for a stalled REST call; the inner [`BASH_PROBE_TIMEOUT_MS`] is + /// the deadline for the command itself. Session cleanup runs outside that + /// deadline under its own bounded timeout. async fn probe_bash_session(sandbox: &daytona_sdk::Sandbox) -> crate::Result<()> { - let deadline = time::Instant::now() + DAYTONA_PROBE_TIMEOUT; + let deadline = time::Instant::now() + DAYTONA_BASH_SESSION_PROBE_TIMEOUT; let timeout_error = || { crate::Error::message(format!( "Daytona Bash session check timed out after {}s", - DAYTONA_PROBE_TIMEOUT.as_secs() + DAYTONA_BASH_SESSION_PROBE_TIMEOUT.as_secs() )) }; let mut session = match time::timeout_at(deadline, DaytonaSession::create(sandbox)).await { @@ -3358,6 +3404,25 @@ mod tests { auth.assert_async().await; } + #[tokio::test] + async fn daytona_credential_probe_reports_configured_timeout() { + let err = daytona_credential_probe_with_timeout( + std::future::pending::>(), + Duration::from_millis(1), + ) + .await + .expect_err("probe should time out"); + let timeout = err + .downcast_ref::() + .expect("timeout should preserve its type"); + + assert_eq!(timeout.timeout(), Duration::from_millis(1)); + assert_eq!( + err.to_string(), + "Daytona credential probe timed out after 1ms" + ); + } + #[tokio::test] async fn daytona_stdin_file_uploads_exact_bytes_and_is_deleted() { let server = MockServer::start_async().await;