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 use runmat_builtins::{IntegerStorage, Tensor};
224
225 use super::*;
226
227 #[test]
228 fn workspace_schema_mismatch_rejects() {
229 let payload = serde_json::json!({
230 "schemaVersion": 99,
231 "kind": WORKSPACE_KIND,
232 "createdAt": "2026-01-01T00:00:00Z",
233 "mode": "auto",
234 "matBase64": ""
235 });
236 let bytes = serde_json::to_vec(&payload).expect("serialize payload");
237 let err = decode_workspace_payload_with_limits(&bytes, ReplayLimits::default())
238 .expect_err("expected schema rejection");
239 assert_eq!(
240 err.identifier(),
241 Some(ReplayErrorKind::UnsupportedSchema.identifier())
242 );
243 }
244
245 #[test]
246 fn workspace_payload_too_large_rejects() {
247 let bytes = vec![0u8; ReplayLimits::default().max_workspace_payload_bytes + 1];
248 let err = decode_workspace_payload_with_limits(&bytes, ReplayLimits::default())
249 .expect_err("expected payload rejection");
250 assert_eq!(
251 err.identifier(),
252 Some(ReplayErrorKind::PayloadTooLarge.identifier())
253 );
254 }
255
256 #[test]
257 fn workspace_variable_count_limit_rejects() {
258 let entries = vec![
259 ("a".to_string(), Value::Num(1.0)),
260 ("b".to_string(), Value::Num(2.0)),
261 ];
262 let bytes = block_on(encode_workspace_payload_with_limits(
263 &entries,
264 "auto",
265 ReplayLimits {
266 max_workspace_variables: 4,
267 ..ReplayLimits::default()
268 },
269 ))
270 .expect("encode workspace payload");
271
272 let err = decode_workspace_payload_with_limits(
273 &bytes,
274 ReplayLimits {
275 max_workspace_variables: 1,
276 ..ReplayLimits::default()
277 },
278 )
279 .expect_err("expected variable limit rejection");
280 assert_eq!(
281 err.identifier(),
282 Some(ReplayErrorKind::ImportRejected.identifier())
283 );
284 }
285
286 #[test]
287 fn workspace_replay_preserves_exact_integer_storage() {
288 let input =
289 Tensor::new_integer(IntegerStorage::U64(vec![1_u64 << 63, u64::MAX]), vec![1, 2])
290 .expect("integer tensor");
291 let entries = vec![("samples".to_string(), Value::Tensor(input))];
292
293 let payload =
294 block_on(encode_workspace_payload(&entries, "auto")).expect("encode workspace payload");
295 let decoded = decode_workspace_payload(&payload).expect("decode workspace payload");
296
297 assert!(matches!(decoded.as_slice(), [(name, Value::Tensor(tensor))]
298 if name == "samples"
299 && tensor.integer_storage()
300 == Some(&IntegerStorage::U64(vec![1_u64 << 63, u64::MAX]))));
301 }
302}