From 1898f67fe0863ed65f016e0a90194c09d782a769 Mon Sep 17 00:00:00 2001
From: Kyle Carberry <kyle@coder.com>
Date: Sun, 6 Nov 2022 14:46:51 -0800
Subject: [PATCH] fix: Ensure the session token is properly passed to instance
 identity (#4923)

Fixes #4921.
---
 cli/agent.go      |  2 +-
 cli/agent_test.go | 20 ++++++++++++++++++++
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/cli/agent.go b/cli/agent.go
index 2c7a4706a1..63455283da 100644
--- a/cli/agent.go
+++ b/cli/agent.go
@@ -165,7 +165,7 @@ func workspaceAgent() *cobra.Command {
 						return "", err
 					}
 					client.SessionToken = resp.SessionToken
-					return "", nil
+					return resp.SessionToken, nil
 				},
 				EnvironmentVariables: map[string]string{
 					"GIT_ASKPASS": executablePath,
diff --git a/cli/agent_test.go b/cli/agent_test.go
index f487ebfc00..a7ccd3b2b3 100644
--- a/cli/agent_test.go
+++ b/cli/agent_test.go
@@ -2,8 +2,11 @@ package cli_test
 
 import (
 	"context"
+	"runtime"
+	"strings"
 	"testing"
 
+	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
@@ -194,6 +197,23 @@ func TestWorkspaceAgent(t *testing.T) {
 			_, err := dialer.Ping(ctx)
 			return err == nil
 		}, testutil.WaitMedium, testutil.IntervalFast)
+
+		sshClient, err := dialer.SSHClient()
+		require.NoError(t, err)
+		defer sshClient.Close()
+		session, err := sshClient.NewSession()
+		require.NoError(t, err)
+		defer session.Close()
+		key := "CODER_AGENT_TOKEN"
+		command := "sh -c 'echo $" + key + "'"
+		if runtime.GOOS == "windows" {
+			command = "cmd.exe /c echo %" + key + "%"
+		}
+		token, err := session.CombinedOutput(command)
+		require.NoError(t, err)
+		_, err = uuid.Parse(strings.TrimSpace(string(token)))
+		require.NoError(t, err)
+
 		cancelFunc()
 		err = <-errC
 		require.NoError(t, err)