Skip to main content

kcl_lib/engine/
engine_manager.rs

1use std::collections::HashMap;
2use std::sync::Arc;
3use std::sync::atomic::Ordering::Relaxed;
4
5use anyhow::Result;
6pub use engine_transport::EngineTransport;
7use indexmap::IndexMap;
8use kcmc::ModelingCmd;
9use kcmc::each_cmd as mcmd;
10use kcmc::shared::Color;
11use kcmc::websocket::BatchResponse;
12use kcmc::websocket::ModelingCmdReq;
13use kcmc::websocket::ModelingSessionData;
14use kcmc::websocket::OkWebSocketResponseData;
15use kcmc::websocket::WebSocketRequest;
16use kcmc::websocket::WebSocketResponse;
17use kittycad_modeling_cmds::ModelingCmdEndpoint;
18use kittycad_modeling_cmds::length_unit::LengthUnit;
19use kittycad_modeling_cmds::ok_response::OkModelingCmdResponse;
20use kittycad_modeling_cmds::websocket::ModelingBatch;
21use kittycad_modeling_cmds::{self as kcmc};
22use tokio::sync::RwLock;
23use uuid::Uuid;
24use web_time::Instant;
25
26use crate::ExecutorSettings;
27use crate::SourceRange;
28use crate::engine::AsyncTasks;
29use crate::engine::DEFAULT_PLANE_INFO;
30use crate::engine::EngineBatchContext;
31use crate::engine::EngineStats;
32use crate::engine::GRID_OBJECT_ID;
33use crate::engine::GRID_SCALE_TEXT_OBJECT_ID;
34use crate::engine::GridScaleBehavior;
35use crate::engine::PlaneName;
36use crate::errors::KclError;
37use crate::errors::KclErrorDetails;
38use crate::execution::DefaultPlanes;
39use crate::execution::IdGenerator;
40use crate::execution::PlaneInfo;
41use crate::settings::types::default_backface_color;
42use crate::settings::types::default_backface_color_struct;
43
44pub enum TransportCloseError {}
45
46mod engine_transport;
47mod mock_transport;
48#[cfg(target_arch = "wasm32")]
49pub mod wasm_transport;
50#[cfg(not(target_arch = "wasm32"))]
51pub mod ws_transport;
52
53/// Information about the responses from the engine.
54#[derive(Clone, Debug)]
55pub struct ResponseInformation {
56    /// The responses from the engine.
57    responses: Arc<RwLock<IndexMap<uuid::Uuid, WebSocketResponse>>>,
58}
59
60impl ResponseInformation {
61    /// Basic constructor.
62    pub fn new(responses: Arc<RwLock<IndexMap<uuid::Uuid, WebSocketResponse>>>) -> Self {
63        Self { responses }
64    }
65
66    /// Add a new response from the engine.
67    pub async fn add(&self, id: Uuid, response: WebSocketResponse) {
68        self.responses.write().await.insert(id, response);
69    }
70}
71
72#[derive(bon::Builder)]
73pub struct EngineManager {
74    // Replaces `engine_req_tx: mpsc::Sender<ToEngineReq>`
75    // from the original native connection type.
76    pub transport: Arc<Box<dyn EngineTransport>>,
77    responses: ResponseInformation,
78    pending_errors: Arc<RwLock<Vec<String>>>,
79    socket_health: Arc<RwLock<SocketHealth>>,
80    ids_of_async_commands: Arc<RwLock<IndexMap<Uuid, SourceRange>>>,
81
82    /// The default planes for the scene.
83    #[builder(default)]
84    default_planes: Arc<RwLock<Option<DefaultPlanes>>>,
85    /// If the server sends session data, it'll be copied to here.
86    session_data: Arc<RwLock<Option<ModelingSessionData>>>,
87
88    #[builder(default)]
89    stats: EngineStats,
90
91    #[builder(default)]
92    async_tasks: AsyncTasks,
93}
94
95impl std::fmt::Debug for EngineManager {
96    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
97        f.debug_struct("EngineManager")
98            .field("responses", &self.responses)
99            .field("pending_errors", &self.pending_errors)
100            .field("socket_health", &self.socket_health)
101            .field("ids_of_async_commands", &self.ids_of_async_commands)
102            .field("default_planes", &self.default_planes)
103            .field("session_data", &self.session_data)
104            .field("stats", &self.stats)
105            .field("async_tasks", &self.async_tasks)
106            .finish()
107    }
108}
109
110impl EngineManager {
111    #[cfg(target_arch = "wasm32")]
112    pub fn new_wasm_transport(
113        manager: wasm_transport::EngineCommandManager,
114        response_context: Arc<wasm_transport::ResponseContext>,
115    ) -> Self {
116        let session_data: Arc<RwLock<Option<ModelingSessionData>>> = Arc::new(RwLock::new(None));
117        let ids_of_async_commands: Arc<RwLock<IndexMap<Uuid, SourceRange>>> = Arc::new(RwLock::new(IndexMap::new()));
118        let socket_health = Arc::new(RwLock::new(SocketHealth::Active));
119        let pending_errors = Arc::new(RwLock::new(Vec::new()));
120        let responses = response_context.response_information();
121
122        Self {
123            transport: Arc::new(Box::new(wasm_transport::WasmTransport::new(manager))),
124            responses,
125            pending_errors,
126            socket_health,
127            ids_of_async_commands,
128            default_planes: Default::default(),
129            session_data,
130            stats: Default::default(),
131            async_tasks: Default::default(),
132        }
133    }
134
135    #[cfg(not(target_arch = "wasm32"))]
136    pub async fn new_websocket_transport(ws: reqwest::Upgraded, heartbeats: Option<u64>) -> Self {
137        Self::new_websocket_transport_with_request_id(ws, heartbeats, None).await
138    }
139
140    #[cfg(not(target_arch = "wasm32"))]
141    pub(crate) async fn new_websocket_transport_with_request_id(
142        ws: reqwest::Upgraded,
143        heartbeats: Option<u64>,
144        request_id: Option<String>,
145    ) -> Self {
146        use crate::engine::engine_manager::ws_transport::WebSocketTransport;
147
148        let session_data: Arc<RwLock<Option<ModelingSessionData>>> = Arc::new(RwLock::new(None));
149        let ids_of_async_commands: Arc<RwLock<IndexMap<Uuid, SourceRange>>> = Arc::new(RwLock::new(IndexMap::new()));
150        let socket_health = Arc::new(RwLock::new(SocketHealth::Active));
151        let pending_errors = Arc::new(RwLock::new(Vec::new()));
152        let responses = ResponseInformation {
153            responses: Arc::new(RwLock::new(IndexMap::new())),
154        };
155
156        let transport = WebSocketTransport::spawn(
157            ws,
158            heartbeats,
159            responses.clone(),
160            Arc::clone(&session_data),
161            Arc::clone(&pending_errors),
162            Arc::clone(&socket_health),
163            request_id,
164        )
165        .await;
166
167        Self {
168            transport: Arc::new(Box::new(transport)),
169            responses,
170            pending_errors,
171            socket_health,
172            ids_of_async_commands,
173            default_planes: Default::default(),
174            session_data,
175            stats: Default::default(),
176            async_tasks: Default::default(),
177        }
178    }
179
180    /// Mock connection that doesn't actually connect to anything.
181    /// Used for testing.
182    pub fn new_mock() -> Self {
183        let session_data: Arc<RwLock<Option<ModelingSessionData>>> = Arc::new(RwLock::new(None));
184        let ids_of_async_commands: Arc<RwLock<IndexMap<Uuid, SourceRange>>> = Arc::new(RwLock::new(IndexMap::new()));
185        let socket_health = Arc::new(RwLock::new(SocketHealth::Active));
186        let pending_errors = Arc::new(RwLock::new(Vec::new()));
187        let responses = ResponseInformation {
188            responses: Arc::new(RwLock::new(IndexMap::new())),
189        };
190        Self {
191            transport: Arc::new(Box::new(mock_transport::MockTransport::new(responses.clone()))),
192            responses,
193            pending_errors,
194            socket_health,
195            ids_of_async_commands,
196            default_planes: Default::default(),
197            session_data,
198            stats: Default::default(),
199            async_tasks: Default::default(),
200        }
201    }
202
203    /// Take the ids of async commands that have accumulated so far and clear them.
204    async fn take_ids_of_async_commands(&self) -> IndexMap<Uuid, SourceRange> {
205        std::mem::take(&mut *self.ids_of_async_commands().write().await)
206    }
207
208    /// Take the responses that have accumulated so far and clear them.
209    pub async fn take_responses(&self) -> IndexMap<Uuid, WebSocketResponse> {
210        std::mem::take(&mut *self.responses().write().await)
211    }
212
213    pub async fn clear_scene(
214        &self,
215        batch_context: &EngineBatchContext,
216        id_generator: &mut IdGenerator,
217        source_range: SourceRange,
218        geometry_only: bool,
219    ) -> Result<(), crate::errors::KclError> {
220        // Clear any batched commands leftover from previous scenes.
221        self.clear_queues(batch_context).await;
222
223        self.batch_modeling_cmd(
224            batch_context,
225            id_generator.next_uuid(),
226            source_range,
227            &ModelingCmd::SceneClearAll(mcmd::SceneClearAll::default()),
228        )
229        .await?;
230
231        // Flush the batch queue, so clear is run right away.
232        // Otherwise the hooks below won't work.
233        self.flush_batch(batch_context, false, source_range).await?;
234
235        // Do the after clear scene hook.
236        self.clear_scene_post_hook(batch_context, id_generator, source_range, geometry_only)
237            .await?;
238
239        Ok(())
240    }
241
242    /// Ensure a specific async command has been completed.
243    pub async fn ensure_async_command_completed(
244        &self,
245        id: uuid::Uuid,
246        source_range: Option<SourceRange>,
247    ) -> Result<OkWebSocketResponseData, KclError> {
248        let source_range = if let Some(source_range) = source_range {
249            source_range
250        } else {
251            // Look it up if we don't have it.
252            self.ids_of_async_commands()
253                .read()
254                .await
255                .get(&id)
256                .cloned()
257                .unwrap_or_default()
258        };
259
260        // The previous 60s ceiling here was too aggressive for long-running
261        // engine commands - notably large STEP / B-rep imports, which the
262        // engine itself routinely takes several minutes to process. When the
263        // ceiling fired first the user got a generic "async command timed
264        // out" message and the eventual engine response (success OR error)
265        // was discarded, masking the real outcome. 600s (10 min) gives the
266        // engine room to finish or to surface its own error.
267        const ASYNC_CMD_TIMEOUT_SECS: u64 = 600;
268        let current_time = Instant::now();
269        while current_time.elapsed().as_secs() < ASYNC_CMD_TIMEOUT_SECS {
270            let responses = self.responses().read().await.clone();
271            let Some(resp) = responses.get(&id) else {
272                // Yield to the event loop so that we don’t block the UI thread.
273                // No seriously WE DO NOT WANT TO PAUSE THE WHOLE APP ON THE JS SIDE.
274                #[cfg(target_arch = "wasm32")]
275                {
276                    let duration = web_time::Duration::from_millis(1);
277                    wasm_timer::Delay::new(duration).await.map_err(|err| {
278                        KclError::new_internal(KclErrorDetails::new(
279                            format!("Failed to sleep: {:?}", err),
280                            vec![source_range],
281                        ))
282                    })?;
283                }
284                #[cfg(not(target_arch = "wasm32"))]
285                tokio::task::yield_now().await;
286                continue;
287            };
288
289            // If the response is an error, return it.
290            // Parsing will do that and we can ignore the result, we don't care.
291            let response = self.parse_websocket_response(resp.clone(), source_range)?;
292            return Ok(response);
293        }
294
295        Err(KclError::new_engine(KclErrorDetails::new(
296            format!(
297                "async command timed out after {ASYNC_CMD_TIMEOUT_SECS}s (client-side ceiling, not an engine error)"
298            ),
299            vec![source_range],
300        )))
301    }
302
303    /// Ensure ALL async commands have been completed.
304    pub async fn ensure_async_commands_completed(&self, batch_context: &EngineBatchContext) -> Result<(), KclError> {
305        // Check if all async commands have been completed.
306        let ids = self.take_ids_of_async_commands().await;
307
308        // Try to get them from the responses.
309        for (id, source_range) in ids {
310            self.ensure_async_command_completed(id, Some(source_range)).await?;
311        }
312
313        // Make sure we check for all async tasks as well.
314        // The reason why we ignore the error here is that, if a model fillets an edge
315        // we previously called something on, it might no longer exist. In which case,
316        // the artifact graph won't care either if its gone since you can't select it
317        // anymore anyways.
318        if let Err(err) = self.async_tasks().join_all().await {
319            crate::log::logln!(
320                "Error waiting for async tasks (this is typically fine and just means that an edge became something else): {:?}",
321                err
322            );
323        }
324
325        // Flush the batch to make sure nothing remains.
326        self.flush_batch(batch_context, true, SourceRange::default()).await?;
327
328        Ok(())
329    }
330
331    /// Set the visibility of edges.
332    async fn set_edge_visibility(
333        &self,
334        batch_context: &EngineBatchContext,
335        visible: bool,
336        source_range: SourceRange,
337        id_generator: &mut IdGenerator,
338    ) -> Result<(), crate::errors::KclError> {
339        self.batch_modeling_cmd(
340            batch_context,
341            id_generator.next_uuid(),
342            source_range,
343            &ModelingCmd::from(mcmd::EdgeLinesVisible::builder().hidden(!visible).build()),
344        )
345        .await?;
346
347        Ok(())
348    }
349
350    /// Re-run the command to apply the settings.
351    pub async fn reapply_settings(
352        &self,
353        batch_context: &EngineBatchContext,
354        settings: &crate::ExecutorSettings,
355        source_range: SourceRange,
356        id_generator: &mut IdGenerator,
357        grid_scale_unit: GridScaleBehavior,
358    ) -> Result<(), crate::errors::KclError> {
359        if settings.geometry_only {
360            return Ok(());
361        }
362        // Set the edge visibility.
363        self.set_edge_visibility(batch_context, settings.highlight_edges, source_range, id_generator)
364            .await?;
365
366        // Send the command to show the grid.
367
368        self.modify_grid(
369            batch_context,
370            !settings.show_grid,
371            grid_scale_unit,
372            source_range,
373            id_generator,
374        )
375        .await?;
376
377        // Set up user's color choices.
378        self.set_user_colors(batch_context, settings, source_range, id_generator)
379            .await?;
380
381        // We do not have commands for changing ssao on the fly.
382
383        // Flush the batch queue, so the settings are applied right away.
384        self.flush_batch(batch_context, false, source_range).await?;
385
386        Ok(())
387    }
388
389    // Add a modeling command to the batch but don't fire it right away.
390    pub async fn batch_modeling_cmd(
391        &self,
392        batch_context: &EngineBatchContext,
393        id: uuid::Uuid,
394        source_range: SourceRange,
395        cmd: &ModelingCmd,
396    ) -> Result<(), crate::errors::KclError> {
397        let req = WebSocketRequest::ModelingCmdReq(ModelingCmdReq {
398            cmd: cmd.clone(),
399            cmd_id: id.into(),
400        });
401
402        // Add cmd to the batch.
403        batch_context.push(req, source_range).await;
404        self.stats().commands_batched.fetch_add(1, Relaxed);
405
406        Ok(())
407    }
408
409    // Add a vector of modeling commands to the batch but don't fire it right away.
410    // This allows you to force them all to be added together in the same order.
411    // When we are running things in parallel this prevents race conditions that might come
412    // if specific commands are run before others.
413    pub async fn batch_modeling_cmds(
414        &self,
415        batch_context: &EngineBatchContext,
416        source_range: SourceRange,
417        cmds: &[ModelingCmdReq],
418    ) -> Result<(), crate::errors::KclError> {
419        // Add cmds to the batch.
420        let mut extended_cmds = Vec::with_capacity(cmds.len());
421        for cmd in cmds {
422            extended_cmds.push((WebSocketRequest::ModelingCmdReq(cmd.clone()), source_range));
423        }
424        self.stats().commands_batched.fetch_add(extended_cmds.len(), Relaxed);
425        batch_context.extend(extended_cmds).await;
426
427        Ok(())
428    }
429
430    /// Add a command to the batch that needs to be executed at the very end.
431    /// This for stuff like fillets or chamfers where if we execute too soon the
432    /// engine will eat the ID and we can't reference it for other commands.
433    pub async fn batch_end_cmd(
434        &self,
435        batch_context: &EngineBatchContext,
436        id: uuid::Uuid,
437        source_range: SourceRange,
438        cmd: &ModelingCmd,
439    ) -> Result<(), crate::errors::KclError> {
440        let req = WebSocketRequest::ModelingCmdReq(ModelingCmdReq {
441            cmd: cmd.clone(),
442            cmd_id: id.into(),
443        });
444
445        // Add cmd to the batch end.
446        batch_context.insert_end(id, req, source_range).await;
447        self.stats().commands_batched.fetch_add(1, Relaxed);
448        Ok(())
449    }
450
451    /// Send the modeling cmd and wait for the response.
452    pub async fn send_modeling_cmd(
453        &self,
454        batch_context: &EngineBatchContext,
455        id: uuid::Uuid,
456        source_range: SourceRange,
457        cmd: &ModelingCmd,
458    ) -> Result<OkWebSocketResponseData, crate::errors::KclError> {
459        let mut requests = batch_context.take_batch().await;
460
461        // Add the command to the batch.
462        requests.push((
463            WebSocketRequest::ModelingCmdReq(ModelingCmdReq {
464                cmd: cmd.clone(),
465                cmd_id: id.into(),
466            }),
467            source_range,
468        ));
469        self.stats().commands_batched.fetch_add(1, Relaxed);
470
471        // Flush the batch queue.
472        self.run_batch(requests, source_range).await
473    }
474
475    /// Send the modeling cmd async and don't wait for the response.
476    /// Add it to our list of async commands.
477    pub async fn async_modeling_cmd(
478        &self,
479        id: uuid::Uuid,
480        source_range: SourceRange,
481        cmd: &ModelingCmd,
482    ) -> Result<(), crate::errors::KclError> {
483        // Add the command ID to the list of async commands.
484        self.ids_of_async_commands().write().await.insert(id, source_range);
485
486        // Fire off the command now, but don't wait for the response, we don't care about it.
487        self.transport
488            .inner_fire_modeling_cmd(
489                id,
490                source_range,
491                WebSocketRequest::ModelingCmdReq(ModelingCmdReq {
492                    cmd: cmd.clone(),
493                    cmd_id: id.into(),
494                }),
495                HashMap::from([(id, source_range)]),
496            )
497            .await?;
498
499        Ok(())
500    }
501
502    /// Run the batch for the specific commands.
503    async fn run_batch(
504        &self,
505        orig_requests: Vec<(WebSocketRequest, SourceRange)>,
506        source_range: SourceRange,
507    ) -> Result<OkWebSocketResponseData, crate::errors::KclError> {
508        // Return early if we have no commands to send.
509        if orig_requests.is_empty() {
510            return Ok(OkWebSocketResponseData::Modeling {
511                modeling_response: OkModelingCmdResponse::Empty {},
512            });
513        }
514
515        let requests: Vec<ModelingCmdReq> = orig_requests
516            .iter()
517            .filter_map(|(val, _)| match val {
518                WebSocketRequest::ModelingCmdReq(ModelingCmdReq { cmd, cmd_id }) => Some(ModelingCmdReq {
519                    cmd: cmd.clone(),
520                    cmd_id: *cmd_id,
521                }),
522                _ => None,
523            })
524            .collect();
525
526        let batched_requests = WebSocketRequest::ModelingCmdBatchReq(ModelingBatch {
527            requests,
528            batch_id: uuid::Uuid::new_v4().into(),
529            responses: true,
530        });
531
532        let final_req = if orig_requests.len() == 1 {
533            // We can unwrap here because we know the batch has only one element.
534            orig_requests.first().unwrap().0.clone()
535        } else {
536            batched_requests
537        };
538
539        // Create the map of original command IDs to source range.
540        // This is for the wasm side, kurt needs it for selections.
541        let mut id_to_source_range = HashMap::new();
542
543        let mut id_to_command = HashMap::new();
544        for (req, range) in orig_requests.iter() {
545            match req {
546                WebSocketRequest::ModelingCmdReq(ModelingCmdReq { cmd, cmd_id }) => {
547                    let id = Uuid::from(*cmd_id);
548                    id_to_source_range.insert(id, *range);
549                    id_to_command.insert(id, ModelingCmdEndpoint::from(cmd));
550                }
551                _ => {
552                    return Err(KclError::new_engine(KclErrorDetails::new(
553                        format!("The request is not a modeling command: {req:?}"),
554                        vec![*range],
555                    )));
556                }
557            }
558        }
559
560        self.stats().batches_sent.fetch_add(1, Relaxed);
561
562        // We pop off the responses to cleanup our mappings.
563        match final_req {
564            WebSocketRequest::ModelingCmdBatchReq(ModelingBatch {
565                ref requests,
566                batch_id,
567                responses: _,
568            }) => {
569                // Get the last command ID.
570                let last_id = requests.last().unwrap().cmd_id;
571                let ws_resp = self
572                    .inner_send_modeling_cmd(batch_id.into(), source_range, final_req, id_to_source_range.clone())
573                    .await?;
574                let response = self.parse_websocket_response(ws_resp, source_range)?;
575
576                // If we have a batch response, we want to return the specific id we care about.
577                if let OkWebSocketResponseData::ModelingBatch { responses } = response {
578                    self.parse_batch_responses(last_id.into(), id_to_source_range, id_to_command, responses)
579                } else {
580                    // We should never get here.
581                    Err(KclError::new_engine(KclErrorDetails::new(
582                        format!("Failed to get batch response: {response:?}"),
583                        vec![source_range],
584                    )))
585                }
586            }
587            WebSocketRequest::ModelingCmdReq(ModelingCmdReq { cmd: _, cmd_id }) => {
588                // You are probably wondering why we can't just return the source range we were
589                // passed with the function. Well this is actually really important.
590                // If this is the last command in the batch and there is only one and we've reached
591                // the end of the file, this will trigger a flush batch function, but it will just
592                // send default or the end of the file as it's source range not the origin of the
593                // request so we need the original request source range in case the engine returns
594                // an error.
595                let source_range = id_to_source_range.get(cmd_id.as_ref()).cloned().ok_or_else(|| {
596                    KclError::new_engine(KclErrorDetails::new(
597                        format!("Failed to get source range for command ID: {cmd_id:?}"),
598                        vec![],
599                    ))
600                })?;
601                let ws_resp = self
602                    .inner_send_modeling_cmd(cmd_id.into(), source_range, final_req, id_to_source_range)
603                    .await?;
604                self.parse_websocket_response(ws_resp, source_range)
605            }
606            _ => Err(KclError::new_engine(KclErrorDetails::new(
607                format!("The final request is not a modeling command: {final_req:?}"),
608                vec![source_range],
609            ))),
610        }
611    }
612
613    /// Force flush the batch queue.
614    pub async fn flush_batch(
615        &self,
616        batch_context: &EngineBatchContext,
617        // Whether or not to flush the end commands as well.
618        // We only do this at the very end of the file.
619        batch_end: bool,
620        source_range: SourceRange,
621    ) -> Result<OkWebSocketResponseData, crate::errors::KclError> {
622        let all_requests = if batch_end {
623            let mut requests = batch_context.take_batch().await;
624            requests.extend(batch_context.take_batch_end().await.values().cloned());
625            requests
626        } else {
627            batch_context.take_batch().await
628        };
629
630        self.run_batch(all_requests, source_range).await
631    }
632
633    async fn make_default_plane(
634        &self,
635        batch_context: &EngineBatchContext,
636        plane_id: uuid::Uuid,
637        info: &PlaneInfo,
638        color: Option<Color>,
639        source_range: SourceRange,
640        id_generator: &mut IdGenerator,
641    ) -> Result<uuid::Uuid, KclError> {
642        // Create new default planes.
643        let default_size = 100.0;
644
645        self.batch_modeling_cmd(
646            batch_context,
647            plane_id,
648            source_range,
649            &ModelingCmd::from(
650                mcmd::MakePlane::builder()
651                    .clobber(false)
652                    .origin(info.origin.into())
653                    .size(LengthUnit(default_size))
654                    .x_axis(info.x_axis.into())
655                    .y_axis(info.y_axis.into())
656                    .hide(true)
657                    .build(),
658            ),
659        )
660        .await?;
661
662        if let Some(color) = color {
663            // Set the color.
664            self.batch_modeling_cmd(
665                batch_context,
666                id_generator.next_uuid(),
667                source_range,
668                &ModelingCmd::from(mcmd::PlaneSetColor::builder().color(color).plane_id(plane_id).build()),
669            )
670            .await?;
671        }
672
673        Ok(plane_id)
674    }
675
676    async fn new_default_planes(
677        &self,
678        batch_context: &EngineBatchContext,
679        id_generator: &mut IdGenerator,
680        source_range: SourceRange,
681        geometry_only: bool,
682    ) -> Result<DefaultPlanes, KclError> {
683        let plane_opacity = 0.1;
684        let plane_color =
685            |red, green, blue| (!geometry_only).then(|| Color::from_rgba(red, green, blue, plane_opacity));
686        let plane_settings: Vec<(PlaneName, Uuid, Option<Color>)> = vec![
687            (PlaneName::Xy, id_generator.next_uuid(), plane_color(0.7, 0.28, 0.28)),
688            (PlaneName::Yz, id_generator.next_uuid(), plane_color(0.28, 0.7, 0.28)),
689            (PlaneName::Xz, id_generator.next_uuid(), plane_color(0.28, 0.28, 0.7)),
690            (PlaneName::NegXy, id_generator.next_uuid(), None),
691            (PlaneName::NegYz, id_generator.next_uuid(), None),
692            (PlaneName::NegXz, id_generator.next_uuid(), None),
693        ];
694
695        let mut planes = HashMap::new();
696        for (name, plane_id, color) in plane_settings {
697            let info = DEFAULT_PLANE_INFO.get(&name).ok_or_else(|| {
698                // We should never get here.
699                KclError::new_engine(KclErrorDetails::new(
700                    format!("Failed to get default plane info for: {name:?}"),
701                    vec![source_range],
702                ))
703            })?;
704            planes.insert(
705                name,
706                self.make_default_plane(batch_context, plane_id, info, color, source_range, id_generator)
707                    .await?,
708            );
709        }
710
711        // Flush the batch queue, so these planes are created right away.
712        self.flush_batch(batch_context, false, source_range).await?;
713
714        Ok(DefaultPlanes {
715            xy: planes[&PlaneName::Xy],
716            neg_xy: planes[&PlaneName::NegXy],
717            xz: planes[&PlaneName::Xz],
718            neg_xz: planes[&PlaneName::NegXz],
719            yz: planes[&PlaneName::Yz],
720            neg_yz: planes[&PlaneName::NegYz],
721        })
722    }
723
724    fn parse_websocket_response(
725        &self,
726        response: WebSocketResponse,
727        source_range: SourceRange,
728    ) -> Result<OkWebSocketResponseData, crate::errors::KclError> {
729        match response {
730            WebSocketResponse::Success(success) => Ok(success.resp),
731            WebSocketResponse::Failure(fail) => {
732                let _request_id = fail.request_id;
733                if fail.errors.is_empty() {
734                    return Err(KclError::new_engine(KclErrorDetails::new(
735                        "Failure response with no error details".to_owned(),
736                        vec![source_range],
737                    )));
738                }
739                Err(KclError::new_engine(KclErrorDetails::new(
740                    fail.errors
741                        .iter()
742                        .map(|e| e.message.clone())
743                        .collect::<Vec<_>>()
744                        .join("\n"),
745                    vec![source_range],
746                )))
747            }
748        }
749    }
750
751    fn parse_batch_responses(
752        &self,
753        // The last response we are looking for.
754        id: uuid::Uuid,
755        // The mapping of source ranges to command IDs.
756        id_to_source_range: HashMap<uuid::Uuid, SourceRange>,
757        // Allows us to print which command failed
758        id_to_command: HashMap<uuid::Uuid, ModelingCmdEndpoint>,
759        // The response from the engine.
760        responses: HashMap<kcmc::id::ModelingCmdId, BatchResponse>,
761    ) -> Result<OkWebSocketResponseData, crate::errors::KclError> {
762        let mut any_err: Option<crate::errors::KclError> = None;
763        let mut target_ok: Option<OkWebSocketResponseData> = None;
764        // Iterate over the responses and check for errors.
765        // Any error takes precedent over any Ok.
766        #[expect(
767            clippy::iter_over_hash_type,
768            reason = "modeling command uses a HashMap and keys are random, so we don't really have a choice"
769        )]
770        for (cmd_id, resp) in responses.iter() {
771            let cmd_id = Uuid::from(*cmd_id);
772            match resp {
773                BatchResponse::Success { response } if cmd_id == id => {
774                    // This is the response we care about.
775                    // Keep looking for errors after locating it.
776                    target_ok = Some(OkWebSocketResponseData::Modeling {
777                        modeling_response: response.clone(),
778                    });
779                }
780                BatchResponse::Success { .. } => continue,
781                BatchResponse::Failure { errors } => {
782                    let command = id_to_command
783                        .get(&cmd_id)
784                        .map(ModelingCmdEndpoint::to_string)
785                        .unwrap_or("[missing entry]".to_string());
786                    // Get the source range for the command.
787                    let source_range = id_to_source_range.get(&cmd_id).cloned().ok_or_else(|| {
788                        KclError::new_engine(KclErrorDetails::new(
789                            format!("Failed to get source range for command {command} with ID: {cmd_id:?}"),
790                            vec![],
791                        ))
792                    })?;
793                    if errors.is_empty() {
794                        any_err = Some(KclError::new_engine(KclErrorDetails::new(
795                            format!("Failure response for batch with no error details at command {command}"),
796                            vec![source_range],
797                        )));
798                        break;
799                    }
800                    let errors = errors.iter().map(|e| e.message.clone()).collect::<Vec<_>>().join("\n");
801                    any_err = Some(KclError::new_engine(KclErrorDetails::new(
802                        format!("command {command} resulted in errors: \n {errors}"),
803                        vec![source_range],
804                    )));
805                    break;
806                }
807            }
808        }
809
810        match (any_err, target_ok) {
811            (Some(err), _) => Err(err),
812            (None, Some(ok)) => Ok(ok),
813            (None, None) => {
814                // Return an error that we did not get an error or the response we wanted.
815                // This should never happen but who knows.
816                Err(KclError::new_engine(KclErrorDetails::new(
817                    format!("Failed to find response for command ID: {id:?}"),
818                    vec![],
819                )))
820            }
821        }
822    }
823
824    async fn set_user_colors(
825        &self,
826        batch_context: &EngineBatchContext,
827        settings: &ExecutorSettings,
828        source_range: SourceRange,
829        id_generator: &mut IdGenerator,
830    ) -> Result<(), KclError> {
831        let bf = settings
832            .default_backface_color
833            .clone()
834            .unwrap_or(default_backface_color());
835        let backface = csscolorparser::parse(&bf)
836            .map(|color| kcmc::shared::Color::from_rgba(color.r, color.g, color.b, color.a))
837            .unwrap_or(default_backface_color_struct());
838        self.batch_modeling_cmd(
839            batch_context,
840            id_generator.next_uuid(),
841            source_range,
842            &ModelingCmd::from(
843                mcmd::SetDefaultSystemProperties::builder()
844                    .backface_color(backface)
845                    .build(),
846            ),
847        )
848        .await?;
849        Ok(())
850    }
851
852    async fn modify_grid(
853        &self,
854        batch_context: &EngineBatchContext,
855        hidden: bool,
856        grid_scale_behavior: GridScaleBehavior,
857        source_range: SourceRange,
858        id_generator: &mut IdGenerator,
859    ) -> Result<(), KclError> {
860        // Hide/show the grid.
861        self.batch_modeling_cmd(
862            batch_context,
863            id_generator.next_uuid(),
864            source_range,
865            &ModelingCmd::from(
866                mcmd::ObjectVisible::builder()
867                    .hidden(hidden)
868                    .object_id(*GRID_OBJECT_ID)
869                    .build(),
870            ),
871        )
872        .await?;
873
874        self.batch_modeling_cmd(
875            batch_context,
876            id_generator.next_uuid(),
877            source_range,
878            &grid_scale_behavior.into_modeling_cmd(),
879        )
880        .await?;
881
882        // Hide/show the grid scale text.
883        self.batch_modeling_cmd(
884            batch_context,
885            id_generator.next_uuid(),
886            source_range,
887            &ModelingCmd::from(
888                mcmd::ObjectVisible::builder()
889                    .hidden(hidden)
890                    .object_id(*GRID_SCALE_TEXT_OBJECT_ID)
891                    .build(),
892            ),
893        )
894        .await?;
895
896        Ok(())
897    }
898
899    pub async fn clear_queues(&self, batch_context: &EngineBatchContext) {
900        batch_context.clear().await;
901        self.ids_of_async_commands().write().await.clear();
902        self.async_tasks().clear().await;
903    }
904
905    fn responses(&self) -> Arc<RwLock<IndexMap<Uuid, WebSocketResponse>>> {
906        self.responses.responses.clone()
907    }
908
909    fn ids_of_async_commands(&self) -> Arc<RwLock<IndexMap<Uuid, SourceRange>>> {
910        self.ids_of_async_commands.clone()
911    }
912
913    fn async_tasks(&self) -> AsyncTasks {
914        self.async_tasks.clone()
915    }
916
917    pub fn stats(&self) -> &EngineStats {
918        &self.stats
919    }
920
921    pub fn get_default_planes(&self) -> Arc<RwLock<Option<DefaultPlanes>>> {
922        self.default_planes.clone()
923    }
924
925    async fn clear_scene_post_hook(
926        &self,
927        batch_context: &EngineBatchContext,
928        id_generator: &mut IdGenerator,
929        source_range: SourceRange,
930        geometry_only: bool,
931    ) -> Result<(), KclError> {
932        // Remake the default planes, since they would have been removed after the scene was cleared.
933        let new_planes = self
934            .new_default_planes(batch_context, id_generator, source_range, geometry_only)
935            .await?;
936        *self.default_planes.write().await = Some(new_planes);
937
938        self.transport.start_new_session(source_range).await?;
939
940        Ok(())
941    }
942
943    async fn inner_send_modeling_cmd(
944        &self,
945        id: uuid::Uuid,
946        source_range: SourceRange,
947        cmd: WebSocketRequest,
948        id_to_source_range: HashMap<Uuid, SourceRange>,
949    ) -> Result<WebSocketResponse, KclError> {
950        let response = self
951            .transport
952            .inner_send_modeling_cmd(id, source_range, cmd, id_to_source_range)
953            .await?;
954
955        self.responses.add(id, response.clone()).await;
956        Ok(response)
957    }
958
959    pub async fn get_session_data(&self) -> Option<ModelingSessionData> {
960        self.session_data.read().await.clone()
961    }
962
963    pub async fn close(&self) {
964        let _ = self.transport.close().await;
965    }
966}
967
968/// State of the connection to the engine.
969#[derive(Debug, PartialEq)]
970pub enum SocketHealth {
971    Active,
972    Inactive,
973}