diff --git a/src/cmd_alias.rs b/src/cmd_alias.rs index f65977cd..111e93d1 100644 --- a/src/cmd_alias.rs +++ b/src/cmd_alias.rs @@ -432,6 +432,7 @@ mod test { io, debug: false, override_host: None, + kcl_retry_config: None, }; let cmd_alias = crate::cmd_alias::CmdAlias { subcmd: t.cmd }; diff --git a/src/cmd_auth.rs b/src/cmd_auth.rs index 63c609e5..d0f9cb08 100644 --- a/src/cmd_auth.rs +++ b/src/cmd_auth.rs @@ -614,6 +614,7 @@ mod test { io, debug: false, override_host: None, + kcl_retry_config: None, }; if let Some(h) = &t.host { ctx.override_host = Some(h.clone()); diff --git a/src/cmd_completion.rs b/src/cmd_completion.rs index 3742cad4..660240f4 100644 --- a/src/cmd_completion.rs +++ b/src/cmd_completion.rs @@ -139,6 +139,7 @@ mod test { io, debug: false, override_host: None, + kcl_retry_config: None, }; cmd.run(&mut ctx).await.unwrap(); diff --git a/src/cmd_config.rs b/src/cmd_config.rs index d6967b22..17c354cd 100644 --- a/src/cmd_config.rs +++ b/src/cmd_config.rs @@ -265,6 +265,7 @@ mod test { io, debug: false, override_host: None, + kcl_retry_config: None, }; let cmd_config = crate::cmd_config::CmdConfig { subcmd: t.cmd }; @@ -301,6 +302,7 @@ mod test { io, debug: false, override_host: None, + kcl_retry_config: None, }; let mut cmd_config = crate::cmd_config::CmdConfig { diff --git a/src/cmd_file.rs b/src/cmd_file.rs index c9f7a991..e56b3c7a 100644 --- a/src/cmd_file.rs +++ b/src/cmd_file.rs @@ -868,6 +868,7 @@ mod test { io, debug: false, override_host: None, + kcl_retry_config: None, }; let cmd_file = crate::cmd_file::CmdFile { subcmd: t.cmd }; diff --git a/src/cmd_generate.rs b/src/cmd_generate.rs index a93205cf..e29fb75f 100644 --- a/src/cmd_generate.rs +++ b/src/cmd_generate.rs @@ -163,6 +163,7 @@ mod test { io, debug: false, override_host: None, + kcl_retry_config: None, }; let cmd = crate::cmd_generate::CmdGenerateMarkdown { dir: "".to_string() }; @@ -192,6 +193,7 @@ mod test { io, debug: false, override_host: None, + kcl_retry_config: None, }; let cmd = crate::cmd_generate::CmdGenerateMarkdown { dir: "".to_string() }; diff --git a/src/cmd_kcl.rs b/src/cmd_kcl.rs index 32e1f17e..9b5edf3e 100644 --- a/src/cmd_kcl.rs +++ b/src/cmd_kcl.rs @@ -159,7 +159,7 @@ impl crate::cmd::Command for CmdKclExport { let units: UnitLength = meta_settings.default_length_units.to_kcmc(); let output_format = get_output_format(&self.output_format, units, self.deterministic); - if crate::context::Context::use_server_kcl_execution() { + let (files, session_data) = if crate::context::Context::use_server_kcl_execution() { let (mut responses, session_data) = ctx .run_server_kcl_then_modeling_cmds( "", @@ -176,37 +176,18 @@ impl crate::cmd::Command for CmdKclExport { anyhow::bail!("Expected Export response from engine"); }; - for file in files { - let path = self.output_dir.join(file.name); - std::fs::write(&path, file.contents)?; - - writeln!(ctx.io.out, "Wrote file: {}", path.display())?; - } - - if self.show_trace { - print_trace_link(&mut ctx.io, &session_data) - } - - return Ok(()); - } - - let client = ctx.api_client("")?; - let ectx = kcl_lib::ExecutorContext::new(&client, settings).await?; - let mut state = kcl_lib::ExecState::new(&ectx); - let session_data = ectx - .run(&program, &mut state) - .await - .map_err(|err| kcl_error_fmt::into_miette(err, &code))? - .1; - kcl_error_fmt::check_exec_state_issues( - &mut ctx.io.err_out, - &filepath.display().to_string(), - &code, - &state, - self.run_options.issue_check(), - )?; - - let files = ectx.export(output_format).await?; + (files, session_data) + } else { + ctx.run_kcl_then_export( + &filepath.display().to_string(), + &code, + &program, + settings, + self.run_options.issue_check(), + output_format, + ) + .await? + }; // Save the files to our export directory. for file in files { diff --git a/src/cmd_say.rs b/src/cmd_say.rs index f2bef2fb..54cbff5e 100644 --- a/src/cmd_say.rs +++ b/src/cmd_say.rs @@ -116,6 +116,7 @@ mod test { io, debug: false, override_host: None, + kcl_retry_config: None, }; let cmd_say = crate::cmd_say::CmdSay { input: t.cmd.input }; diff --git a/src/cmd_user.rs b/src/cmd_user.rs index 90782278..47a88e3e 100644 --- a/src/cmd_user.rs +++ b/src/cmd_user.rs @@ -85,6 +85,7 @@ mod test { io, debug: false, override_host: None, + kcl_retry_config: None, }; let cmd_user = crate::cmd_user::CmdUser { subcmd: t.cmd }; diff --git a/src/context.rs b/src/context.rs index 34bc2122..4bc370c3 100644 --- a/src/context.rs +++ b/src/context.rs @@ -9,7 +9,7 @@ use kittycad_modeling_cmds::{ ModelingCmd, each_cmd as mcmd, output::TakeSnapshot, websocket::{ - FailureWebSocketResponse, ModelingCmdReq, ModelingSessionData, OkWebSocketResponseData, + FailureWebSocketResponse, ModelingCmdReq, ModelingSessionData, OkWebSocketResponseData, RawFile, SuccessWebSocketResponse, WebSocketRequest, WebSocketResponse, }, }; @@ -36,9 +36,169 @@ pub struct Context<'a> { pub debug: bool, // If set, override the host used when commands don't specify one. pub(crate) override_host: Option, + /// Overrides retry behavior for KCL execution tests. + #[cfg(test)] + pub(crate) kcl_retry_config: Option, +} + +/// Controls how retryable KCL execution failures are retried. +#[derive(Debug, Clone)] +pub(crate) struct RetryConfig { + retries: usize, + print_retries: bool, +} + +impl RetryConfig { + /// Returns retry settings that make each KCL execution attempt run only + /// once. + fn no_retries() -> Self { + Self { + retries: 0, + print_retries: false, + } + } +} + +#[cfg(test)] +impl Default for RetryConfig { + /// Returns the retry settings used by tests that exercise retry handling. + fn default() -> Self { + Self { + retries: 2, + print_retries: true, + } + } +} + +/// Result data captured from a single KCL program execution attempt. +struct KclProgramRun { + exec_ctx: kcl_lib::ExecutorContext, + exec_state: kcl_lib::ExecState, + session_data: Option, +} + +/// Error type used to preserve KCL diagnostics across retry attempts. +enum KclExecError { + /// A KCL error that should be formatted with filename and source context. + Err(Box), + /// A KCL execution error that already includes command outputs. + WithOutputs(Box), + /// Any non-KCL error from setup or issue checking. + Other(anyhow::Error), +} + +impl KclExecError { + /// Converts the stored error into the diagnostic form expected by callers. + fn into_anyhow(self, filename: &str, code: &str) -> anyhow::Error { + match self { + Self::Err(err) => kcl_error_fmt::into_miette_for_parse(filename, code, *err), + Self::WithOutputs(err) => kcl_error_fmt::into_miette(*err, code), + Self::Other(err) => err, + } + } +} + +impl std::fmt::Display for KclExecError { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Err(err) => write!(formatter, "{err}"), + Self::WithOutputs(err) => write!(formatter, "{err}"), + Self::Other(err) => write!(formatter, "{err}"), + } + } +} + +impl kcl_lib::IsRetryable for KclExecError { + fn is_retryable(&self) -> bool { + match self { + Self::Err(err) => kcl_lib::IsRetryable::is_retryable(err.as_ref()), + Self::WithOutputs(err) => kcl_lib::IsRetryable::is_retryable(err.as_ref()), + Self::Other(_) => false, + } + } +} + +/// Runs an async operation until it succeeds, fails fatally, or exhausts +/// retries. +#[cfg(test)] +async fn execute_with_retries(config: &RetryConfig, mut execute: F) -> std::result::Result +where + F: FnMut() -> Fut, + Fut: std::future::Future>, + E: kcl_lib::IsRetryable + std::fmt::Display, +{ + let mut retries_remaining = config.retries; + loop { + let exec_result = execute().await; + + if should_retry_kcl_attempt(config, &mut retries_remaining, &exec_result) { + continue; + } + + return exec_result; + } +} + +/// Adjusts the retries remaining and returns whether a failed KCL attempt +/// should be retried. This handles all the book-keeping of retrying. +fn should_retry_kcl_attempt( + config: &RetryConfig, + retries_remaining: &mut usize, + result: &std::result::Result, +) -> bool +where + E: kcl_lib::IsRetryable + std::fmt::Display, +{ + if *retries_remaining > 0 + && let Err(error) = result + && kcl_lib::IsRetryable::is_retryable(error) + { + if config.print_retries { + eprintln!("Execute got {error}; retrying..."); + } + *retries_remaining -= 1; + true + } else { + false + } +} + +/// Runs one KCL execution attempt and returns state needed for follow-up +/// commands. +async fn run_kcl_program_once_with_client( + client: &kittycad::Client, + program: &kcl_lib::Program, + settings: kcl_lib::ExecutorSettings, +) -> std::result::Result { + let ctx = kcl_lib::ExecutorContext::new(client, settings) + .await + .map_err(KclExecError::Other)?; + let mut exec_state = kcl_lib::ExecState::new(&ctx); + let session_data = ctx + .run(program, &mut exec_state) + .await + .map_err(|err| KclExecError::WithOutputs(Box::new(err)))? + .1; + Ok(KclProgramRun { + exec_ctx: ctx, + exec_state, + session_data, + }) } impl<'a> Context<'a> { + /// Returns the retry settings to use for local KCL execution. + fn kcl_retry_config(&self) -> RetryConfig { + #[cfg(test)] + { + self.kcl_retry_config.clone().unwrap_or_else(RetryConfig::no_retries) + } + #[cfg(not(test))] + { + RetryConfig::no_retries() + } + } + fn resolve_api_host_and_baseurl(&self, hostname: &str) -> Result<(String, String)> { let host = if !hostname.is_empty() { hostname.to_string() @@ -112,6 +272,8 @@ impl<'a> Context<'a> { io, debug: false, override_host: None, + #[cfg(test)] + kcl_retry_config: None, } } @@ -445,44 +607,64 @@ impl<'a> Context<'a> { let program = kcl_lib::Program::parse_no_errs(code) .map_err(|err| kcl_error_fmt::into_miette_for_parse(filename, code, err))?; - let ctx = kcl_lib::ExecutorContext::new(&client, cmd_kcl::with_heartbeats(settings)).await?; - let mut state = kcl_lib::ExecState::new(&ctx); - let session_data = ctx - .run(&program, &mut state) - .await - .map_err(|err| kcl_error_fmt::into_miette(err, code))? - .1; - kcl_error_fmt::check_exec_state_issues(&mut self.io.err_out, filename, code, &state, issue_check)?; - - let batch_context = kcl_lib::EngineBatchContext::new(); + let settings = cmd_kcl::with_heartbeats(settings); + let retry_config = self.kcl_retry_config(); + let mut retries_remaining = retry_config.retries; + loop { + let result: std::result::Result<(OkWebSocketResponseData, Option), KclExecError> = + async { + let run = run_kcl_program_once_with_client(&client, &program, settings.clone()).await?; + + kcl_error_fmt::check_exec_state_issues( + &mut self.io.err_out, + filename, + code, + &run.exec_state, + issue_check, + ) + .map_err(KclExecError::Other)?; + + let batch_context = kcl_lib::EngineBatchContext::new(); + + // Zoom on the object. + run.exec_ctx + .engine + .send_modeling_cmd( + &batch_context, + uuid::Uuid::new_v4(), + kcl_lib::SourceRange::default(), + &ModelingCmd::from( + mcmd::ZoomToFit::builder() + .animated(false) + .object_ids(Default::default()) + .padding(0.1) + .build(), + ), + ) + .await + .map_err(|err| KclExecError::Err(Box::new(err)))?; + + let resp = run + .exec_ctx + .engine + .send_modeling_cmd( + &batch_context, + uuid::Uuid::new_v4(), + kcl_lib::SourceRange::default(), + &cmd, + ) + .await + .map_err(|err| KclExecError::Err(Box::new(err)))?; + Ok((resp, run.session_data)) + } + .await; - // Zoom on the object. - ctx.engine - .send_modeling_cmd( - &batch_context, - uuid::Uuid::new_v4(), - kcl_lib::SourceRange::default(), - &ModelingCmd::from( - mcmd::ZoomToFit::builder() - .animated(false) - .object_ids(Default::default()) - .padding(0.1) - .build(), - ), - ) - .await?; + if should_retry_kcl_attempt(&retry_config, &mut retries_remaining, &result) { + continue; + } - let resp = ctx - .engine - .send_modeling_cmd( - &batch_context, - uuid::Uuid::new_v4(), - kcl_lib::SourceRange::default(), - &cmd, - ) - .await - .map_err(|err| kcl_error_fmt::into_miette_for_parse(filename, code, err))?; - Ok((resp, session_data)) + return result.map_err(|err| err.into_anyhow(filename, code)); + } } pub(crate) async fn run_kcl_then_modeling_cmds( @@ -505,32 +687,50 @@ impl<'a> Context<'a> { let program = kcl_lib::Program::parse_no_errs(code) .map_err(|err| kcl_error_fmt::into_miette_for_parse(filename, code, err))?; - let ctx = kcl_lib::ExecutorContext::new(&client, cmd_kcl::with_heartbeats(settings)).await?; - let mut state = kcl_lib::ExecState::new(&ctx); - let session_data = ctx - .run(&program, &mut state) - .await - .map_err(|err| kcl_error_fmt::into_miette(err, code))? - .1; - kcl_error_fmt::check_exec_state_issues(&mut self.io.err_out, filename, code, &state, issue_check)?; + let settings = cmd_kcl::with_heartbeats(settings); + let retry_config = self.kcl_retry_config(); + let mut retries_remaining = retry_config.retries; + loop { + let result: std::result::Result<(Vec, Option), KclExecError> = + async { + let run = run_kcl_program_once_with_client(&client, &program, settings.clone()).await?; + + kcl_error_fmt::check_exec_state_issues( + &mut self.io.err_out, + filename, + code, + &run.exec_state, + issue_check, + ) + .map_err(KclExecError::Other)?; + + let batch_context = kcl_lib::EngineBatchContext::new(); + let mut responses = Vec::with_capacity(cmds.len()); + for cmd in cmds.clone() { + let resp = run + .exec_ctx + .engine + .send_modeling_cmd( + &batch_context, + uuid::Uuid::new_v4(), + kcl_lib::SourceRange::default(), + &cmd, + ) + .await + .map_err(|err| KclExecError::Err(Box::new(err)))?; + responses.push(resp); + } - let batch_context = kcl_lib::EngineBatchContext::new(); - let mut responses = Vec::with_capacity(cmds.len()); - for cmd in cmds { - let resp = ctx - .engine - .send_modeling_cmd( - &batch_context, - uuid::Uuid::new_v4(), - kcl_lib::SourceRange::default(), - &cmd, - ) - .await - .map_err(|err| kcl_error_fmt::into_miette_for_parse(filename, code, err))?; - responses.push(resp); - } + Ok((responses, run.session_data)) + } + .await; - Ok((responses, session_data)) + if should_retry_kcl_attempt(&retry_config, &mut retries_remaining, &result) { + continue; + } + + return result.map_err(|err| err.into_anyhow(filename, code)); + } } /// Run the given KCL program, then after, run the given extra modeling commands. @@ -573,37 +773,98 @@ impl<'a> Context<'a> { let program = kcl_lib::Program::parse_no_errs(code) .map_err(|err| kcl_error_fmt::into_miette_for_parse(filename, code, err))?; - let ctx = kcl_lib::ExecutorContext::new(&client, cmd_kcl::with_heartbeats(settings)).await?; - let mut state = kcl_lib::ExecState::new(&ctx); - let session_data = ctx - .run(&program, &mut state) - .await - .map_err(|err| kcl_error_fmt::into_miette(err, code))? - .1; - kcl_error_fmt::check_exec_state_issues(&mut self.io.err_out, filename, code, &state, issue_check)?; + let settings = cmd_kcl::with_heartbeats(settings); + let retry_config = self.kcl_retry_config(); + let mut retries_remaining = retry_config.retries; + loop { + let result: std::result::Result<(Vec, Option), KclExecError> = async { + let run = run_kcl_program_once_with_client(&client, &program, settings.clone()).await?; - let batch_context = kcl_lib::EngineBatchContext::new(); - let mut snapshot_resps = Vec::new(); - for snapshot_cmd in snapshot_cmds { - let resp = ctx - .engine - .send_modeling_cmd( - &batch_context, - uuid::Uuid::new_v4(), - kcl_lib::SourceRange::default(), - &snapshot_cmd, + kcl_error_fmt::check_exec_state_issues( + &mut self.io.err_out, + filename, + code, + &run.exec_state, + issue_check, ) - .await - .map_err(|err| kcl_error_fmt::into_miette_for_parse(filename, code, err))?; - if let OkWebSocketResponseData::Modeling { - modeling_response: kittycad_modeling_cmds::ok_response::OkModelingCmdResponse::TakeSnapshot(snap), - } = resp - { - snapshot_resps.push(snap); + .map_err(KclExecError::Other)?; + + let batch_context = kcl_lib::EngineBatchContext::new(); + let mut snapshot_resps = Vec::new(); + for snapshot_cmd in snapshot_cmds.clone() { + let resp = run + .exec_ctx + .engine + .send_modeling_cmd( + &batch_context, + uuid::Uuid::new_v4(), + kcl_lib::SourceRange::default(), + &snapshot_cmd, + ) + .await + .map_err(|err| KclExecError::Err(Box::new(err)))?; + if let OkWebSocketResponseData::Modeling { + modeling_response: + kittycad_modeling_cmds::ok_response::OkModelingCmdResponse::TakeSnapshot(snap), + } = resp + { + snapshot_resps.push(snap); + } + } + + Ok((snapshot_resps, run.session_data)) } + .await; + + if should_retry_kcl_attempt(&retry_config, &mut retries_remaining, &result) { + continue; + } + + return result.map_err(|err| err.into_anyhow(filename, code)); } + } + + /// Runs KCL, checks execution issues, and exports the resulting files. + pub(crate) async fn run_kcl_then_export( + &mut self, + filename: &str, + code: &str, + program: &kcl_lib::Program, + settings: kcl_lib::ExecutorSettings, + issue_check: kcl_error_fmt::KclIssueCheck, + output_format: kittycad_modeling_cmds::format::OutputFormat3d, + ) -> Result<(Vec, Option)> { + let client = self.api_client("")?; + let retry_config = self.kcl_retry_config(); + let mut retries_remaining = retry_config.retries; + loop { + let result: std::result::Result<(Vec, Option), KclExecError> = async { + let run = run_kcl_program_once_with_client(&client, program, settings.clone()).await?; + + kcl_error_fmt::check_exec_state_issues( + &mut self.io.err_out, + filename, + code, + &run.exec_state, + issue_check, + ) + .map_err(KclExecError::Other)?; + + let files = run + .exec_ctx + .export(output_format.clone()) + .await + .map_err(|err| KclExecError::Err(Box::new(err)))?; + Ok((files, run.session_data)) + } + .await; + + if should_retry_kcl_attempt(&retry_config, &mut retries_remaining, &result) { + continue; + } - Ok((snapshot_resps, session_data)) + return result.map_err(|err| err.into_anyhow(filename, code)); + } } /// Create and poll a plain text-to-CAD job. @@ -1388,6 +1649,30 @@ mod test { use super::*; + /// Test error used to verify retryable and fatal retry paths. + #[derive(Debug, Clone, Copy, PartialEq, Eq)] + enum RetryTestError { + /// Error variant that should be retried. + Retryable, + /// Error variant that should stop retrying. + Fatal, + } + + impl std::fmt::Display for RetryTestError { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Retryable => write!(formatter, "retryable"), + Self::Fatal => write!(formatter, "fatal"), + } + } + } + + impl kcl_lib::IsRetryable for RetryTestError { + fn is_retryable(&self) -> bool { + matches!(self, Self::Retryable) + } + } + pub struct TestItem { name: String, zoo_pager_env: String, @@ -1624,6 +1909,69 @@ mod test { assert_eq!(md, "Hello world"); } + /// Verifies that retryable failures are retried until the operation + /// succeeds. + #[tokio::test] + async fn execute_with_retries_retries_retryable_errors() { + let retry_config = RetryConfig { + retries: 2, + print_retries: false, + }; + let mut attempts = 0; + + let result: std::result::Result<&str, RetryTestError> = execute_with_retries(&retry_config, || { + attempts += 1; + let attempt = attempts; + async move { + if attempt < 3 { + Err(RetryTestError::Retryable) + } else { + Ok("ok") + } + } + }) + .await; + + assert_eq!(result, Ok("ok")); + assert_eq!(attempts, 3); + } + + /// Verifies that fatal failures are returned without retrying. + #[tokio::test] + async fn execute_with_retries_does_not_retry_fatal_errors() { + let retry_config = RetryConfig { + retries: 2, + print_retries: false, + }; + let mut attempts = 0; + + let result: std::result::Result<(), RetryTestError> = execute_with_retries(&retry_config, || { + attempts += 1; + async { Err(RetryTestError::Fatal) } + }) + .await; + + assert_eq!(result, Err(RetryTestError::Fatal)); + assert_eq!(attempts, 1); + } + + /// Verifies that retryable failures are not retried when retries are + /// disabled. + #[tokio::test] + async fn execute_with_retries_no_retries_does_not_retry_retryable_errors() { + let retry_config = RetryConfig::no_retries(); + let mut attempts = 0; + + let result: std::result::Result<(), RetryTestError> = execute_with_retries(&retry_config, || { + attempts += 1; + async { Err(RetryTestError::Retryable) } + }) + .await; + + assert_eq!(result, Err(RetryTestError::Retryable)); + assert_eq!(attempts, 1); + } + #[test] fn resolve_host_prefers_explicit_then_global() { let mut config = crate::config::new_blank_config().unwrap(); @@ -1634,6 +1982,7 @@ mod test { io, debug: false, override_host: None, + kcl_retry_config: None, }; // No override: falls back to default host in config (which will be DEFAULT_HOST initially) @@ -1667,6 +2016,7 @@ mod test { io, debug: false, override_host: None, + kcl_retry_config: None, }; let (files, filepath) = ctx diff --git a/src/tests.rs b/src/tests.rs index ee8681a9..b255478a 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -308,6 +308,7 @@ async fn run_test_item_with_config(config: &mut TestConfig, item: TestItem) { io, debug: false, override_host: None, + kcl_retry_config: Some(crate::context::RetryConfig::default()), }; let _cwd_guard = CurrentDirGuard::change_to(item.current_directory.as_deref())