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