From 76b4369aa670ff4281629daaaa2ec76d22159e59 Mon Sep 17 00:00:00 2001
From: Justin Bradford <jabradford@gmail.com>
Date: Wed, 19 Nov 2025 22:49:37 -0800
Subject: [PATCH] feat: store machine id on connection entries in Uncloud
 config so it can be removed with machine (#182)

* fix: Store machine id on connection entries so it can be removed when machine is

* Add "GetContextOverrideOrCurrent" helper method to cli and use in `uc machine rm` to clean up connections
---
 cmd/uncloud/machine/rm.go              | 17 +++++++++++++++--
 internal/cli/cli.go                    | 15 +++++++++++----
 internal/cli/config/connection.go      |  1 +
 internal/cli/config/connection_test.go |  9 +++++++++
 4 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/cmd/uncloud/machine/rm.go b/cmd/uncloud/machine/rm.go
index 0516ceb7..a50a9eb7 100644
--- a/cmd/uncloud/machine/rm.go
+++ b/cmd/uncloud/machine/rm.go
@@ -152,8 +152,21 @@ func remove(ctx context.Context, uncli *cli.CLI, nameOrID string, opts removeOpt
 		}
 	}
 
-	// TODO: remove the connection to the machine from the uncloud config if it exists. We need a way to associate
-	//  the machine with its connection in the config, e.g. by storing the machine name in the connection metadata.
+	// Remove the connection to the machine from the uncloud config if it exists.
+	if uncli.Config != nil {
+		contextName := uncli.GetContextOverrideOrCurrent()
+		if context, ok := uncli.Config.Contexts[contextName]; ok {
+			for i, c := range context.Connections {
+				if c.MachineID == m.Id {
+					context.Connections = slices.Delete(context.Connections, i, i+1)
+					break
+				}
+			}
+			if err := uncli.Config.Save(); err != nil {
+				return fmt.Errorf("save config: %w", err)
+			}
+		}
+	}
 
 	// TODO: If Caddy was running on this machine and a cluster domain is reserved,
 	//  let the user know that the DNS records should be updated.
diff --git a/internal/cli/cli.go b/internal/cli/cli.go
index 1bc93fba..446197ae 100644
--- a/internal/cli/cli.go
+++ b/internal/cli/cli.go
@@ -71,6 +71,14 @@ func (cli *CLI) SetCurrentContext(name string) error {
 	return cli.Config.Save()
 }
 
+func (cli *CLI) GetContextOverrideOrCurrent() string {
+	contextName := cli.contextOverride
+	if contextName == "" {
+		contextName = cli.Config.CurrentContext
+	}
+	return contextName
+}
+
 // ConnectCluster connects to a cluster using the context override or the current context if not specified.
 // If the CLI was initialised with a machine connection, the config is ignored and the connection is used instead.
 func (cli *CLI) ConnectCluster(ctx context.Context) (*client.Client, error) {
@@ -235,6 +243,7 @@ func (cli *CLI) initRemoteMachine(ctx context.Context, opts InitClusterOptions)
 	// Save the machine's SSH connection details in the context config.
 	connCfg := config.MachineConnection{
 		SSHKeyFile: opts.RemoteMachine.KeyPath,
+		MachineID:  resp.Machine.Id,
 	}
 	if opts.RemoteMachine.UseSSHCLI {
 		connCfg.SSHCLI = config.NewSSHDestination(opts.RemoteMachine.User, opts.RemoteMachine.Host, opts.RemoteMachine.Port)
@@ -287,10 +296,7 @@ type AddMachineOptions struct {
 // cluster. The machine client is connected to the new machine and can be used to interact with it.
 // Both client should be closed after use by the caller.
 func (cli *CLI) AddMachine(ctx context.Context, opts AddMachineOptions) (*client.Client, *client.Client, error) {
-	contextName := cli.contextOverride
-	if contextName == "" {
-		contextName = cli.Config.CurrentContext
-	}
+	contextName := cli.GetContextOverrideOrCurrent()
 	c, err := cli.ConnectCluster(ctx)
 	if err != nil {
 		return nil, nil, fmt.Errorf("connect to cluster (context '%s'): %w", contextName, err)
@@ -406,6 +412,7 @@ func (cli *CLI) AddMachine(ctx context.Context, opts AddMachineOptions) (*client
 	// Save the machine's SSH connection details in the context config.
 	connCfg := config.MachineConnection{
 		SSHKeyFile: opts.RemoteMachine.KeyPath,
+		MachineID:  addResp.Machine.Id,
 	}
 	if opts.RemoteMachine.UseSSHCLI {
 		connCfg.SSHCLI = config.NewSSHDestination(opts.RemoteMachine.User, opts.RemoteMachine.Host, opts.RemoteMachine.Port)
diff --git a/internal/cli/config/connection.go b/internal/cli/config/connection.go
index c1be9c93..42f374ac 100644
--- a/internal/cli/config/connection.go
+++ b/internal/cli/config/connection.go
@@ -25,6 +25,7 @@ type MachineConnection struct {
 	TCP       *netip.AddrPort `yaml:"tcp,omitempty"`
 	Host      string          `yaml:"host,omitempty"`
 	PublicKey secret.Secret   `yaml:"public_key,omitempty"`
+	MachineID string          `yaml:"machine_id,omitempty"`
 }
 
 func (c MachineConnection) String() string {
diff --git a/internal/cli/config/connection_test.go b/internal/cli/config/connection_test.go
index 1dc3118c..3e6e9160 100644
--- a/internal/cli/config/connection_test.go
+++ b/internal/cli/config/connection_test.go
@@ -58,6 +58,15 @@ func TestMachineConnection_String(t *testing.T) {
 			conn: MachineConnection{},
 			want: "unknown connection",
 		},
+		{
+			name: "ssh connection with machine id",
+			conn: MachineConnection{
+				SSH:       "user@host.com",
+				MachineID: "ed98b9f7575308c340263cd279e3b568",
+			},
+			// String() does not use MachineID; just verifying entry with it is valid
+			want: "ssh://user@host.com",
+		},
 	}
 
 	for _, tt := range tests {