feat: add system update and uninstall commands to the ggo CLI, enhance agent connection handling and worker restart logic

Author: Code2Life

cmd/ggo/main.go

@@ -10,6 +10,7 @@ import (
 	"github.com/NexusGPU/gpu-go/cmd/ggo/libs"
 	"github.com/NexusGPU/gpu-go/cmd/ggo/share"
 	"github.com/NexusGPU/gpu-go/cmd/ggo/studio"
+	"github.com/NexusGPU/gpu-go/cmd/ggo/system"
 	"github.com/NexusGPU/gpu-go/cmd/ggo/use"
 	"github.com/NexusGPU/gpu-go/cmd/ggo/version"
 	"github.com/NexusGPU/gpu-go/cmd/ggo/worker"
@@ -57,6 +58,8 @@ func init() {
 	rootCmd.AddCommand(deps.NewDepsCmd())
 	rootCmd.AddCommand(studio.NewStudioCmd())
 	rootCmd.AddCommand(libs.NewLibsCmd())
+	rootCmd.AddCommand(system.NewUpdateCmd())
+	rootCmd.AddCommand(system.NewUninstallCmd())
 
 	// Auth commands (login/logout at root level for convenience)
 	rootCmd.AddCommand(auth.NewLoginCmd())

cmd/ggo/system/commands.go

@@ -0,0 +1,54 @@
+package system
+
+import (
+	"os"
+	"os/exec"
+	"runtime"
+
+	"github.com/spf13/cobra"
+)
+
+// NewUpdateCmd creates the update command.
+func NewUpdateCmd() *cobra.Command {
+	return newScriptCmd(
+		scriptActionUpdate,
+		"update",
+		"Update ggo to the latest version",
+		"Downloads and runs the platform install script from the CDN.",
+		"ggo update",
+	)
+}
+
+// NewUninstallCmd creates the uninstall command.
+func NewUninstallCmd() *cobra.Command {
+	return newScriptCmd(
+		scriptActionUninstall,
+		"uninstall",
+		"Uninstall ggo from this machine",
+		"Downloads and runs the platform uninstall script from the CDN.",
+		"ggo uninstall",
+	)
+}
+
+func newScriptCmd(action scriptAction, use, short, long, example string) *cobra.Command {
+	cmd := &cobra.Command{
+		Use:     use,
+		Short:   short,
+		Long:    long,
+		Example: example,
+		Args:    cobra.NoArgs,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			command, cmdArgs, err := buildScriptCommand(action, runtime.GOOS)
+			if err != nil {
+				return err
+			}
+			execCmd := exec.Command(command, cmdArgs...)
+			execCmd.Stdout = os.Stdout
+			execCmd.Stderr = os.Stderr
+			execCmd.Stdin = os.Stdin
+			return execCmd.Run()
+		},
+	}
+
+	return cmd
+}

cmd/ggo/system/system.go

@@ -0,0 +1,49 @@
+package system
+
+import (
+	"fmt"
+)
+
+type scriptAction string
+
+const (
+	scriptActionUpdate    scriptAction = "update"
+	scriptActionUninstall scriptAction = "uninstall"
+)
+
+const cdnBaseURL = "https://cdn.tensor-fusion.ai/archive/gpugo"
+
+func buildScriptCommand(action scriptAction, goos string) (string, []string, error) {
+	scriptBase, err := scriptBaseName(action)
+	if err != nil {
+		return "", nil, err
+	}
+
+	switch goos {
+	case "linux", "darwin":
+		url := fmt.Sprintf("%s/%s.sh", cdnBaseURL, scriptBase)
+		return "sh", []string{"-c", fmt.Sprintf("curl -sfL %s | sh", url)}, nil
+	case "windows":
+		url := fmt.Sprintf("%s/%s.ps1", cdnBaseURL, scriptBase)
+		return "powershell", []string{
+			"-NoProfile",
+			"-ExecutionPolicy",
+			"Bypass",
+			"-Command",
+			fmt.Sprintf("irm %s | iex", url),
+		}, nil
+	default:
+		return "", nil, fmt.Errorf("unsupported OS: %s", goos)
+	}
+}
+
+func scriptBaseName(action scriptAction) (string, error) {
+	switch action {
+	case scriptActionUpdate:
+		return "install", nil
+	case scriptActionUninstall:
+		return "uninstall", nil
+	default:
+		return "", fmt.Errorf("unsupported action: %s", action)
+	}
+}

cmd/ggo/system/system_test.go

@@ -0,0 +1,59 @@
+package system
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestBuildScriptCommand_UninstallLinux(t *testing.T) {
+	cmd, args, err := buildScriptCommand(scriptActionUninstall, "linux")
+	require.NoError(t, err)
+	assert.Equal(t, "sh", cmd)
+	assert.Equal(t, []string{
+		"-c",
+		"curl -sfL https://cdn.tensor-fusion.ai/archive/gpugo/uninstall.sh | sh",
+	}, args)
+}
+
+func TestBuildScriptCommand_UninstallWindows(t *testing.T) {
+	cmd, args, err := buildScriptCommand(scriptActionUninstall, "windows")
+	require.NoError(t, err)
+	assert.Equal(t, "powershell", cmd)
+	assert.Equal(t, []string{
+		"-NoProfile",
+		"-ExecutionPolicy",
+		"Bypass",
+		"-Command",
+		"irm https://cdn.tensor-fusion.ai/archive/gpugo/uninstall.ps1 | iex",
+	}, args)
+}
+
+func TestBuildScriptCommand_UpdateDarwin(t *testing.T) {
+	cmd, args, err := buildScriptCommand(scriptActionUpdate, "darwin")
+	require.NoError(t, err)
+	assert.Equal(t, "sh", cmd)
+	assert.Equal(t, []string{
+		"-c",
+		"curl -sfL https://cdn.tensor-fusion.ai/archive/gpugo/install.sh | sh",
+	}, args)
+}
+
+func TestBuildScriptCommand_UpdateWindows(t *testing.T) {
+	cmd, args, err := buildScriptCommand(scriptActionUpdate, "windows")
+	require.NoError(t, err)
+	assert.Equal(t, "powershell", cmd)
+	assert.Equal(t, []string{
+		"-NoProfile",
+		"-ExecutionPolicy",
+		"Bypass",
+		"-Command",
+		"irm https://cdn.tensor-fusion.ai/archive/gpugo/install.ps1 | iex",
+	}, args)
+}
+
+func TestBuildScriptCommand_UnsupportedOS(t *testing.T) {
+	_, _, err := buildScriptCommand(scriptActionUpdate, "plan9")
+	require.Error(t, err)
+}

internal/agent/agent.go

@@ -446,6 +446,7 @@ func (a *Agent) convertToWorkerInfos(apiWorkers []api.WorkerConfig) ([]*hvApi.Wo
 		envVars["TF_ENABLE_LOG"] = "1"
 		envVars[EnvURLAuth] = "1"
 		envVars[EnvAuthorizedKeyPath] = filepath.Join(a.paths.ConfigDir(), w.WorkerID+"_share_codes")
+		envVars[EnvConnectionInfoPath] = a.connectionsDir
 
 		// Set hard limiter environment variables for Fractional GPU support
 		// TODO: use MIG for partitioned
@@ -1070,7 +1071,7 @@ func (a *Agent) collectWorkerStatus(
 		return a.collectWorkerStatusFromHypervisor(forceRefresh, connectionChanges, currentConnections, gpuChanges)
 	}
 	// Fallback to config-based status
-	return a.collectWorkerStatusFromConfig(forceRefresh, gpuChanges)
+	return a.collectWorkerStatusFromConfig(forceRefresh, currentConnections, gpuChanges)
 }
 
 // collectWorkerStatusFromHypervisor gets worker status from hypervisor
@@ -1162,7 +1163,11 @@ func (a *Agent) collectWorkerStatusFromHypervisor(
 }
 
 // collectWorkerStatusFromConfig gets worker status from local config (fallback)
-func (a *Agent) collectWorkerStatusFromConfig(forceRefresh bool, gpuChanges map[string]bool) ([]api.WorkerStatus, error) {
+func (a *Agent) collectWorkerStatusFromConfig(
+	forceRefresh bool,
+	currentConnections map[string][]string,
+	gpuChanges map[string]bool,
+) ([]api.WorkerStatus, error) {
 	workerConfigs, err := a.config.LoadWorkers()
 	if err != nil {
 		return nil, err
@@ -1179,6 +1184,10 @@ func (a *Agent) collectWorkerStatusFromConfig(forceRefresh bool, gpuChanges map[
 		workerChanged := true
 		connectionChanged := true
 		gpuChanged := forceRefresh || len(gpuChanges) > 0
+		connections := w.Connections
+		if connLines, ok := currentConnections[w.WorkerID]; ok {
+			connections = parseConnectionsToAPI(connLines)
+		}
 
 		gpuIndices := resolveWorkerGPUIndices(w.WorkerID, w.GPUIndices, w.GPUIDs, gpuIndexByID)
 		workerStatuses[i] = api.WorkerStatus{
@@ -1187,7 +1196,7 @@ func (a *Agent) collectWorkerStatusFromConfig(forceRefresh bool, gpuChanges map[
 			PID:               w.PID,
 			GPUIDs:            w.GPUIDs,
 			GPUIndices:        gpuIndices,
-			Connections:       w.Connections,
+			Connections:       connections,
 			WorkerChanged:     &workerChanged,
 			ConnectionChanged: &connectionChanged,
 			GPUChanged:        &gpuChanged,

internal/agent/agent_test.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"net/http"
 	"net/http/httptest"
+	"os"
 	"path/filepath"
 	"sync"
 	"testing"
@@ -277,6 +278,70 @@ func TestAgent_ReportStatus(t *testing.T) {
 	assert.Equal(t, 12345, receivedReq.Workers[0].PID)
 }
 
+func TestAgent_ReportStatus_ReadsConnectionFilesEveryTime(t *testing.T) {
+	tmpDir := t.TempDir()
+	configDir := filepath.Join(tmpDir, "config")
+	stateDir := filepath.Join(tmpDir, "state")
+	connectionsDir := filepath.Join(tmpDir, "connections")
+	require.NoError(t, os.MkdirAll(connectionsDir, 0755))
+
+	var mu sync.Mutex
+	var receivedReqs []api.AgentStatusRequest
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		var req api.AgentStatusRequest
+		_ = json.NewDecoder(r.Body).Decode(&req)
+		mu.Lock()
+		receivedReqs = append(receivedReqs, req)
+		mu.Unlock()
+		w.Header().Set("Content-Type", "application/json")
+		_ = json.NewEncoder(w).Encode(api.SuccessResponse{Success: true})
+	}))
+	defer server.Close()
+
+	configMgr := config.NewManager(configDir, stateDir)
+	err := configMgr.SaveGPUs([]config.GPUConfig{
+		{GPUID: "GPU-0", GPUIndex: 0, Vendor: "nvidia", Model: "RTX 4090", VRAMMb: 24576},
+	})
+	require.NoError(t, err)
+	err = configMgr.SaveWorkers([]config.WorkerConfig{
+		{WorkerID: "worker_1", GPUIDs: []string{"GPU-0"}, ListenPort: 9001, Enabled: true, Status: "running", PID: 12345},
+	})
+	require.NoError(t, err)
+	err = os.WriteFile(filepath.Join(connectionsDir, "worker_1.txt"), []byte("10.1.1.1,1234,11\n"), 0644)
+	require.NoError(t, err)
+
+	client := api.NewClient(
+		api.WithBaseURL(server.URL),
+		api.WithAgentSecret("gpugo_secret123"),
+	)
+	agent := &Agent{
+		client:         client,
+		config:         configMgr,
+		ctx:            context.Background(),
+		agentID:        "agent_test123",
+		paths:          platform.DefaultPaths().WithConfigDir(configDir),
+		connectionsDir: connectionsDir,
+	}
+
+	err = agent.reportStatus()
+	require.NoError(t, err)
+
+	err = os.WriteFile(filepath.Join(connectionsDir, "worker_1.txt"), []byte("10.2.2.2,5678,22\n"), 0644)
+	require.NoError(t, err)
+	err = agent.reportStatus()
+	require.NoError(t, err)
+
+	mu.Lock()
+	defer mu.Unlock()
+	require.Len(t, receivedReqs, 2)
+	require.Len(t, receivedReqs[0].Workers, 1)
+	require.Len(t, receivedReqs[0].Workers[0].Connections, 1)
+	assert.Equal(t, "10.1.1.1", receivedReqs[0].Workers[0].Connections[0].ClientIP)
+	require.Len(t, receivedReqs[1].Workers, 1)
+	require.Len(t, receivedReqs[1].Workers[0].Connections, 1)
+	assert.Equal(t, "10.2.2.2", receivedReqs[1].Workers[0].Connections[0].ClientIP)
+}
+
 func TestAgent_HandleHeartbeatResponse(t *testing.T) {
 	tmpDir := t.TempDir()
 	configDir := filepath.Join(tmpDir, "config")
@@ -492,6 +557,44 @@ func TestAgent_WithHypervisor(t *testing.T) {
 	agent.Stop()
 }
 
+func TestAgent_ConvertToWorkerInfos_IncludesConnectionInfoPath(t *testing.T) {
+	tmpDir := t.TempDir()
+	configDir := filepath.Join(tmpDir, "config")
+	stateDir := filepath.Join(tmpDir, "state")
+
+	configMgr := config.NewManager(configDir, stateDir)
+	cfg := &config.Config{
+		ConfigVersion: 1,
+		AgentID:       "agent_test123",
+		AgentSecret:   "gpugo_secret123",
+		ServerURL:     "http://localhost",
+		License: api.License{
+			Plain:     "test|pro|9999999999",
+			Encrypted: "enc",
+		},
+	}
+	err := configMgr.SaveConfig(cfg)
+	require.NoError(t, err)
+
+	agent := NewAgent(api.NewClient(), configMgr)
+	agent.workerBinaryPath = "/bin/true"
+	agent.connectionsDir = filepath.Join(tmpDir, "connections")
+
+	infos, err := agent.convertToWorkerInfos([]api.WorkerConfig{
+		{
+			WorkerID:   "worker_1",
+			GPUIDs:     []string{"gpu-0"},
+			ListenPort: 9001,
+			Enabled:    true,
+		},
+	})
+	require.NoError(t, err)
+	require.Len(t, infos, 1)
+	require.NotNil(t, infos[0].WorkerRunningInfo)
+
+	assert.Equal(t, agent.connectionsDir, infos[0].WorkerRunningInfo.Env[EnvConnectionInfoPath])
+}
+
 func TestAgent_LicenseParsing(t *testing.T) {
 	tmpDir := t.TempDir()
 	configDir := filepath.Join(tmpDir, "config")