Skip to main content

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