1use base64::Engine;
2use chrono::Utc;
3use runmat_builtins::Value;
4use serde::{Deserialize, Serialize};
5
6use crate::builtins::io::mat::load::decode_workspace_from_mat_bytes;
7use crate::builtins::io::mat::save::encode_workspace_to_mat_bytes;
8use crate::replay::limits::ReplayLimits;
9use crate::runtime_error::{replay_error, replay_error_with_source, ReplayErrorKind};
10use crate::{BuiltinResult, RuntimeError};
11
12const WORKSPACE_SCHEMA_VERSION: u32 = 1;
13const WORKSPACE_KIND: &str = "workspace-state";
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16pub enum WorkspaceReplayMode {
17 Auto,
18 Force,
19 Off,
20}
21
22impl WorkspaceReplayMode {
23 pub fn as_str(self) -> &'static str {
24 match self {
25 Self::Auto => "auto",
26 Self::Force => "force",
27 Self::Off => "off",
28 }
29 }
30}
31
32#[derive(Debug, Clone, Serialize, Deserialize)]
33#[serde(rename_all = "camelCase")]
34struct WorkspaceReplayPayload {
35 schema_version: u32,
36 kind: String,
37 created_at: String,
38 mode: String,
39 mat_base64: String,
40}
41
42pub async fn encode_workspace_payload(
43 entries: &[(String, Value)],
44 mode: &str,
45) -> BuiltinResult<Vec<u8>> {
46 encode_workspace_payload_with_limits(entries, mode, ReplayLimits::default()).await
47}
48
49pub async fn export_workspace_state(
50 entries: &[(String, Value)],
51 mode: WorkspaceReplayMode,
52) -> BuiltinResult<Option<Vec<u8>>> {
53 if matches!(mode, WorkspaceReplayMode::Off) {
54 return Ok(None);
55 }
56 encode_workspace_payload(entries, mode.as_str())
57 .await
58 .map(Some)
59}
60
61pub async fn encode_workspace_payload_with_limits(
62 entries: &[(String, Value)],
63 mode: &str,
64 limits: ReplayLimits,
65) -> BuiltinResult<Vec<u8>> {
66 validate_workspace_mode(mode)?;
67 if entries.len() > limits.max_workspace_variables {
68 return Err(replay_error(
69 ReplayErrorKind::ImportRejected,
70 format!(
71 "workspace export includes {} variables, exceeding limit {}",
72 entries.len(),
73 limits.max_workspace_variables
74 ),
75 ));
76 }
77
78 let mat_bytes = encode_workspace_to_mat_bytes(entries).await?;
79 if mat_bytes.len() > limits.max_workspace_mat_bytes {
80 return Err(replay_error(
81 ReplayErrorKind::PayloadTooLarge,
82 format!(
83 "workspace MAT payload is {} bytes, exceeding limit {}",
84 mat_bytes.len(),
85 limits.max_workspace_mat_bytes
86 ),
87 ));
88 }
89
90 let payload = WorkspaceReplayPayload {
91 schema_version: WORKSPACE_SCHEMA_VERSION,
92 kind: WORKSPACE_KIND.to_string(),
93 created_at: Utc::now().to_rfc3339(),
94 mode: mode.to_string(),
95 mat_base64: base64::engine::general_purpose::STANDARD.encode(mat_bytes),
96 };
97
98 let encoded = serde_json::to_vec(&payload).map_err(|err| {
99 replay_error_with_source(
100 ReplayErrorKind::DecodeFailed,
101 "failed to encode workspace replay payload",
102 err,
103 )
104 })?;
105
106 if encoded.len() > limits.max_workspace_payload_bytes {
107 return Err(replay_error(
108 ReplayErrorKind::PayloadTooLarge,
109 format!(
110 "workspace replay payload is {} bytes, exceeding limit {}",
111 encoded.len(),
112 limits.max_workspace_payload_bytes
113 ),
114 ));
115 }
116
117 Ok(encoded)
118}
119
120pub fn decode_workspace_payload(bytes: &[u8]) -> BuiltinResult<Vec<(String, Value)>> {
121 decode_workspace_payload_with_limits(bytes, ReplayLimits::default())
122}
123
124pub fn import_workspace_state(bytes: &[u8]) -> BuiltinResult<Vec<(String, Value)>> {
125 decode_workspace_payload(bytes)
126}
127
128pub fn decode_workspace_payload_with_limits(
129 bytes: &[u8],
130 limits: ReplayLimits,
131) -> BuiltinResult<Vec<(String, Value)>> {
132 if bytes.len() > limits.max_workspace_payload_bytes {
133 return Err(replay_error(
134 ReplayErrorKind::PayloadTooLarge,
135 format!(
136 "workspace replay payload is {} bytes, exceeding limit {}",
137 bytes.len(),
138 limits.max_workspace_payload_bytes
139 ),
140 ));
141 }
142
143 let payload: WorkspaceReplayPayload = serde_json::from_slice(bytes).map_err(|err| {
144 replay_error_with_source(
145 ReplayErrorKind::DecodeFailed,
146 "failed to decode workspace replay payload",
147 err,
148 )
149 })?;
150
151 if payload.schema_version != WORKSPACE_SCHEMA_VERSION {
152 return Err(replay_error(
153 ReplayErrorKind::UnsupportedSchema,
154 format!(
155 "unsupported workspace replay schema version {}",
156 payload.schema_version
157 ),
158 ));
159 }
160 if payload.kind != WORKSPACE_KIND {
161 return Err(replay_error(
162 ReplayErrorKind::ImportRejected,
163 format!("unexpected replay payload kind '{}'", payload.kind),
164 ));
165 }
166 validate_workspace_mode(&payload.mode)?;
167
168 let mat_bytes = base64::engine::general_purpose::STANDARD
169 .decode(payload.mat_base64.as_bytes())
170 .map_err(|err| {
171 replay_error_with_source(
172 ReplayErrorKind::DecodeFailed,
173 "failed to decode workspace replay MAT bytes",
174 err,
175 )
176 })?;
177
178 if mat_bytes.len() > limits.max_workspace_mat_bytes {
179 return Err(replay_error(
180 ReplayErrorKind::PayloadTooLarge,
181 format!(
182 "workspace MAT payload is {} bytes, exceeding limit {}",
183 mat_bytes.len(),
184 limits.max_workspace_mat_bytes
185 ),
186 ));
187 }
188
189 let entries = decode_workspace_from_mat_bytes(&mat_bytes).map_err(|err| {
190 replay_error_with_source(
191 ReplayErrorKind::DecodeFailed,
192 "failed to decode workspace MAT payload",
193 err,
194 )
195 })?;
196 if entries.len() > limits.max_workspace_variables {
197 return Err(replay_error(
198 ReplayErrorKind::ImportRejected,
199 format!(
200 "workspace payload includes {} variables, exceeding limit {}",
201 entries.len(),
202 limits.max_workspace_variables
203 ),
204 ));
205 }
206 Ok(entries)
207}
208
209fn validate_workspace_mode(mode: &str) -> Result<(), RuntimeError> {
210 if matches!(mode, "auto" | "force") {
211 Ok(())
212 } else {
213 Err(replay_error(
214 ReplayErrorKind::ImportRejected,
215 format!("workspace replay mode '{mode}' is not supported"),
216 ))
217 }
218}
219
220#[cfg(test)]
221mod tests {
222 use futures::executor::block_on;
223
224 use super::*;
225
226 #[test]
227 fn workspace_schema_mismatch_rejects() {
228 let payload = serde_json::json!({
229 "schemaVersion": 99,
230 "kind": WORKSPACE_KIND,
231 "createdAt": "2026-01-01T00:00:00Z",
232 "mode": "auto",
233 "matBase64": ""
234 });
235 let bytes = serde_json::to_vec(&payload).expect("serialize payload");
236 let err = decode_workspace_payload_with_limits(&bytes, ReplayLimits::default())
237 .expect_err("expected schema rejection");
238 assert_eq!(
239 err.identifier(),
240 Some(ReplayErrorKind::UnsupportedSchema.identifier())
241 );
242 }
243
244 #[test]
245 fn workspace_payload_too_large_rejects() {
246 let bytes = vec![0u8; ReplayLimits::default().max_workspace_payload_bytes + 1];
247 let err = decode_workspace_payload_with_limits(&bytes, ReplayLimits::default())
248 .expect_err("expected payload rejection");
249 assert_eq!(
250 err.identifier(),
251 Some(ReplayErrorKind::PayloadTooLarge.identifier())
252 );
253 }
254
255 #[test]
256 fn workspace_variable_count_limit_rejects() {
257 let entries = vec![
258 ("a".to_string(), Value::Num(1.0)),
259 ("b".to_string(), Value::Num(2.0)),
260 ];
261 let bytes = block_on(encode_workspace_payload_with_limits(
262 &entries,
263 "auto",
264 ReplayLimits {
265 max_workspace_variables: 4,
266 ..ReplayLimits::default()
267 },
268 ))
269 .expect("encode workspace payload");
270
271 let err = decode_workspace_payload_with_limits(
272 &bytes,
273 ReplayLimits {
274 max_workspace_variables: 1,
275 ..ReplayLimits::default()
276 },
277 )
278 .expect_err("expected variable limit rejection");
279 assert_eq!(
280 err.identifier(),
281 Some(ReplayErrorKind::ImportRejected.identifier())
282 );
283 }
284}