feat: add agent metadata (#6614)

agent/agent_test.go

@@ -33,6 +33,7 @@ import (
 	"github.com/stretchr/testify/require"
 	"go.uber.org/goleak"
 	"golang.org/x/crypto/ssh"
+	"golang.org/x/exp/maps"
 	"golang.org/x/xerrors"
 	"tailscale.com/net/speedtest"
 	"tailscale.com/tailcfg"
@@ -61,7 +62,7 @@ func TestAgent_Stats_SSH(t *testing.T) {
 	defer cancel()
 
 	//nolint:dogsled
-	conn, _, stats, _, _ := setupAgent(t, agentsdk.Metadata{}, 0)
+	conn, _, stats, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
 
 	sshClient, err := conn.SSHClient(ctx)
 	require.NoError(t, err)
@@ -94,7 +95,7 @@ func TestAgent_Stats_ReconnectingPTY(t *testing.T) {
 	defer cancel()
 
 	//nolint:dogsled
-	conn, _, stats, _, _ := setupAgent(t, agentsdk.Metadata{}, 0)
+	conn, _, stats, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
 
 	ptyConn, err := conn.ReconnectingPTY(ctx, uuid.New(), 128, 128, "/bin/bash")
 	require.NoError(t, err)
@@ -124,7 +125,7 @@ func TestAgent_Stats_Magic(t *testing.T) {
 		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
 		defer cancel()
 		//nolint:dogsled
-		conn, _, _, _, _ := setupAgent(t, agentsdk.Metadata{}, 0)
+		conn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
 		sshClient, err := conn.SSHClient(ctx)
 		require.NoError(t, err)
 		defer sshClient.Close()
@@ -151,7 +152,7 @@ func TestAgent_Stats_Magic(t *testing.T) {
 		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
 		defer cancel()
 		//nolint:dogsled
-		conn, _, stats, _, _ := setupAgent(t, agentsdk.Metadata{}, 0)
+		conn, _, stats, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
 		sshClient, err := conn.SSHClient(ctx)
 		require.NoError(t, err)
 		defer sshClient.Close()
@@ -186,7 +187,7 @@ func TestAgent_Stats_Magic(t *testing.T) {
 
 func TestAgent_SessionExec(t *testing.T) {
 	t.Parallel()
-	session := setupSSHSession(t, agentsdk.Metadata{})
+	session := setupSSHSession(t, agentsdk.Manifest{})
 
 	command := "echo test"
 	if runtime.GOOS == "windows" {
@@ -199,7 +200,7 @@ func TestAgent_SessionExec(t *testing.T) {
 
 func TestAgent_GitSSH(t *testing.T) {
 	t.Parallel()
-	session := setupSSHSession(t, agentsdk.Metadata{})
+	session := setupSSHSession(t, agentsdk.Manifest{})
 	command := "sh -c 'echo $GIT_SSH_COMMAND'"
 	if runtime.GOOS == "windows" {
 		command = "cmd.exe /c echo %GIT_SSH_COMMAND%"
@@ -219,7 +220,7 @@ func TestAgent_SessionTTYShell(t *testing.T) {
 		// it seems like it could be either.
 		t.Skip("ConPTY appears to be inconsistent on Windows.")
 	}
-	session := setupSSHSession(t, agentsdk.Metadata{})
+	session := setupSSHSession(t, agentsdk.Manifest{})
 	command := "sh"
 	if runtime.GOOS == "windows" {
 		command = "cmd.exe"
@@ -242,7 +243,7 @@ func TestAgent_SessionTTYShell(t *testing.T) {
 
 func TestAgent_SessionTTYExitCode(t *testing.T) {
 	t.Parallel()
-	session := setupSSHSession(t, agentsdk.Metadata{})
+	session := setupSSHSession(t, agentsdk.Manifest{})
 	command := "areallynotrealcommand"
 	err := session.RequestPty("xterm", 128, 128, ssh.TerminalModes{})
 	require.NoError(t, err)
@@ -281,7 +282,7 @@ func TestAgent_Session_TTY_MOTD(t *testing.T) {
 	// Set HOME so we can ensure no ~/.hushlogin is present.
 	t.Setenv("HOME", tmpdir)
 
-	session := setupSSHSession(t, agentsdk.Metadata{
+	session := setupSSHSession(t, agentsdk.Manifest{
 		MOTDFile: name,
 	})
 	err = session.RequestPty("xterm", 128, 128, ssh.TerminalModes{})
@@ -327,7 +328,7 @@ func TestAgent_Session_TTY_Hushlogin(t *testing.T) {
 	// Set HOME so we can ensure ~/.hushlogin is present.
 	t.Setenv("HOME", tmpdir)
 
-	session := setupSSHSession(t, agentsdk.Metadata{
+	session := setupSSHSession(t, agentsdk.Manifest{
 		MOTDFile: name,
 	})
 	err = session.RequestPty("xterm", 128, 128, ssh.TerminalModes{})
@@ -357,7 +358,7 @@ func TestAgent_Session_TTY_FastCommandHasOutput(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
 	defer cancel()
 	//nolint:dogsled
-	conn, _, _, _, _ := setupAgent(t, agentsdk.Metadata{}, 0)
+	conn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
 	sshClient, err := conn.SSHClient(ctx)
 	require.NoError(t, err)
 	defer sshClient.Close()
@@ -407,7 +408,7 @@ func TestAgent_Session_TTY_HugeOutputIsNotLost(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
 	defer cancel()
 	//nolint:dogsled
-	conn, _, _, _, _ := setupAgent(t, agentsdk.Metadata{}, 0)
+	conn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
 	sshClient, err := conn.SSHClient(ctx)
 	require.NoError(t, err)
 	defer sshClient.Close()
@@ -706,7 +707,7 @@ func TestAgent_SFTP(t *testing.T) {
 		home = "/" + strings.ReplaceAll(home, "\\", "/")
 	}
 	//nolint:dogsled
-	conn, _, _, _, _ := setupAgent(t, agentsdk.Metadata{}, 0)
+	conn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
 	sshClient, err := conn.SSHClient(ctx)
 	require.NoError(t, err)
 	defer sshClient.Close()
@@ -738,7 +739,7 @@ func TestAgent_SCP(t *testing.T) {
 	defer cancel()
 
 	//nolint:dogsled
-	conn, _, _, _, _ := setupAgent(t, agentsdk.Metadata{}, 0)
+	conn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
 	sshClient, err := conn.SSHClient(ctx)
 	require.NoError(t, err)
 	defer sshClient.Close()
@@ -757,7 +758,7 @@ func TestAgent_EnvironmentVariables(t *testing.T) {
 	t.Parallel()
 	key := "EXAMPLE"
 	value := "value"
-	session := setupSSHSession(t, agentsdk.Metadata{
+	session := setupSSHSession(t, agentsdk.Manifest{
 		EnvironmentVariables: map[string]string{
 			key: value,
 		},
@@ -774,7 +775,7 @@ func TestAgent_EnvironmentVariables(t *testing.T) {
 func TestAgent_EnvironmentVariableExpansion(t *testing.T) {
 	t.Parallel()
 	key := "EXAMPLE"
-	session := setupSSHSession(t, agentsdk.Metadata{
+	session := setupSSHSession(t, agentsdk.Manifest{
 		EnvironmentVariables: map[string]string{
 			key: "$SOMETHINGNOTSET",
 		},
@@ -801,7 +802,7 @@ func TestAgent_CoderEnvVars(t *testing.T) {
 		t.Run(key, func(t *testing.T) {
 			t.Parallel()
 
-			session := setupSSHSession(t, agentsdk.Metadata{})
+			session := setupSSHSession(t, agentsdk.Manifest{})
 			command := "sh -c 'echo $" + key + "'"
 			if runtime.GOOS == "windows" {
 				command = "cmd.exe /c echo %" + key + "%"
@@ -824,7 +825,7 @@ func TestAgent_SSHConnectionEnvVars(t *testing.T) {
 		t.Run(key, func(t *testing.T) {
 			t.Parallel()
 
-			session := setupSSHSession(t, agentsdk.Metadata{})
+			session := setupSSHSession(t, agentsdk.Manifest{})
 			command := "sh -c 'echo $" + key + "'"
 			if runtime.GOOS == "windows" {
 				command = "cmd.exe /c echo %" + key + "%"
@@ -848,7 +849,7 @@ func TestAgent_StartupScript(t *testing.T) {
 		client := &client{
 			t:       t,
 			agentID: uuid.New(),
-			metadata: agentsdk.Metadata{
+			manifest: agentsdk.Manifest{
 				StartupScript: command,
 				DERPMap:       &tailcfg.DERPMap{},
 			},
@@ -879,7 +880,7 @@ func TestAgent_StartupScript(t *testing.T) {
 		client := &client{
 			t:       t,
 			agentID: uuid.New(),
-			metadata: agentsdk.Metadata{
+			manifest: agentsdk.Manifest{
 				StartupScript: command,
 				DERPMap:       &tailcfg.DERPMap{},
 			},
@@ -912,13 +913,125 @@ func TestAgent_StartupScript(t *testing.T) {
 	})
 }
 
+func TestAgent_Metadata(t *testing.T) {
+	t.Parallel()
+
+	t.Run("Once", func(t *testing.T) {
+		t.Parallel()
+		script := "echo -n hello"
+		if runtime.GOOS == "windows" {
+			script = "powershell " + script
+		}
+		//nolint:dogsled
+		_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
+			Metadata: []codersdk.WorkspaceAgentMetadataDescription{
+				{
+					Key:      "greeting",
+					Interval: 0,
+					Script:   script,
+				},
+			},
+		}, 0)
+
+		var gotMd map[string]agentsdk.PostMetadataRequest
+		require.Eventually(t, func() bool {
+			gotMd = client.getMetadata()
+			return len(gotMd) == 1
+		}, testutil.WaitShort, testutil.IntervalMedium)
+
+		collectedAt := gotMd["greeting"].CollectedAt
+
+		require.Never(t, func() bool {
+			gotMd = client.getMetadata()
+			if len(gotMd) != 1 {
+				panic("unexpected number of metadata")
+			}
+			return !gotMd["greeting"].CollectedAt.Equal(collectedAt)
+		}, testutil.WaitShort, testutil.IntervalMedium)
+	})
+
+	t.Run("Many", func(t *testing.T) {
+		if runtime.GOOS == "windows" {
+			// Shell scripting in Windows is a pain, and we have already tested
+			// that the OS logic works in the simpler "Once" test above.
+			t.Skip()
+		}
+		t.Parallel()
+
+		dir := t.TempDir()
+
+		const reportInterval = 2
+		const intervalUnit = 100 * time.Millisecond
+		var (
+			greetingPath = filepath.Join(dir, "greeting")
+			script       = "echo hello | tee -a " + greetingPath
+		)
+		_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
+			Metadata: []codersdk.WorkspaceAgentMetadataDescription{
+				{
+					Key:      "greeting",
+					Interval: reportInterval,
+					Script:   script,
+				},
+				{
+					Key:      "bad",
+					Interval: reportInterval,
+					Script:   "exit 1",
+				},
+			},
+		}, 0)
+
+		require.Eventually(t, func() bool {
+			return len(client.getMetadata()) == 2
+		}, testutil.WaitShort, testutil.IntervalMedium)
+
+		for start := time.Now(); time.Since(start) < testutil.WaitMedium; time.Sleep(testutil.IntervalMedium) {
+			md := client.getMetadata()
+			if len(md) != 2 {
+				panic("unexpected number of metadata entries")
+			}
+
+			require.Equal(t, "hello\n", md["greeting"].Value)
+			require.Equal(t, "exit status 1", md["bad"].Error)
+
+			greetingByt, err := os.ReadFile(greetingPath)
+			require.NoError(t, err)
+
+			var (
+				numGreetings      = bytes.Count(greetingByt, []byte("hello"))
+				idealNumGreetings = time.Since(start) / (reportInterval * intervalUnit)
+				// We allow a 50% error margin because the report loop may backlog
+				// in CI and other toasters. In production, there is no hard
+				// guarantee on timing either, and the frontend gives similar
+				// wiggle room to the staleness of the value.
+				upperBound = int(idealNumGreetings) + 1
+				lowerBound = (int(idealNumGreetings) / 2)
+			)
+
+			if idealNumGreetings < 50 {
+				// There is an insufficient sample size.
+				continue
+			}
+
+			t.Logf("numGreetings: %d, idealNumGreetings: %d", numGreetings, idealNumGreetings)
+			// The report loop may slow down on load, but it should never, ever
+			// speed up.
+			if numGreetings > upperBound {
+				t.Fatalf("too many greetings: %d > %d in %v", numGreetings, upperBound, time.Since(start))
+			} else if numGreetings < lowerBound {
+				t.Fatalf("too few greetings: %d < %d", numGreetings, lowerBound)
+			}
+		}
+	})
+}
+
 func TestAgent_Lifecycle(t *testing.T) {
 	t.Parallel()
 
 	t.Run("StartTimeout", func(t *testing.T) {
 		t.Parallel()
 
-		_, client, _, _, _ := setupAgent(t, agentsdk.Metadata{
+		_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
 			StartupScript:        "sleep 5",
 			StartupScriptTimeout: time.Nanosecond,
 		}, 0)
@@ -947,7 +1060,7 @@ func TestAgent_Lifecycle(t *testing.T) {
 	t.Run("StartError", func(t *testing.T) {
 		t.Parallel()
 
-		_, client, _, _, _ := setupAgent(t, agentsdk.Metadata{
+		_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
 			StartupScript:        "false",
 			StartupScriptTimeout: 30 * time.Second,
 		}, 0)
@@ -976,7 +1089,7 @@ func TestAgent_Lifecycle(t *testing.T) {
 	t.Run("Ready", func(t *testing.T) {
 		t.Parallel()
 
-		_, client, _, _, _ := setupAgent(t, agentsdk.Metadata{
+		_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
 			StartupScript:        "true",
 			StartupScriptTimeout: 30 * time.Second,
 		}, 0)
@@ -1005,7 +1118,7 @@ func TestAgent_Lifecycle(t *testing.T) {
 	t.Run("ShuttingDown", func(t *testing.T) {
 		t.Parallel()
 
-		_, client, _, _, closer := setupAgent(t, agentsdk.Metadata{
+		_, client, _, _, closer := setupAgent(t, agentsdk.Manifest{
 			ShutdownScript:       "sleep 5",
 			StartupScriptTimeout: 30 * time.Second,
 		}, 0)
@@ -1043,7 +1156,7 @@ func TestAgent_Lifecycle(t *testing.T) {
 	t.Run("ShutdownTimeout", func(t *testing.T) {
 		t.Parallel()
 
-		_, client, _, _, closer := setupAgent(t, agentsdk.Metadata{
+		_, client, _, _, closer := setupAgent(t, agentsdk.Manifest{
 			ShutdownScript:        "sleep 5",
 			ShutdownScriptTimeout: time.Nanosecond,
 		}, 0)
@@ -1090,7 +1203,7 @@ func TestAgent_Lifecycle(t *testing.T) {
 	t.Run("ShutdownError", func(t *testing.T) {
 		t.Parallel()
 
-		_, client, _, _, closer := setupAgent(t, agentsdk.Metadata{
+		_, client, _, _, closer := setupAgent(t, agentsdk.Manifest{
 			ShutdownScript:        "false",
 			ShutdownScriptTimeout: 30 * time.Second,
 		}, 0)
@@ -1141,7 +1254,7 @@ func TestAgent_Lifecycle(t *testing.T) {
 		client := &client{
 			t:       t,
 			agentID: uuid.New(),
-			metadata: agentsdk.Metadata{
+			manifest: agentsdk.Manifest{
 				DERPMap:        tailnettest.RunDERPAndSTUN(t),
 				StartupScript:  "echo 1",
 				ShutdownScript: "echo " + expected,
@@ -1194,7 +1307,7 @@ func TestAgent_Startup(t *testing.T) {
 	t.Run("EmptyDirectory", func(t *testing.T) {
 		t.Parallel()
 
-		_, client, _, _, _ := setupAgent(t, agentsdk.Metadata{
+		_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
 			StartupScript:        "true",
 			StartupScriptTimeout: 30 * time.Second,
 			Directory:            "",
@@ -1208,7 +1321,7 @@ func TestAgent_Startup(t *testing.T) {
 	t.Run("HomeDirectory", func(t *testing.T) {
 		t.Parallel()
 
-		_, client, _, _, _ := setupAgent(t, agentsdk.Metadata{
+		_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
 			StartupScript:        "true",
 			StartupScriptTimeout: 30 * time.Second,
 			Directory:            "~",
@@ -1224,7 +1337,7 @@ func TestAgent_Startup(t *testing.T) {
 	t.Run("HomeEnvironmentVariable", func(t *testing.T) {
 		t.Parallel()
 
-		_, client, _, _, _ := setupAgent(t, agentsdk.Metadata{
+		_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
 			StartupScript:        "true",
 			StartupScriptTimeout: 30 * time.Second,
 			Directory:            "$HOME",
@@ -1251,7 +1364,7 @@ func TestAgent_ReconnectingPTY(t *testing.T) {
 	defer cancel()
 
 	//nolint:dogsled
-	conn, _, _, _, _ := setupAgent(t, agentsdk.Metadata{}, 0)
+	conn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
 	id := uuid.New()
 	netConn, err := conn.ReconnectingPTY(ctx, id, 100, 100, "/bin/bash")
 	require.NoError(t, err)
@@ -1353,7 +1466,7 @@ func TestAgent_Dial(t *testing.T) {
 			}()
 
 			//nolint:dogsled
-			conn, _, _, _, _ := setupAgent(t, agentsdk.Metadata{}, 0)
+			conn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
 			require.True(t, conn.AwaitReachable(context.Background()))
 			conn1, err := conn.DialContext(context.Background(), l.Addr().Network(), l.Addr().String())
 			require.NoError(t, err)
@@ -1375,7 +1488,7 @@ func TestAgent_Speedtest(t *testing.T) {
 	defer cancel()
 	derpMap := tailnettest.RunDERPAndSTUN(t)
 	//nolint:dogsled
-	conn, _, _, _, _ := setupAgent(t, agentsdk.Metadata{
+	conn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{
 		DERPMap: derpMap,
 	}, 0)
 	defer conn.Close()
@@ -1397,7 +1510,7 @@ func TestAgent_Reconnect(t *testing.T) {
 	client := &client{
 		t:       t,
 		agentID: agentID,
-		metadata: agentsdk.Metadata{
+		manifest: agentsdk.Manifest{
 			DERPMap: derpMap,
 		},
 		statsChan:   statsCh,
@@ -1432,7 +1545,7 @@ func TestAgent_WriteVSCodeConfigs(t *testing.T) {
 	client := &client{
 		t:       t,
 		agentID: uuid.New(),
-		metadata: agentsdk.Metadata{
+		manifest: agentsdk.Manifest{
 			GitAuthConfigs: 1,
 			DERPMap:        &tailcfg.DERPMap{},
 		},
@@ -1461,7 +1574,7 @@ func TestAgent_WriteVSCodeConfigs(t *testing.T) {
 
 func setupSSHCommand(t *testing.T, beforeArgs []string, afterArgs []string) *exec.Cmd {
 	//nolint:dogsled
-	agentConn, _, _, _, _ := setupAgent(t, agentsdk.Metadata{}, 0)
+	agentConn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
 	listener, err := net.Listen("tcp", "127.0.0.1:0")
 	require.NoError(t, err)
 	waitGroup := sync.WaitGroup{}
@@ -1504,7 +1617,7 @@ func setupSSHCommand(t *testing.T, beforeArgs []string, afterArgs []string) *exe
 	return exec.Command("ssh", args...)
 }
 
-func setupSSHSession(t *testing.T, options agentsdk.Metadata) *ssh.Session {
+func setupSSHSession(t *testing.T, options agentsdk.Manifest) *ssh.Session {
 	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
 	defer cancel()
 	//nolint:dogsled
@@ -1528,7 +1641,7 @@ func (c closeFunc) Close() error {
 	return c()
 }
 
-func setupAgent(t *testing.T, metadata agentsdk.Metadata, ptyTimeout time.Duration) (
+func setupAgent(t *testing.T, metadata agentsdk.Manifest, ptyTimeout time.Duration) (
 	*codersdk.WorkspaceAgentConn,
 	*client,
 	<-chan *agentsdk.Stats,
@@ -1548,7 +1661,7 @@ func setupAgent(t *testing.T, metadata agentsdk.Metadata, ptyTimeout time.Durati
 	c := &client{
 		t:           t,
 		agentID:     agentID,
-		metadata:    metadata,
+		manifest:    metadata,
 		statsChan:   statsCh,
 		coordinator: coordinator,
 	}
@@ -1631,7 +1744,8 @@ func assertWritePayload(t *testing.T, w io.Writer, payload []byte) {
 type client struct {
 	t                  *testing.T
 	agentID            uuid.UUID
-	metadata           agentsdk.Metadata
+	manifest           agentsdk.Manifest
+	metadata           map[string]agentsdk.PostMetadataRequest
 	statsChan          chan *agentsdk.Stats
 	coordinator        tailnet.Coordinator
 	lastWorkspaceAgent func()
@@ -1643,8 +1757,8 @@ type client struct {
 	logs            []agentsdk.StartupLog
 }
 
-func (c *client) Metadata(_ context.Context) (agentsdk.Metadata, error) {
-	return c.metadata, nil
+func (c *client) Manifest(_ context.Context) (agentsdk.Manifest, error) {
+	return c.manifest, nil
 }
 
 func (c *client) Listen(_ context.Context) (net.Conn, error) {
@@ -1718,6 +1832,22 @@ func (c *client) getStartup() agentsdk.PostStartupRequest {
 	return c.startup
 }
 
+func (c *client) getMetadata() map[string]agentsdk.PostMetadataRequest {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return maps.Clone(c.metadata)
+}
+
+func (c *client) PostMetadata(_ context.Context, key string, req agentsdk.PostMetadataRequest) error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if c.metadata == nil {
+		c.metadata = make(map[string]agentsdk.PostMetadataRequest)
+	}
+	c.metadata[key] = req
+	return nil
+}
+
 func (c *client) PostStartup(_ context.Context, startup agentsdk.PostStartupRequest) error {
 	c.mu.Lock()
 	defer c.mu.Unlock()