From ca53e5005333874a64299c4417fbb52722f514f7 Mon Sep 17 00:00:00 2001
From: Dreaded_X <tim@huizinga.dev>
Date: Thu, 1 Oct 2026 02:54:17 +0200
Subject: [PATCH] feat: Retain context

---
 src/lib.rs         |  1 +
 src/main.rs        |  7 +++++++
 src/talosconfig.rs | 34 ++++++++++++++++++++++++++++++++++
 3 files changed, 42 insertions(+)
 create mode 100644 src/talosconfig.rs

diff --git a/src/lib.rs b/src/lib.rs
index ce4e59a..e8987b7 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -10,6 +10,7 @@ pub mod node;
 pub mod patch;
 pub mod schematic;
 pub mod secret;
+pub mod talosconfig;
 pub mod tftp;
 
 pub(crate) static REPO_PATH: OnceLock<PathBuf> = OnceLock::new();
diff --git a/src/main.rs b/src/main.rs
index 6cd6231..5e07e00 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -6,6 +6,7 @@ use clap::{CommandFactory, Parser};
 use clap_complete::{Shell, generate as generate_complete};
 use crete::cluster::Cluster;
 use crete::environment::PathEnvironment;
+use crete::talosconfig::{get_current_context, set_context};
 use crete::{get_configs_path, get_repo_path, set_repo_path};
 use futures::StreamExt;
 use futures::future::join_all;
@@ -53,6 +54,8 @@ async fn generate(opts: &GlobalOpts) -> Result<(), Error> {
             .unwrap();
     }
 
+    let context = get_current_context().await;
+
     // Remove existing config files and create output directory
     let path = get_configs_path();
     tokio::fs::remove_dir_all(&path).await.ok();
@@ -80,6 +83,10 @@ async fn generate(opts: &GlobalOpts) -> Result<(), Error> {
         cluster.talosconfig_merge().await;
     }
 
+    if let Some(context) = context {
+        set_context(context).await;
+    }
+
     Ok(())
 }
 
diff --git a/src/talosconfig.rs b/src/talosconfig.rs
new file mode 100644
index 0000000..60d630b
--- /dev/null
+++ b/src/talosconfig.rs
@@ -0,0 +1,34 @@
+use serde::Deserialize;
+use tokio::process::Command;
+
+use crate::{get_configs_path, run_command};
+
+#[derive(Debug, Deserialize)]
+struct TalosConfig {
+    pub context: String,
+}
+
+pub async fn get_current_context() -> Option<String> {
+    let path = get_configs_path().join("talosconfig");
+
+    tokio::fs::read(&path)
+        .await
+        .ok()
+        .and_then(|config| serde_yaml::from_slice(&config).ok())
+        .map(|config: TalosConfig| config.context)
+}
+
+pub async fn set_context(context: String) {
+    let path = get_configs_path().join("talosconfig");
+
+    let mut command = Command::new("talosctl");
+    command.args([
+        "config",
+        "--talosconfig",
+        path.to_str().expect("Path should be valid utf-8"),
+        "context",
+        &context,
+    ]);
+
+    run_command(command).await;
+}