kcl_lib/engine/
conn_mock.rs

1//! Functions for setting up our WebSocket and WebRTC connections for communications with the
2//! engine.
3
4use std::{collections::HashMap, sync::Arc};
5
6use anyhow::Result;
7use indexmap::IndexMap;
8use kcmc::{
9    ok_response::OkModelingCmdResponse,
10    websocket::{
11        BatchResponse, ModelingBatch, OkWebSocketResponseData, SuccessWebSocketResponse, WebSocketRequest,
12        WebSocketResponse,
13    },
14};
15use kittycad_modeling_cmds::{self as kcmc, websocket::ModelingCmdReq, ImportFiles, ModelingCmd};
16use tokio::sync::RwLock;
17use uuid::Uuid;
18
19#[cfg(feature = "artifact-graph")]
20use crate::execution::ArtifactCommand;
21use crate::{
22    engine::{AsyncTasks, EngineStats},
23    errors::KclError,
24    exec::DefaultPlanes,
25    execution::IdGenerator,
26    SourceRange,
27};
28
29#[derive(Debug, Clone)]
30pub struct EngineConnection {
31    batch: Arc<RwLock<Vec<(WebSocketRequest, SourceRange)>>>,
32    batch_end: Arc<RwLock<IndexMap<uuid::Uuid, (WebSocketRequest, SourceRange)>>>,
33    #[cfg(feature = "artifact-graph")]
34    artifact_commands: Arc<RwLock<Vec<ArtifactCommand>>>,
35    ids_of_async_commands: Arc<RwLock<IndexMap<Uuid, SourceRange>>>,
36    responses: Arc<RwLock<IndexMap<Uuid, WebSocketResponse>>>,
37    /// The default planes for the scene.
38    default_planes: Arc<RwLock<Option<DefaultPlanes>>>,
39    stats: EngineStats,
40    async_tasks: AsyncTasks,
41}
42
43impl EngineConnection {
44    pub async fn new() -> Result<EngineConnection> {
45        Ok(EngineConnection {
46            batch: Arc::new(RwLock::new(Vec::new())),
47            batch_end: Arc::new(RwLock::new(IndexMap::new())),
48            #[cfg(feature = "artifact-graph")]
49            artifact_commands: Arc::new(RwLock::new(Vec::new())),
50            ids_of_async_commands: Arc::new(RwLock::new(IndexMap::new())),
51            responses: Arc::new(RwLock::new(IndexMap::new())),
52            default_planes: Default::default(),
53            stats: Default::default(),
54            async_tasks: AsyncTasks::new(),
55        })
56    }
57}
58
59#[async_trait::async_trait]
60impl crate::engine::EngineManager for EngineConnection {
61    fn batch(&self) -> Arc<RwLock<Vec<(WebSocketRequest, SourceRange)>>> {
62        self.batch.clone()
63    }
64
65    fn batch_end(&self) -> Arc<RwLock<IndexMap<uuid::Uuid, (WebSocketRequest, SourceRange)>>> {
66        self.batch_end.clone()
67    }
68
69    fn responses(&self) -> Arc<RwLock<IndexMap<Uuid, WebSocketResponse>>> {
70        self.responses.clone()
71    }
72
73    fn stats(&self) -> &EngineStats {
74        &self.stats
75    }
76
77    #[cfg(feature = "artifact-graph")]
78    fn artifact_commands(&self) -> Arc<RwLock<Vec<ArtifactCommand>>> {
79        self.artifact_commands.clone()
80    }
81
82    fn ids_of_async_commands(&self) -> Arc<RwLock<IndexMap<Uuid, SourceRange>>> {
83        self.ids_of_async_commands.clone()
84    }
85
86    fn async_tasks(&self) -> AsyncTasks {
87        self.async_tasks.clone()
88    }
89
90    fn get_default_planes(&self) -> Arc<RwLock<Option<DefaultPlanes>>> {
91        self.default_planes.clone()
92    }
93
94    async fn clear_scene_post_hook(
95        &self,
96        _id_generator: &mut IdGenerator,
97        _source_range: SourceRange,
98    ) -> Result<(), KclError> {
99        Ok(())
100    }
101
102    async fn get_debug(&self) -> Option<OkWebSocketResponseData> {
103        None
104    }
105
106    async fn fetch_debug(&self) -> Result<(), KclError> {
107        unimplemented!();
108    }
109
110    async fn inner_fire_modeling_cmd(
111        &self,
112        id: uuid::Uuid,
113        source_range: SourceRange,
114        cmd: WebSocketRequest,
115        id_to_source_range: HashMap<Uuid, SourceRange>,
116    ) -> Result<(), KclError> {
117        // Pop off the id we care about.
118        self.ids_of_async_commands.write().await.swap_remove(&id);
119
120        // Add the response to our responses.
121        let response = self
122            .inner_send_modeling_cmd(id, source_range, cmd, id_to_source_range)
123            .await?;
124        self.responses().write().await.insert(id, response);
125
126        Ok(())
127    }
128
129    async fn inner_send_modeling_cmd(
130        &self,
131        id: uuid::Uuid,
132        _source_range: SourceRange,
133        cmd: WebSocketRequest,
134        _id_to_source_range: HashMap<Uuid, SourceRange>,
135    ) -> Result<WebSocketResponse, KclError> {
136        match cmd {
137            WebSocketRequest::ModelingCmdBatchReq(ModelingBatch {
138                ref requests,
139                batch_id: _,
140                responses: _,
141            }) => {
142                // Create the empty responses.
143                let mut responses = HashMap::with_capacity(requests.len());
144                for request in requests {
145                    responses.insert(
146                        request.cmd_id,
147                        BatchResponse::Success {
148                            response: OkModelingCmdResponse::Empty {},
149                        },
150                    );
151                }
152                Ok(WebSocketResponse::Success(SuccessWebSocketResponse {
153                    request_id: Some(id),
154                    resp: OkWebSocketResponseData::ModelingBatch { responses },
155                    success: true,
156                }))
157            }
158            WebSocketRequest::ModelingCmdReq(ModelingCmdReq {
159                cmd: ModelingCmd::ImportFiles(ImportFiles { .. }),
160                cmd_id,
161            }) => Ok(WebSocketResponse::Success(SuccessWebSocketResponse {
162                request_id: Some(id),
163                resp: OkWebSocketResponseData::Modeling {
164                    modeling_response: OkModelingCmdResponse::ImportFiles(
165                        kittycad_modeling_cmds::output::ImportFiles {
166                            object_id: cmd_id.into(),
167                        },
168                    ),
169                },
170                success: true,
171            })),
172            WebSocketRequest::ModelingCmdReq(_) => Ok(WebSocketResponse::Success(SuccessWebSocketResponse {
173                request_id: Some(id),
174                resp: OkWebSocketResponseData::Modeling {
175                    modeling_response: OkModelingCmdResponse::Empty {},
176                },
177                success: true,
178            })),
179            _ => Ok(WebSocketResponse::Success(SuccessWebSocketResponse {
180                request_id: Some(id),
181                resp: OkWebSocketResponseData::Modeling {
182                    modeling_response: OkModelingCmdResponse::Empty {},
183                },
184                success: true,
185            })),
186        }
187    }
188
189    async fn close(&self) {}
190}
kcl_lib/engine/conn_mock.rs

kcl_lib/engine/
conn_mock.rs
