Azure
diff --git a/‎.github/workflows/agent-e2e-kind.yaml‎
Lines changed: 6 additions & 0 deletions b/‎.github/workflows/agent-e2e-kind.yaml‎
Lines changed: 6 additions & 0 deletions
diff --git a/‎cmd/agent/internal/cmd/agentupgrade.go‎
Lines changed: 28 additions & 0 deletions b/‎cmd/agent/internal/cmd/agentupgrade.go‎
Lines changed: 28 additions & 0 deletions
diff --git a/‎cmd/agent/internal/cmd/agentupgrade_test.go‎
Lines changed: 32 additions & 0 deletions b/‎cmd/agent/internal/cmd/agentupgrade_test.go‎
Lines changed: 32 additions & 0 deletions
diff --git a/‎cmd/agent/internal/cmd/cmd.go‎
Lines changed: 1 addition & 0 deletions b/‎cmd/agent/internal/cmd/cmd.go‎
Lines changed: 1 addition & 0 deletions
diff --git a/‎cmd/agent/internal/daemon/agentupgrade.go‎
Lines changed: 164 additions & 0 deletions b/‎cmd/agent/internal/daemon/agentupgrade.go‎
Lines changed: 164 additions & 0 deletions
@@ -179,6 +179,12 @@ jobs:
       - name: Validate node restart operation
         run: python3 ./hack/agent/e2e-kind/e2e.py --verbose validate-node-reboot-operation
 
+      - name: Validate agent upgrade operation
+        run: python3 ./hack/agent/e2e-kind/e2e.py --verbose validate-agent-upgrade-operation
+
+      - name: Validate agent upgrade rollback
+        run: python3 ./hack/agent/e2e-kind/e2e.py --verbose validate-agent-upgrade-rollback
+
       - name: Validate workload on agent node
         run: python3 ./hack/agent/e2e-kind/e2e.py --verbose validate-workload
 
 
@@ -0,0 +1,28 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+package cmd
+
+import (
+	"github.com/spf13/cobra"
+
+	"github.com/Azure/unbounded/cmd/agent/internal/daemon"
+)
+
+func newCmdRecordAgentUpgradeFailureSignal() *cobra.Command {
+	var message string
+
+	cmd := &cobra.Command{
+		Use:    "record-agent-upgrade-failure-signal",
+		Short:  "Record AgentUpgrade daemon recovery failure signal",
+		Hidden: true,
+		Args:   cobra.NoArgs,
+		RunE: func(*cobra.Command, []string) error {
+			return daemon.RecordAgentUpgradeFailureSignal(message)
+		},
+	}
+
+	cmd.Flags().StringVar(&message, "message", "", "failure message to record")
+
+	return cmd
+}
@@ -0,0 +1,32 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+package cmd
+
+import (
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/Azure/unbounded/pkg/agent/goalstates"
+)
+
+func TestRecordAgentUpgradeFailureSignalCommand(t *testing.T) {
+	dir := t.TempDir()
+	signalPath := filepath.Join(dir, "agent-upgrade-signal")
+	t.Setenv(goalstates.EnvDaemonAgentUpgradeSignalPath, signalPath)
+	require.NoError(t, os.WriteFile(signalPath, []byte(`{"operationName":"op-1"}`+"\n"), 0o600))
+
+	cmd := newCmdRecordAgentUpgradeFailureSignal()
+	cmd.SetArgs([]string{
+		"--message", "rolled back to last good",
+	})
+	require.NoError(t, cmd.Execute())
+
+	data, err := os.ReadFile(signalPath)
+	require.NoError(t, err)
+	assert.JSONEq(t, `{"operationName":"op-1","failureMessage":"rolled back to last good"}`, string(data))
+}
@@ -29,6 +29,7 @@ func Run() {
 		newCmdDaemon(cmdCtx),
 		newCmdReset(cmdCtx),
 		newCmdVersion(),
+		newCmdRecordAgentUpgradeFailureSignal(),
 	)
 
 	if err := root.Execute(); err != nil {
 
@@ -0,0 +1,164 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+package daemon
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"log/slog"
+	"os"
+	"strings"
+
+	"github.com/Azure/unbounded/pkg/agent/agentbinary"
+	"github.com/Azure/unbounded/pkg/agent/goalstates"
+)
+
+const (
+	agentUpgradeDownloadURLParameter = "downloadURL"
+	agentUpgradeBinaryMode           = 0o755
+)
+
+// agentUpgradeSignal is the JSON payload for pending and failure signals.
+type agentUpgradeSignal struct {
+	OperationName             string `json:"operationName"`
+	ObservedMachineGeneration int64  `json:"observedMachineGeneration,omitempty"`
+	// FailureMessage is set only after recovery reports a failed upgraded daemon.
+	FailureMessage string `json:"failureMessage,omitempty"`
+}
+
+// agentUpgradeSignalOperator manages persistent AgentUpgrade signal files.
+type agentUpgradeSignalOperator interface {
+	RecordPending(operationName string, observedMachineGeneration int64) error
+	RecordFailure(message string) error
+	Read() (*agentUpgradeSignal, error)
+	Clear() error
+}
+
+// fileAgentUpgradeSignalOperator stores AgentUpgrade signals on disk.
+type fileAgentUpgradeSignalOperator struct {
+	path string
+}
+
+func agentUpgradeDownloadURL(parameters map[string]string) (string, error) {
+	downloadURL := strings.TrimSpace(parameters[agentUpgradeDownloadURLParameter])
+	if downloadURL == "" {
+		return "", fmt.Errorf("missing required parameter %q", agentUpgradeDownloadURLParameter)
+	}
+
+	return downloadURL, nil
+}
+
+func upgradeDaemonBinary(ctx context.Context, log *slog.Logger, downloadURL string) error {
+	paths, err := goalstates.ResolvedAgentUpgradePaths()
+	if err != nil {
+		return fmt.Errorf("resolve current daemon binary symlink: %w", err)
+	}
+	targetPath := paths.NextTargetPath()
+	if err := agentbinary.InstallAndSwitchFromTarGz(ctx, downloadURL, paths, agentUpgradeBinaryMode); err != nil {
+		return err
+	}
+
+	log.Info("staged upgraded daemon binary",
+		"url", downloadURL,
+		"previous", paths.CurrentTargetPath,
+		"current", targetPath,
+	)
+
+	return nil
+}
+
+func newAgentUpgradeSignalOperator() (agentUpgradeSignalOperator, error) {
+	paths, err := goalstates.ResolvedAgentUpgradePaths()
+	if err != nil {
+		return nil, fmt.Errorf("resolve AgentUpgrade signal path: %w", err)
+	}
+
+	return newAgentUpgradeSignalOperatorForPath(paths.SignalPath), nil
+}
+
+func newAgentUpgradeSignalOperatorForPath(path string) agentUpgradeSignalOperator {
+	return fileAgentUpgradeSignalOperator{path: path}
+}
+
+func (o fileAgentUpgradeSignalOperator) RecordPending(operationName string, observedMachineGeneration int64) error {
+	return o.write(agentUpgradeSignal{
+		OperationName:             operationName,
+		ObservedMachineGeneration: observedMachineGeneration,
+	})
+}
+
+func (o fileAgentUpgradeSignalOperator) RecordFailure(message string) error {
+	pending, err := o.Read()
+	if err != nil {
+		return fmt.Errorf("read pending AgentUpgrade operation signal: %w", err)
+	}
+	if pending == nil {
+		slog.Warn("no pending AgentUpgrade operation signal found; skipping failure signal", "path", o.path)
+		return nil
+	}
+
+	message = strings.TrimSpace(message)
+	if message == "" {
+		message = "AgentUpgrade daemon failed after switching binary"
+	}
+
+	return o.write(agentUpgradeSignal{
+		OperationName:  pending.OperationName,
+		FailureMessage: message,
+	})
+}
+
+func (o fileAgentUpgradeSignalOperator) Clear() error {
+	if err := os.Remove(o.path); err != nil && !errors.Is(err, os.ErrNotExist) {
+		return err
+	}
+
+	return nil
+}
+
+func (o fileAgentUpgradeSignalOperator) write(signal agentUpgradeSignal) error {
+	data, err := json.Marshal(signal)
+	if err != nil {
+		return err
+	}
+
+	return writeFile(o.path, append(data, '\n'), 0o600)
+}
+
+func (o fileAgentUpgradeSignalOperator) Read() (*agentUpgradeSignal, error) {
+	data, err := os.ReadFile(o.path)
+	if err != nil {
+		if errors.Is(err, os.ErrNotExist) {
+			return nil, nil
+		}
+
+		return nil, err
+	}
+
+	var signal agentUpgradeSignal
+	if err := json.Unmarshal(data, &signal); err != nil {
+		return nil, fmt.Errorf("decode AgentUpgrade signal %s: %w", o.path, err)
+	}
+	signal.OperationName = strings.TrimSpace(signal.OperationName)
+	signal.FailureMessage = strings.TrimSpace(signal.FailureMessage)
+	if signal.OperationName == "" {
+		slog.Warn("AgentUpgrade signal missing operation name; ignoring signal", "path", o.path)
+		return nil, nil
+	}
+
+	return &signal, nil
+}
+
+// RecordAgentUpgradeFailureSignal records that the daemon failed after an
+// AgentUpgrade.
+func RecordAgentUpgradeFailureSignal(message string) error {
+	signals, err := newAgentUpgradeSignalOperator()
+	if err != nil {
+		return err
+	}
+
+	return signals.RecordFailure(message)
+}
Original file line number	Diff line number	Diff line change
`@@ -29,6 +29,7 @@ func Run() {`
`29`	`29`	`newCmdDaemon(cmdCtx),`
`30`	`30`	`newCmdReset(cmdCtx),`
`31`	`31`	`newCmdVersion(),`
	`32`	`+ newCmdRecordAgentUpgradeFailureSignal(),`
`32`	`33`	`)`
`33`	`34`
`34`	`35`	`if err := root.Execute(); err != nil {`