diff --git a/Cargo.lock b/Cargo.lock index f6a4b668..9a3c7ad5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -524,6 +524,15 @@ dependencies = [ "alloc-stdlib", ] +[[package]] +name = "bs58" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf88ba1141d185c399bee5288d850d63b8369520c1eafc32a0430b5b6c287bf4" +dependencies = [ + "tinyvec", +] + [[package]] name = "bstr" version = "1.12.1" @@ -1157,6 +1166,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ececcb659e7ba858fb4f10388c250a7252eb0a27373f1a72b8748afdd248e587" dependencies = [ "powerfmt", + "serde_core", ] [[package]] @@ -2131,6 +2141,12 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fbd780fe5cc30f81464441920d82ac8740e2e46b29a6fad543ddd075229ce37e" +[[package]] +name = "hex" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" + [[package]] name = "home" version = "0.5.11" @@ -2785,16 +2801,15 @@ dependencies = [ [[package]] name = "kcl-api" version = "0.2.167" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4d0782f9f12d26f00ae5c0a28a670e4e60ae8ea8f7eb8fe2dc644edbffc06b7d" +source = "git+https://github.com/KittyCAD/modeling-app?branch=main#16dee6595e8585a4aeaddcde435cda2d85b489ed" dependencies = [ "indexmap 2.14.0", - "kcl-error", + "kcl-error 0.2.167 (git+https://github.com/KittyCAD/modeling-app?branch=main)", "kittycad-measurements", "kittycad-unit-conversion-derive", "parse-display 0.11.0", "parse-display-derive 0.11.0", - "schemars", + "schemars 0.8.22", "serde", "ts-rs", "uuid", @@ -2803,8 +2818,7 @@ dependencies = [ [[package]] name = "kcl-derive-docs" version = "0.2.167" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "570fc50f5925bf4b10014fe3403789bdb83da16406c59dccd72d40eccd546e13" +source = "git+https://github.com/KittyCAD/modeling-app?branch=main#16dee6595e8585a4aeaddcde435cda2d85b489ed" dependencies = [ "proc-macro2", "quote", @@ -2818,7 +2832,20 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5df43f5977901ffe0036cca941f2e542f6d04f70889d8d95735188c0e16944c6" dependencies = [ "miette", - "schemars", + "schemars 0.8.22", + "serde", + "serde_json", + "thiserror 2.0.18", + "ts-rs", +] + +[[package]] +name = "kcl-error" +version = "0.2.167" +source = "git+https://github.com/KittyCAD/modeling-app?branch=main#16dee6595e8585a4aeaddcde435cda2d85b489ed" +dependencies = [ + "miette", + "schemars 0.8.22", "serde", "serde_json", "thiserror 2.0.18", @@ -2828,8 +2855,7 @@ dependencies = [ [[package]] name = "kcl-lib" version = "0.2.167" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b8554618da616fdb606bf0777271c644aa3779482f8ee2a47bd921c1f18c7114" +source = "git+https://github.com/KittyCAD/modeling-app?branch=main#16dee6595e8585a4aeaddcde435cda2d85b489ed" dependencies = [ "ahash", "anyhow", @@ -2859,7 +2885,7 @@ dependencies = [ "js-sys", "kcl-api", "kcl-derive-docs", - "kcl-error", + "kcl-error 0.2.167 (git+https://github.com/KittyCAD/modeling-app?branch=main)", "kittycad", "kittycad-modeling-cmds", "lazy_static", @@ -2876,7 +2902,7 @@ dependencies = [ "rgba_simple", "rmp-serde", "ropey", - "schemars", + "schemars 0.8.22", "serde", "serde_json", "sha2 0.11.0", @@ -2905,8 +2931,7 @@ dependencies = [ [[package]] name = "kcl-test-server" version = "0.2.167" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "efc65b9bf2c113192810400f572c4c79b6b51724e089274c9bcafa4287904a4c" +source = "git+https://github.com/KittyCAD/modeling-app?branch=main#16dee6595e8585a4aeaddcde435cda2d85b489ed" dependencies = [ "anyhow", "hyper 0.14.32", @@ -2946,7 +2971,7 @@ dependencies = [ "reqwest-middleware", "reqwest-retry", "reqwest-tracing", - "schemars", + "schemars 0.8.22", "serde", "serde_bytes", "serde_json", @@ -2970,9 +2995,9 @@ dependencies = [ [[package]] name = "kittycad-modeling-cmds" -version = "0.2.208" +version = "0.2.209" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b6e1f60a7440459009f291de2287e3c9b030120c3827ecb8a467a31370103d14" +checksum = "f1068db9b2fe4d8fbf20f9002d2433837a9f7afded7510cf7ed973b4cb69dd24" dependencies = [ "anyhow", "bon", @@ -2982,17 +3007,20 @@ dependencies = [ "enum-iterator-derive", "euler", "http 1.4.2", + "kcl-error 0.2.167 (registry+https://github.com/rust-lang/crates.io-index)", "kittycad-measurements", "kittycad-modeling-cmds-macros", "kittycad-unit-conversion-derive", "parse-display 0.11.0", "parse-display-derive 0.11.0", - "schemars", + "schemars 0.8.22", "serde", "serde_bytes", "serde_json", + "serde_with", "tabled", "ts-rs", + "typed-path", "uuid", ] @@ -3652,7 +3680,7 @@ dependencies = [ "reqwest 0.12.28", "reqwest-middleware", "rustfmt-wrapper", - "schemars", + "schemars 0.8.22", "serde", "serde_json", "serde_yaml 0.9.34+deprecated", @@ -4540,6 +4568,26 @@ dependencies = [ "thiserror 2.0.18", ] +[[package]] +name = "ref-cast" +version = "1.0.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f354300ae66f76f1c85c5f84693f0ce81d747e2c3f21a45fef496d89c960bf7d" +dependencies = [ + "ref-cast-impl", +] + +[[package]] +name = "ref-cast-impl" +version = "1.0.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7186006dcb21920990093f30e3dea63b7d6e977bf1256be20c3563a5db070da" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + [[package]] name = "regex" version = "1.12.4" @@ -5001,6 +5049,30 @@ dependencies = [ "uuid", ] +[[package]] +name = "schemars" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4cd191f9397d57d581cddd31014772520aa448f65ef991055d7f61582c65165f" +dependencies = [ + "dyn-clone", + "ref-cast", + "serde", + "serde_json", +] + +[[package]] +name = "schemars" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2b42f36aa1cd011945615b92222f6bf73c599a102a300334cd7f8dbeec726cc" +dependencies = [ + "dyn-clone", + "ref-cast", + "serde", + "serde_json", +] + [[package]] name = "schemars_derive" version = "0.8.22" @@ -5191,6 +5263,38 @@ dependencies = [ "serde", ] +[[package]] +name = "serde_with" +version = "3.21.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76a5c54c7310e7b8b9577c286d7e399ddd876c3e12b3ed917a8aabc4b96e9e8c" +dependencies = [ + "base64 0.22.1", + "bs58", + "chrono", + "hex", + "indexmap 1.9.3", + "indexmap 2.14.0", + "schemars 0.9.0", + "schemars 1.2.1", + "serde_core", + "serde_json", + "serde_with_macros", + "time", +] + +[[package]] +name = "serde_with_macros" +version = "3.21.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "84d57bc0c8b9a17920c178daa6bb924850d54a9c97ab45194bb8c17ad66bb660" +dependencies = [ + "darling 0.23.0", + "proc-macro2", + "quote", + "syn 2.0.118", +] + [[package]] name = "serde_yaml" version = "0.8.26" @@ -7403,6 +7507,7 @@ dependencies = [ "ignore", "image", "itertools 0.15.0", + "kcl-error 0.2.167 (registry+https://github.com/rust-lang/crates.io-index)", "kcl-lib", "kcl-test-server", "kittycad", @@ -7421,6 +7526,7 @@ dependencies = [ "regex", "reqwest 0.12.28", "ring", + "rmp-serde", "serde", "serde_json", "serde_yaml 0.9.34+deprecated", diff --git a/Cargo.toml b/Cargo.toml index 4cafa649..9306866d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -39,6 +39,7 @@ image = { version = "0.25", default-features = false, features = [ "jpeg", ] } itertools = "0.15.0" +kcl-error = "=0.2.167" kcl-lib = { version = "=0.2.167", features = ["disable-println"] } kcl-test-server = "=0.2.167" kittycad = { version = "0.4.12", features = [ @@ -47,7 +48,7 @@ kittycad = { version = "0.4.12", features = [ "requests", "retry", ] } -kittycad-modeling-cmds = { version = "0.2.208", features = [ +kittycad-modeling-cmds = { version = "0.2.209", features = [ "websocket", "tabled", ] } @@ -73,6 +74,7 @@ reqwest = { version = "0.12", default-features = false, features = [ "deflate", ] } ring = "0.17.14" +rmp-serde = "1" serde = { version = "1", features = ["derive"] } serde_json = "1" serde_yaml = "0.9" @@ -121,5 +123,6 @@ incremental = true debug = 0 [patch.crates-io] -# kcl-lib = { git = "https://github.com/KittyCAD/modeling-app", branch = "main" } +kcl-lib = { git = "https://github.com/KittyCAD/modeling-app", branch = "main" } +kcl-test-server = { git = "https://github.com/KittyCAD/modeling-app", branch = "main" } # kittycad-modeling-cmds = { git = "https://github.com/KittyCAD/modeling-api", branch = "achalmers/remove-cruft"} diff --git a/src/cmd_kcl.rs b/src/cmd_kcl.rs index 70692231..ad6c4eb9 100644 --- a/src/cmd_kcl.rs +++ b/src/cmd_kcl.rs @@ -157,6 +157,39 @@ impl crate::cmd::Command for CmdKclExport { .map_err(|err| kcl_error_fmt::into_miette_for_parse(&filepath.display().to_string(), &code, err))?; let meta_settings = program.meta_settings()?.unwrap_or_default(); 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 (mut responses, session_data) = ctx + .run_server_kcl_then_modeling_cmds( + "", + &filepath.display().to_string(), + &filepath, + &code, + vec![kcmc::ModelingCmd::Export( + kcmc::Export::builder().entity_ids(vec![]).format(output_format).build(), + )], + settings, + self.run_options.issue_check(), + ) + .await?; + let Some(kcmc::websocket::OkWebSocketResponseData::Export { files }) = responses.pop() else { + 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?; @@ -174,9 +207,7 @@ impl crate::cmd::Command for CmdKclExport { self.run_options.issue_check(), )?; - let files = ectx - .export(get_output_format(&self.output_format, units, self.deterministic)) - .await?; + let files = ectx.export(output_format).await?; // Save the files to our export directory. for file in files { diff --git a/src/context.rs b/src/context.rs index 33f98307..1c64e9a6 100644 --- a/src/context.rs +++ b/src/context.rs @@ -1,15 +1,30 @@ -use std::{io::Write, str::FromStr, time::Duration}; +use std::{io::Write, path::Path, str::FromStr, time::Duration}; use anyhow::{anyhow, Result}; -use futures::StreamExt; +use futures::{SinkExt, StreamExt}; use kcl_lib::engine_connection::EngineManager; -use kcmc::{each_cmd as mcmd, websocket::OkWebSocketResponseData}; +use kcmc::{ + each_cmd as mcmd, + exec_kcl::{KclFile, KclProject}, + shared::safe_filepath::SafeFilepath, + websocket::{ + FailureWebSocketResponse, ModelingCmdReq, ModelingSessionData, OkWebSocketResponseData, + SuccessWebSocketResponse, WebSocketRequest, WebSocketResponse, + }, +}; use kittycad::types::{ApiCallStatus, AsyncApiCallOutput, TextToCad, TextToCadCreateBody, TextToCadMultiFileIteration}; -use kittycad_modeling_cmds::{self as kcmc, output::TakeSnapshot, websocket::ModelingSessionData, ModelingCmd}; -use tokio_tungstenite::{tungstenite::protocol::Role, WebSocketStream}; +use kittycad_modeling_cmds::{self as kcmc, output::TakeSnapshot, ModelingCmd}; +use tokio_tungstenite::{tungstenite::protocol::Role, tungstenite::Message as WsMsg, WebSocketStream}; use crate::{cmd_kcl, config::Config, config_file::get_env_var, kcl_error_fmt, types::FormatOutput}; +type DirectWs = WebSocketStream; +type DirectWsRead = futures::stream::SplitStream; +type DirectWsWrite = futures::stream::SplitSink; + +const ENGINE_EXECUTION_ENV: &str = "ENGINE_EXECUTION"; +const WS_RESPONSE_TIMEOUT_SECS: u64 = 600; + pub struct Context<'a> { pub config: &'a mut (dyn Config + Send + Sync + 'a), pub io: crate::iostreams::IoStreams, @@ -213,6 +228,145 @@ impl Context<'_> { Ok(engine) } + /// Should KCL be executed on the server (true)? + /// Or locally (false)? + pub(crate) fn use_server_kcl_execution() -> bool { + std::env::var(ENGINE_EXECUTION_ENV) + .map(|value| !value.is_empty()) + .unwrap_or_default() + } + + async fn engine_ws_with_settings( + &self, + hostname: &str, + settings: &kcl_lib::ExecutorSettings, + ) -> Result { + let client = self.api_client(hostname)?; + let pr = std::env::var("ZOO_ENGINE_PR").ok().and_then(|s| s.parse().ok()); + let (ws, _headers) = client + .modeling() + .commands_ws(kittycad::modeling::CommandsWsParams { + api_call_id: None, + fps: None, + order_independent_transparency: None, + pool: None, + post_effect: if settings.enable_ssao { + Some(kittycad::types::PostEffectType::Ssao) + } else { + None + }, + pr, + replay: settings.replay.clone(), + show_grid: if settings.show_grid { Some(true) } else { None }, + unlocked_framerate: None, + video_res_height: None, + video_res_width: None, + webrtc: Some(false), + }) + .await?; + Ok(ws) + } + + #[allow(clippy::too_many_arguments)] + pub(crate) async fn run_server_kcl_then_modeling_cmds( + &mut self, + hostname: &str, + _filename: &str, + filepath: &Path, + code: &str, + cmds: Vec, + settings: kcl_lib::ExecutorSettings, + issue_check: kcl_error_fmt::KclIssueCheck, + ) -> Result<(Vec, Option)> { + let project = build_kcl_project(filepath, code)?; + let ws = self.engine_ws_with_settings(hostname, &settings).await?; + let wsconfig = tokio_tungstenite::tungstenite::protocol::WebSocketConfig::default() + .max_message_size(Some(usize::MAX)) + .max_frame_size(Some(usize::MAX)); + let ws_stream = WebSocketStream::from_raw_socket(ws, Role::Client, Some(wsconfig)).await; + let (mut write, mut read) = ws_stream.split(); + let mut session_data = None; + let mut heartbeat = + tokio::time::interval(Duration::from_secs(settings.heartbeats.unwrap_or(cmd_kcl::HEARTBEATS))); + + let exec_request_id = uuid::Uuid::new_v4(); + send_ws_request( + &mut write, + WebSocketRequest::ExecKclProject { + request_id: exec_request_id, + project, + }, + ) + .await?; + + loop { + let resp = read_ws_response_with_heartbeat(&mut read, &mut write, &mut heartbeat).await?; + if let Some(session) = update_session_data(&resp) { + session_data = Some(session); + continue; + } + + if response_request_id(&resp) != Some(exec_request_id) { + continue; + } + + let WebSocketResponse::Success(SuccessWebSocketResponse { + resp: OkWebSocketResponseData::ExecKclProject { result }, + .. + }) = resp + else { + return Err(websocket_failure_to_anyhow(resp)); + }; + + match result { + Ok(_) => break, + Err(err) => { + check_server_compilation_issues(&mut self.io.err_out, &err.non_fatal, issue_check)?; + if let Some(error) = err.error { + return Err(anyhow!("KCL execution failed: {}", error.get_message())); + } + break; + } + } + } + + let mut responses = Vec::with_capacity(cmds.len()); + for cmd in cmds { + let cmd_id = uuid::Uuid::new_v4(); + send_ws_request( + &mut write, + WebSocketRequest::ModelingCmdReq(ModelingCmdReq { + cmd, + cmd_id: cmd_id.into(), + }), + ) + .await?; + + loop { + let resp = read_ws_response_with_heartbeat(&mut read, &mut write, &mut heartbeat).await?; + if let Some(session) = update_session_data(&resp) { + session_data = Some(session); + continue; + } + + if response_request_id(&resp) != Some(cmd_id) { + continue; + } + + match resp { + WebSocketResponse::Success(SuccessWebSocketResponse { resp, .. }) => { + responses.push(resp); + break; + } + WebSocketResponse::Failure(_) => return Err(websocket_failure_to_anyhow(resp)), + } + } + } + + let _ = write.send(WsMsg::Close(None)).await; + Ok((responses, session_data)) + } + pub async fn send_kcl_modeling_cmd( &mut self, hostname: &str, @@ -222,6 +376,33 @@ impl Context<'_> { settings: kcl_lib::ExecutorSettings, issue_check: kcl_error_fmt::KclIssueCheck, ) -> Result<(OkWebSocketResponseData, Option)> { + if Self::use_server_kcl_execution() { + let (mut responses, session_data) = self + .run_server_kcl_then_modeling_cmds( + hostname, + filename, + Path::new(filename), + code, + vec![ + ModelingCmd::from( + mcmd::ZoomToFit::builder() + .animated(false) + .object_ids(Default::default()) + .padding(0.1) + .build(), + ), + cmd, + ], + settings, + issue_check, + ) + .await?; + let resp = responses + .pop() + .ok_or_else(|| anyhow!("Expected response from engine after executing KCL"))?; + return Ok((resp, session_data)); + } + let client = self.api_client(hostname)?; let program = kcl_lib::Program::parse_no_errs(code) @@ -276,6 +457,20 @@ impl Context<'_> { settings: kcl_lib::ExecutorSettings, issue_check: kcl_error_fmt::KclIssueCheck, ) -> Result<(Vec, Option)> { + if Self::use_server_kcl_execution() { + return self + .run_server_kcl_then_modeling_cmds( + hostname, + filename, + Path::new(filename), + code, + cmds, + settings, + issue_check, + ) + .await; + } + let client = self.api_client(hostname)?; let program = kcl_lib::Program::parse_no_errs(code) @@ -320,6 +515,31 @@ impl Context<'_> { settings: kcl_lib::ExecutorSettings, issue_check: kcl_error_fmt::KclIssueCheck, ) -> Result<(Vec, Option)> { + if Self::use_server_kcl_execution() { + let (responses, session_data) = self + .run_server_kcl_then_modeling_cmds( + hostname, + filename, + Path::new(filename), + code, + snapshot_cmds, + settings, + issue_check, + ) + .await?; + let mut snapshot_resps = Vec::new(); + for resp in responses { + if let OkWebSocketResponseData::Modeling { + modeling_response: kittycad_modeling_cmds::ok_response::OkModelingCmdResponse::TakeSnapshot(snap), + } = resp + { + snapshot_resps.push(snap); + } + } + + return Ok((snapshot_resps, session_data)); + } + let client = self.api_client(hostname)?; let program = kcl_lib::Program::parse_no_errs(code) @@ -1012,6 +1232,178 @@ pub(crate) fn reasoning_to_markdown(reason: &kittycad::types::ReasoningMessage) } } +fn build_kcl_project(filepath: &Path, code: &str) -> Result { + if filepath.to_str() == Some("-") { + let entrypoint = safe_filepath("main.kcl")?; + let file = KclFile::builder() + .path(entrypoint.clone()) + .contents(code.as_bytes().to_vec()) + .build(); + return Ok(KclProject::builder().files(vec![file]).entrypoint(entrypoint).build()); + } + + let project_root = filepath + .parent() + .filter(|path| !path.as_os_str().is_empty()) + .unwrap_or_else(|| Path::new(".")); + let entrypoint_path = relative_project_path(project_root, filepath)?; + let entrypoint = safe_filepath(&entrypoint_path)?; + let mut files = Vec::new(); + let mut found_entrypoint = false; + + for entry in ignore::WalkBuilder::new(project_root).hidden(false).build() { + let entry = entry?; + if !entry.file_type().is_some_and(|file_type| file_type.is_file()) { + continue; + } + if entry.path().extension().and_then(|ext| ext.to_str()) != Some("kcl") { + continue; + } + + let relative_path = relative_project_path(project_root, entry.path())?; + let contents = if relative_path == entrypoint_path { + found_entrypoint = true; + code.as_bytes().to_vec() + } else { + std::fs::read(entry.path())? + }; + + files.push( + KclFile::builder() + .path(safe_filepath(&relative_path)?) + .contents(contents) + .build(), + ); + } + + if !found_entrypoint { + files.push( + KclFile::builder() + .path(entrypoint.clone()) + .contents(code.as_bytes().to_vec()) + .build(), + ); + } + + Ok(KclProject::builder().files(files).entrypoint(entrypoint).build()) +} + +fn check_server_compilation_issues( + err_out: &mut impl std::io::Write, + issues: &[kcl_error::CompilationIssue], + issue_check: kcl_error_fmt::KclIssueCheck, +) -> Result<()> { + if issue_check == kcl_error_fmt::KclIssueCheck::Ignore || issues.is_empty() { + return Ok(()); + } + + for issue in issues { + writeln!(err_out, "{:?}: {}", issue.severity, issue.message)?; + } + + if issue_check == kcl_error_fmt::KclIssueCheck::DenyErrors && issues.iter().any(|issue| issue.is_err()) { + anyhow::bail!( + "KCL execution reported errors. Please fix your KCL program before continuing. If you really want to proceed anyway, rerun this command with `--allow-errors`." + ); + } + + Ok(()) +} + +fn relative_project_path(project_root: &Path, path: &Path) -> Result { + let relative = if project_root == Path::new(".") { + path + } else { + path.strip_prefix(project_root)? + }; + Ok(relative.to_string_lossy().replace('\\', "/")) +} + +fn safe_filepath(path: &str) -> Result { + path.parse::() + .map_err(|err| anyhow!("invalid KCL project path `{path}`: {err}")) +} + +async fn send_ws_request(write: &mut DirectWsWrite, request: WebSocketRequest) -> Result<()> { + let msg = serde_json::to_string(&request)?; + write + .send(WsMsg::Text(msg.into())) + .await + .map_err(|err| anyhow!("could not send request to engine websocket: {err}"))?; + Ok(()) +} + +async fn read_ws_response_with_heartbeat( + read: &mut DirectWsRead, + write: &mut DirectWsWrite, + heartbeat: &mut tokio::time::Interval, +) -> Result { + let timeout = tokio::time::sleep(Duration::from_secs(WS_RESPONSE_TIMEOUT_SECS)); + tokio::pin!(timeout); + + loop { + tokio::select! { + maybe_msg = read.next() => { + let Some(msg) = maybe_msg else { + anyhow::bail!("engine websocket closed before sending a response"); + }; + return parse_ws_msg(msg?); + } + _ = heartbeat.tick() => { + send_ws_request(write, WebSocketRequest::Ping {}).await?; + } + _ = &mut timeout => { + anyhow::bail!("engine websocket response timed out after {WS_RESPONSE_TIMEOUT_SECS}s"); + } + } + } +} + +fn parse_ws_msg(msg: WsMsg) -> Result { + match msg { + WsMsg::Text(text) => Ok(serde_json::from_str(&text)?), + WsMsg::Binary(bin) => Ok(rmp_serde::from_slice(&bin)?), + other => anyhow::bail!("unexpected engine websocket message: {other}"), + } +} + +fn update_session_data(response: &WebSocketResponse) -> Option { + match response { + WebSocketResponse::Success(SuccessWebSocketResponse { + resp: OkWebSocketResponseData::ModelingSessionData { session }, + .. + }) => Some(session.clone()), + _ => None, + } +} + +fn response_request_id(response: &WebSocketResponse) -> Option { + match response { + WebSocketResponse::Success(SuccessWebSocketResponse { request_id, .. }) => *request_id, + WebSocketResponse::Failure(FailureWebSocketResponse { request_id, .. }) => *request_id, + } +} + +fn websocket_failure_to_anyhow(response: WebSocketResponse) -> anyhow::Error { + match response { + WebSocketResponse::Failure(FailureWebSocketResponse { errors, .. }) => { + if errors.is_empty() { + anyhow!("engine websocket request failed with no error details") + } else { + anyhow!( + "{}", + errors + .into_iter() + .map(|error| error.message) + .collect::>() + .join("\n") + ) + } + } + other => anyhow!("unexpected engine websocket response: {other:?}"), + } +} + fn indent_block(s: &str) -> String { let mut out = String::new(); for line in s.lines() {