netget/utils/
save_load.rs

1//! Save/load utility for persisting server and client configurations
2//!
3//! This module handles serializing server/client state to action arrays
4//! and saving them to `.netget` files, as well as loading and parsing them.
5
6use anyhow::{Context, Result};
7use serde_json::{json, Value};
8use std::path::{Path, PathBuf};
9
10use crate::state::app_state::AppState;
11use crate::state::client::ClientId;
12use crate::state::server::ServerId;
13
14/// File extension for NetGet save files
15pub const NETGET_EXTENSION: &str = ".netget";
16pub const MAX_SESSION_BYTES: usize = 16 * 1024 * 1024;
17
18/// Normalize a filename to ensure it has the correct extension
19/// Strips any existing extension and adds .netget
20pub fn normalize_filename(name: &str) -> String {
21    let name = name.trim();
22    if name.is_empty() {
23        return NETGET_EXTENSION.to_string();
24    }
25    if name.ends_with(NETGET_EXTENSION) {
26        return name.to_string();
27    }
28    // Replace only the extension, retaining the caller's directory. Using file_stem()
29    // alone redirected `/some/directory/config.json` into the working directory.
30    Path::new(name)
31        .with_extension(&NETGET_EXTENSION[1..])
32        .to_string_lossy()
33        .into_owned()
34}
35
36/// Convert a server instance to an open_server action
37fn server_to_action(server: &crate::state::server::ServerInstance) -> Value {
38    let mut action = json!({
39        "type": "open_server",
40        "port": server.port,
41        "base_stack": server.protocol_name,
42        "instruction": server.instruction,
43    });
44
45    if let Some(address) = server.local_addr {
46        action["host"] = json!(address.ip().to_string());
47    }
48
49    // Add optional fields if present
50    if !server.memory.is_empty() {
51        action["initial_memory"] = json!(server.memory);
52    }
53
54    if let Some(ref params) = server.startup_params {
55        action["startup_params"] = params.clone();
56    }
57
58    if let Some(ref config) = server.event_handler_config {
59        action["event_handlers"] = json!(config.handlers);
60    }
61    if let Some(ref instruction) = server.feedback_instructions {
62        action["feedback_instructions"] = json!(instruction);
63    }
64
65    // Note: send_first is protocol-specific and not stored in ServerInstance
66    // It will use protocol defaults when recreated
67
68    action
69}
70
71/// Convert a client instance to an open_client action
72fn client_to_action(client: &crate::state::client::ClientInstance) -> Value {
73    let mut action = json!({
74        "type": "open_client",
75        "protocol": client.protocol_name,
76        "remote_addr": client.remote_addr,
77        "instruction": client.instruction,
78    });
79
80    // Add optional fields if present
81    if !client.memory.is_empty() {
82        action["initial_memory"] = json!(client.memory);
83    }
84
85    if let Some(ref params) = client.startup_params {
86        action["startup_params"] = params.clone();
87    }
88
89    if let Some(ref config) = client.event_handler_config {
90        action["event_handlers"] = json!(config.handlers);
91    }
92    if let Some(ref instruction) = client.feedback_instructions {
93        action["feedback_instructions"] = json!(instruction);
94    }
95
96    action
97}
98
99/// Save all servers and clients to a file
100pub async fn save_all(state: &AppState, filename: &str) -> Result<PathBuf> {
101    let mut servers = state.get_all_servers().await;
102    let mut clients = state.get_all_clients().await;
103    servers.sort_by_key(|item| item.id.as_u32());
104    clients.sort_by_key(|item| item.id.as_u32());
105    let tasks = state.get_all_tasks().await;
106    let pipes = state.list_pipes().await;
107    let easy = state.get_all_easy_instances().await;
108    let mut actions = Vec::new();
109    let mut resources = Vec::new();
110    for server in servers {
111        let mut action = server_to_action(&server);
112        attach_tasks(
113            &mut action,
114            &tasks,
115            &crate::state::task::TaskScope::Server(server.id),
116        )?;
117        resources.push(json!({"key": server.id.as_u32(), "action": action}));
118        actions.push(action);
119    }
120    for client in clients {
121        let mut action = client_to_action(&client);
122        attach_tasks(
123            &mut action,
124            &tasks,
125            &crate::state::task::TaskScope::Client(client.id),
126        )?;
127        resources.push(json!({"key": client.id.as_u32(), "action": action}));
128        actions.push(action);
129    }
130    let global_tasks: Vec<_> = tasks
131        .iter()
132        .filter(|task| matches!(task.scope, crate::state::task::TaskScope::Global))
133        .filter_map(task_definition)
134        .collect();
135    if !pipes.is_empty() || !easy.is_empty() || !global_tasks.is_empty() {
136        let wrappers: Vec<_> = easy
137            .into_iter()
138            .filter(|item| {
139                item.underlying_server_id.is_some() || item.underlying_client_id.is_some()
140            })
141            .map(|item| {
142                json!({
143                    "protocol": item.protocol_name, "underlying_protocol": item.underlying_protocol,
144                    "instruction": item.user_instruction,
145                    "server": item.underlying_server_id.map(|id| id.as_u32()),
146                    "client": item.underlying_client_id.map(|id| id.as_u32()),
147                })
148            })
149            .collect();
150        actions = vec![json!({"type": "restore_session", "session": {
151            "version": 2, "resources": resources, "pipes": pipes,
152            "easy": wrappers, "global_tasks": global_tasks,
153        }})];
154    }
155    save_actions(actions, filename).await
156}
157
158/// Save a specific server to a file
159pub async fn save_server(state: &AppState, server_id: ServerId, filename: &str) -> Result<PathBuf> {
160    let server = state
161        .get_server(server_id)
162        .await
163        .context("Server not found")?;
164
165    let mut action = server_to_action(&server);
166    attach_tasks(
167        &mut action,
168        &state.get_all_tasks().await,
169        &crate::state::task::TaskScope::Server(server_id),
170    )?;
171    let actions = vec![action];
172    save_actions(actions, filename).await
173}
174
175/// Save a specific client to a file
176pub async fn save_client(state: &AppState, client_id: ClientId, filename: &str) -> Result<PathBuf> {
177    let client = state
178        .get_client(client_id)
179        .await
180        .context("Client not found")?;
181
182    let mut action = client_to_action(&client);
183    attach_tasks(
184        &mut action,
185        &state.get_all_tasks().await,
186        &crate::state::task::TaskScope::Client(client_id),
187    )?;
188    let actions = vec![action];
189    save_actions(actions, filename).await
190}
191
192/// Save an array of actions to a file
193async fn save_actions(actions: Vec<Value>, filename: &str) -> Result<PathBuf> {
194    let filename = normalize_filename(filename);
195    let path = PathBuf::from(&filename);
196
197    // Wrap actions in the standard LLM format: {"actions": [...]}
198    let wrapped = json!({
199        "actions": actions
200    });
201
202    // Serialize to pretty JSON
203    let json =
204        serde_json::to_string_pretty(&wrapped).context("Failed to serialize actions to JSON")?;
205
206    if json.len() > MAX_SESSION_BYTES {
207        anyhow::bail!("session exceeds the 16 MiB save limit");
208    }
209    crate::utils::file_io::write_atomic_async(&path, json.into_bytes())
210        .await
211        .with_context(|| format!("Failed to save {}", path.display()))?;
212
213    Ok(path)
214}
215
216/// Load actions from a file
217pub async fn load_actions(filename: &str) -> Result<Vec<Value>> {
218    // Normalize filename (add .netget if missing)
219    let filename = normalize_filename(filename);
220
221    // Read file
222    let content = crate::utils::file_io::read_text_async(Path::new(&filename), MAX_SESSION_BYTES)
223        .await
224        .context(format!("Failed to read file: {}", filename))?;
225
226    // Parse JSON - expect {"actions": [...]} format
227    let parsed: Value = serde_json::from_str(&content)
228        .context(format!("Failed to parse JSON from file: {}", filename))?;
229
230    // Extract actions array
231    let actions = parsed
232        .get("actions")
233        .and_then(|v| v.as_array())
234        .ok_or_else(|| anyhow::anyhow!("File must contain {{\"actions\": [...]}} format"))?
235        .clone();
236
237    Ok(actions)
238}
239
240/// Check if a string is a valid actions JSON in {"actions": [...]} format
241pub fn is_actions_json(input: &str) -> bool {
242    // Try to parse as JSON object
243    if let Ok(value) = serde_json::from_str::<Value>(input) {
244        // Check for {"actions": [...]} format
245        if let Some(actions) = value.get("actions").and_then(|v| v.as_array()) {
246            // Check if all elements have a "type" field
247            return !actions.is_empty()
248                && actions.iter().all(|item| {
249                    item.as_object()
250                        .and_then(|obj| obj.get("type"))
251                        .and_then(|t| t.as_str())
252                        .is_some()
253                });
254        }
255    }
256    false
257}
258
259fn task_definition(
260    task: &crate::state::task::ScheduledTask,
261) -> Option<crate::llm::actions::common::ServerTaskDefinition> {
262    use crate::state::task::{TaskStatus, TaskType};
263    if !matches!(task.status, TaskStatus::Scheduled | TaskStatus::Executing) {
264        return None;
265    }
266    let (recurring, interval_secs, max_executions) = match task.task_type {
267        TaskType::OneShot { .. } => (false, None, None),
268        TaskType::Recurring {
269            interval_secs,
270            max_executions,
271            executions_count,
272        } => {
273            let remaining = max_executions.map(|max| max.saturating_sub(executions_count));
274            if remaining == Some(0) {
275                return None;
276            }
277            (true, Some(interval_secs), remaining)
278        }
279    };
280    let delay = task
281        .next_execution
282        .saturating_duration_since(crate::utils::clock::Instant::now());
283    let delay_secs = delay
284        .as_secs()
285        .saturating_add(u64::from(delay.subsec_nanos() != 0));
286    Some(crate::llm::actions::common::ServerTaskDefinition {
287        task_id: task.name.clone(),
288        recurring,
289        delay_secs: Some(delay_secs),
290        interval_secs,
291        max_executions,
292        instruction: task.instruction.clone(),
293        context: task.context.clone(),
294    })
295}
296
297fn attach_tasks(
298    action: &mut Value,
299    tasks: &[crate::state::task::ScheduledTask],
300    scope: &crate::state::task::TaskScope,
301) -> Result<()> {
302    use crate::state::task::TaskScope;
303    let definitions: Vec<_> = tasks
304        .iter()
305        .filter(|task| match (&task.scope, scope) {
306            (TaskScope::Server(a), TaskScope::Server(b)) => a == b,
307            (TaskScope::Client(a), TaskScope::Client(b)) => a == b,
308            _ => false,
309        })
310        .filter_map(task_definition)
311        .collect();
312    if !definitions.is_empty() {
313        action["scheduled_tasks"] = serde_json::to_value(definitions)?;
314    }
315    Ok(())
316}
317
318#[derive(serde::Deserialize)]
319#[serde(deny_unknown_fields)]
320struct SavedResource {
321    key: u32,
322    action: Value,
323}
324#[derive(serde::Deserialize)]
325#[serde(deny_unknown_fields)]
326struct SavedEasy {
327    protocol: String,
328    underlying_protocol: String,
329    instruction: Option<String>,
330    server: Option<u32>,
331    client: Option<u32>,
332}
333#[derive(serde::Deserialize)]
334#[serde(deny_unknown_fields)]
335struct SavedSession {
336    version: u32,
337    resources: Vec<SavedResource>,
338    #[serde(default)]
339    pipes: Vec<crate::pipe::PipeSpec>,
340    #[serde(default)]
341    easy: Vec<SavedEasy>,
342    #[serde(default)]
343    global_tasks: Vec<crate::llm::actions::common::ServerTaskDefinition>,
344}
345
346/// Restore configuration relationships using newly allocated IDs. Connections,
347/// transient request queues and in-flight responses are deliberately not resumed.
348/// On startup failure, remove resources created by this restore; existing state
349/// is never used as a rollback target.
350pub async fn restore_session(
351    state: &AppState,
352    llm: &crate::llm::OllamaClient,
353    value: &Value,
354) -> Result<()> {
355    use crate::cli::management::{ClientForm, ServerForm};
356    use crate::state::task::{prepare_tasks, TaskScope};
357    use crate::state::{EasyId, EasyInstance, EasyStatus};
358    use std::collections::{HashMap, HashSet};
359    if !crate::utils::json_budget::within_budget(value, MAX_SESSION_BYTES * 4, 100_000, 64) {
360        anyhow::bail!("session exceeds the structural budget");
361    }
362    let session: SavedSession =
363        serde_json::from_value(value.clone()).context("invalid session format")?;
364    if session.version != 2 {
365        anyhow::bail!("unsupported session version {}", session.version);
366    }
367    enum Form {
368        Server(u32, ServerForm),
369        Client(u32, ClientForm),
370    }
371    let mut forms = Vec::new();
372    let mut server_keys = HashSet::new();
373    let mut client_keys = HashSet::new();
374    let mut all_keys = HashSet::new();
375    let mut protocols = HashMap::new();
376    for resource in session.resources {
377        if !all_keys.insert(resource.key) {
378            anyhow::bail!("duplicate saved resource key {}", resource.key);
379        }
380        let mut action = resource.action;
381        let object = action
382            .as_object_mut()
383            .context("saved resource action must be an object")?;
384        if !object.contains_key("protocol") {
385            object.insert(
386                "protocol".into(),
387                object.get("base_stack").cloned().unwrap_or(Value::Null),
388            );
389        }
390        match action["type"].as_str() {
391            Some("open_server") => {
392                let form: ServerForm = serde_json::from_value(action)?;
393                prepare_tasks(form.scheduled_tasks.as_deref(), TaskScope::Global)?;
394                protocols.insert(resource.key, form.protocol.clone());
395                server_keys.insert(resource.key);
396                forms.push(Form::Server(resource.key, form));
397            }
398            Some("open_client") => {
399                let form: ClientForm = serde_json::from_value(action)?;
400                prepare_tasks(form.scheduled_tasks.as_deref(), TaskScope::Global)?;
401                protocols.insert(resource.key, form.protocol.clone());
402                client_keys.insert(resource.key);
403                forms.push(Form::Client(resource.key, form));
404            }
405            _ => anyhow::bail!("session resources must be open_server or open_client"),
406        }
407    }
408    let global_tasks = prepare_tasks(Some(&session.global_tasks), TaskScope::Global)?;
409    let mut checked_pipes = Vec::new();
410    for pipe in &session.pipes {
411        if !server_keys.contains(&pipe.from) || !server_keys.contains(&pipe.to) {
412            anyhow::bail!("saved pipe references a missing server");
413        }
414        if crate::pipe::would_create_cycle(&checked_pipes, pipe.from, pipe.to) {
415            anyhow::bail!("saved pipe graph contains a cycle");
416        }
417        checked_pipes.push(pipe.clone());
418    }
419    let mut wrapped_keys = HashSet::new();
420    for easy in &session.easy {
421        let known = crate::protocol::EASY_REGISTRY
422            .get_by_name(&easy.protocol)
423            .with_context(|| format!("unknown saved easy protocol {}", easy.protocol))?;
424        if known.underlying_protocol() != easy.underlying_protocol {
425            anyhow::bail!("saved easy underlying protocol does not match its registry");
426        }
427        if easy.server.is_some() == easy.client.is_some()
428            || easy.server.is_some_and(|key| !server_keys.contains(&key))
429            || easy.client.is_some_and(|key| !client_keys.contains(&key))
430        {
431            anyhow::bail!("saved easy wrapper must reference exactly one existing resource");
432        }
433        let key = easy
434            .server
435            .or(easy.client)
436            .expect("validated resource reference");
437        if !protocols[&key].eq_ignore_ascii_case(&easy.underlying_protocol) {
438            anyhow::bail!("saved easy wrapper protocol differs from its resource");
439        }
440        if !wrapped_keys.insert(key) {
441            anyhow::bail!("multiple easy wrappers reference the same resource");
442        }
443    }
444    let mut servers = HashMap::new();
445    let mut clients = HashMap::new();
446    let mut ownership = RestoreOwnership::new(state.clone());
447    let (status_tx, mut status_rx) = tokio::sync::mpsc::unbounded_channel();
448    ownership.printer = Some(tokio::spawn(async move {
449        while let Some(line) = status_rx.recv().await {
450            tracing::info!("{line}");
451        }
452    }));
453    let result: Result<()> = async {
454        for form in forms {
455            match form {
456                Form::Server(key, form) => {
457                    let id = form.create(state, status_tx.clone()).await?;
458                    ownership.servers.push(id);
459                    servers.insert(key, id);
460                }
461                Form::Client(key, form) => {
462                    let id = form.create(state, llm.clone(), status_tx.clone()).await?;
463                    ownership.clients.push(id);
464                    clients.insert(key, id);
465                }
466            }
467        }
468        for pipe in session.pipes {
469            state
470                .add_pipe(
471                    servers[&pipe.from],
472                    pipe.on,
473                    servers[&pipe.to],
474                    pipe.as_action,
475                    pipe.map,
476                )
477                .await?;
478        }
479        for easy in session.easy {
480            let instance = EasyInstance::new(
481                EasyId::new(0),
482                easy.protocol,
483                easy.underlying_protocol,
484                easy.instruction,
485            );
486            let id = state.add_easy_instance(instance).await;
487            ownership.easy.push(id);
488            if let Some(key) = easy.server {
489                state.link_server_to_easy(servers[&key], id).await;
490            }
491            if let Some(key) = easy.client {
492                state.link_client_to_easy(clients[&key], id).await;
493            }
494            state.update_easy_status(id, EasyStatus::Running).await;
495        }
496        for task in global_tasks {
497            ownership.tasks.push(state.add_task(task).await);
498        }
499        Ok(())
500    }
501    .await;
502    drop(status_tx);
503    if let Some(printer) = ownership.printer.take() {
504        printer.abort();
505    }
506    if result.is_err() {
507        ownership.rollback().await;
508    } else {
509        ownership.committed = true;
510    }
511    result
512}
513
514struct RestoreOwnership {
515    state: AppState,
516    servers: Vec<ServerId>,
517    clients: Vec<ClientId>,
518    easy: Vec<crate::state::EasyId>,
519    tasks: Vec<crate::state::task::TaskId>,
520    printer: Option<tokio::task::JoinHandle<()>>,
521    committed: bool,
522}
523impl RestoreOwnership {
524    fn new(state: AppState) -> Self {
525        Self {
526            state,
527            servers: Vec::new(),
528            clients: Vec::new(),
529            easy: Vec::new(),
530            tasks: Vec::new(),
531            printer: None,
532            committed: false,
533        }
534    }
535    async fn rollback(&mut self) {
536        while let Some(id) = self.tasks.last().copied() {
537            self.state.remove_task(id).await;
538            self.tasks.pop();
539        }
540        while let Some(id) = self.easy.last().copied() {
541            self.state.remove_easy_instance(id).await;
542            self.easy.pop();
543        }
544        while let Some(id) = self.clients.last().copied() {
545            self.state.remove_client(id).await;
546            self.clients.pop();
547        }
548        while let Some(id) = self.servers.last().copied() {
549            self.state.remove_server(id).await;
550            self.servers.pop();
551        }
552        self.committed = true;
553    }
554}
555impl Drop for RestoreOwnership {
556    fn drop(&mut self) {
557        if let Some(printer) = self.printer.take() {
558            printer.abort();
559        }
560        if self.committed {
561            return;
562        }
563        let mut cleanup = Self::new(self.state.clone());
564        cleanup.servers = std::mem::take(&mut self.servers);
565        cleanup.clients = std::mem::take(&mut self.clients);
566        cleanup.easy = std::mem::take(&mut self.easy);
567        cleanup.tasks = std::mem::take(&mut self.tasks);
568        cleanup.committed = true; // Runtime shutdown must not recursively spawn cleanup.
569        if let Ok(runtime) = tokio::runtime::Handle::try_current() {
570            runtime.spawn(async move {
571                cleanup.rollback().await;
572            });
573        }
574    }
575}
netget/utils/save_load.rs

netget/utils/
save_load.rs
