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