coder · f0ssel · Jun 25, 2024 · Jun 22, 2024 · Jun 22, 2024 · Jun 22, 2024
diff --git a/cli/ssh.go b/cli/ssh.go
@@ -12,6 +12,7 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
+	"slices"
 	"strings"
 	"sync"
 	"time"
@@ -57,6 +58,7 @@ func (r *RootCmd) ssh() *serpent.Command {
 		logDirPath       string
 		remoteForwards   []string
 		env              []string
+		usageApp         string
 		disableAutostart bool
 	)
 	client := new(codersdk.Client)
@@ -251,6 +253,15 @@ func (r *RootCmd) ssh() *serpent.Command {
 			stopPolling := tryPollWorkspaceAutostop(ctx, client, workspace)
 			defer stopPolling()
 
+			usageAppName := getUsageAppName(usageApp)
+			if usageAppName != "" {
+				closeUsage := client.UpdateWorkspaceUsageWithBodyContext(ctx, workspace.ID, codersdk.PostWorkspaceUsageRequest{
+					AgentID: workspaceAgent.ID,
+					AppName: usageAppName,
+				})
+				defer closeUsage()
+			}
+
 			if stdio {
 				rawSSH, err := conn.SSH(ctx)
 				if err != nil {
@@ -509,6 +520,12 @@ func (r *RootCmd) ssh() *serpent.Command {
 			FlagShorthand: "e",
 			Value:         serpent.StringArrayOf(&env),
 		},
+		{
+			Flag:        "usage-app",
+			Description: "Specifies the usage app to use for workspace activity tracking.",
+			Env:         "CODER_SSH_USAGE_APP",
+			Value:       serpent.StringOf(&usageApp),
+		},
 		sshDisableAutostartOption(serpent.BoolOf(&disableAutostart)),
 	}
 	return cmd
@@ -1044,3 +1061,20 @@ func (r stdioErrLogReader) Read(_ []byte) (int, error) {
 	r.l.Error(context.Background(), "reading from stdin in stdio mode is not allowed")
 	return 0, io.EOF
 }
+
+func getUsageAppName(usageApp string) codersdk.UsageAppName {
+	if usageApp == "" {
+		usageApp = "ssh"
+	}
+
+	allowedUsageApps := []string{
+		string(codersdk.UsageAppNameSSH),
+		string(codersdk.UsageAppNameVscode),
+		string(codersdk.UsageAppNameJetbrains),
+	}
+	if slices.Contains(allowedUsageApps, usageApp) {
+		return codersdk.UsageAppName(usageApp)
+	}
+
+	return ""
+}
diff --git a/cli/ssh_test.go b/cli/ssh_test.go
@@ -36,13 +36,15 @@ import (
 	"github.com/coder/coder/v2/agent"
 	"github.com/coder/coder/v2/agent/agentssh"
 	"github.com/coder/coder/v2/agent/agenttest"
+	agentproto "github.com/coder/coder/v2/agent/proto"
 	"github.com/coder/coder/v2/cli/clitest"
 	"github.com/coder/coder/v2/cli/cliui"
 	"github.com/coder/coder/v2/coderd/coderdtest"
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbfake"
 	"github.com/coder/coder/v2/coderd/database/dbtestutil"
 	"github.com/coder/coder/v2/coderd/rbac"
+	"github.com/coder/coder/v2/coderd/workspacestats/workspacestatstest"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/provisioner/echo"
 	"github.com/coder/coder/v2/provisionersdk/proto"
@@ -1292,6 +1294,108 @@ func TestSSH(t *testing.T) {
 		require.NoError(t, err)
 		require.Len(t, ents, 1, "expected one file in logdir %s", logDir)
 	})
+	t.Run("UpdateUsage", func(t *testing.T) {
+		t.Parallel()
+
+		type testCase struct {
+			name                   string
+			experiment             bool
+			usageAppName           string
+			expectedCalls          int
+			expectedCountSSH       int
+			expectedCountJetbrains int
+			expectedCountVscode    int
+		}
+		tcs := []testCase{
+			{
+				name: "NoExperiment",
+			},
+			{
+				name:             "Empty",
+				experiment:       true,
+				expectedCalls:    1,
+				expectedCountSSH: 1,
+			},
+			{
+				name:             "SSH",
+				experiment:       true,
+				usageAppName:     "ssh",
+				expectedCalls:    1,
+				expectedCountSSH: 1,
+			},
+			{
+				name:                   "Jetbrains",
+				experiment:             true,
+				usageAppName:           "jetbrains",
+				expectedCalls:          1,
+				expectedCountJetbrains: 1,
+			},
+			{
+				name:                "Vscode",
+				experiment:          true,
+				usageAppName:        "vscode",
+				expectedCalls:       1,
+				expectedCountVscode: 1,
+			},
+			{
+				name:         "Invalid",
+				experiment:   true,
+				usageAppName: "invalid",
+			},
+		}
+
+		for _, tc := range tcs {
+			tc := tc
+			t.Run(tc.name, func(t *testing.T) {
+				t.Parallel()
+
+				dv := coderdtest.DeploymentValues(t)
+				if tc.experiment {
+					dv.Experiments = []string{string(codersdk.ExperimentWorkspaceUsage)}
+				}
+				batcher := &workspacestatstest.StatsBatcher{
+					LastStats: &agentproto.Stats{},
+				}
+				admin, store := coderdtest.NewWithDatabase(t, &coderdtest.Options{
+					DeploymentValues: dv,
+					StatsBatcher:     batcher,
+				})
+				admin.SetLogger(slogtest.Make(t, nil).Named("client").Leveled(slog.LevelDebug))
+				first := coderdtest.CreateFirstUser(t, admin)
+				client, user := coderdtest.CreateAnotherUser(t, admin, first.OrganizationID)
+				r := dbfake.WorkspaceBuild(t, store, database.Workspace{
+					OrganizationID: first.OrganizationID,
+					OwnerID:        user.ID,
+				}).WithAgent().Do()
+				workspace := r.Workspace
+				agentToken := r.AgentToken
+				inv, root := clitest.New(t, "ssh", workspace.Name, fmt.Sprintf("--usage-app=%s", tc.usageAppName))
+				clitest.SetupConfig(t, client, root)
+				pty := ptytest.New(t).Attach(inv)
+
+				ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+				defer cancel()
+
+				cmdDone := tGo(t, func() {
+					err := inv.WithContext(ctx).Run()
+					assert.NoError(t, err)
+				})
+				pty.ExpectMatch("Waiting")
+
+				_ = agenttest.New(t, client.URL, agentToken)
+				coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
+
+				// Shells on Mac, Windows, and Linux all exit shells with the "exit" command.
+				pty.WriteLine("exit")
+				<-cmdDone
+
+				require.EqualValues(t, tc.expectedCalls, batcher.Called)
+				require.EqualValues(t, tc.expectedCountSSH, batcher.LastStats.SessionCountSsh)
+				require.EqualValues(t, tc.expectedCountJetbrains, batcher.LastStats.SessionCountJetbrains)
+				require.EqualValues(t, tc.expectedCountVscode, batcher.LastStats.SessionCountVscode)
+			})
+		}
+	})
 }
 
 //nolint:paralleltest // This test uses t.Setenv, parent test MUST NOT be parallel.

diff --git a/cli/testdata/coder_ssh_--help.golden b/cli/testdata/coder_ssh_--help.golden
@@ -42,6 +42,9 @@ OPTIONS:
       --stdio bool, $CODER_SSH_STDIO
           Specifies whether to emit SSH output over stdin/stdout.
 
+      --usage-app string, $CODER_SSH_USAGE_APP
+          Specifies the usage app to use for workspace activity tracking.
+
       --wait yes|no|auto, $CODER_SSH_WAIT (default: auto)
           Specifies whether or not to wait for the startup script to finish
           executing. Auto means that the agent startup script behavior

diff --git a/cli/vscodessh.go b/cli/vscodessh.go
@@ -110,7 +110,7 @@ func (r *RootCmd) vscodeSSH() *serpent.Command {
 			// will call this command after the workspace is started.
 			autostart := false
 
-			_, workspaceAgent, err := getWorkspaceAndAgent(ctx, inv, client, autostart, fmt.Sprintf("%s/%s", owner, name))
+			workspace, workspaceAgent, err := getWorkspaceAndAgent(ctx, inv, client, autostart, fmt.Sprintf("%s/%s", owner, name))
 			if err != nil {
 				return xerrors.Errorf("find workspace and agent: %w", err)
 			}
@@ -176,6 +176,13 @@ func (r *RootCmd) vscodeSSH() *serpent.Command {
 			defer agentConn.Close()
 
 			agentConn.AwaitReachable(ctx)
+
+			closeUsage := client.UpdateWorkspaceUsageWithBodyContext(ctx, workspace.ID, codersdk.PostWorkspaceUsageRequest{
+				AgentID: workspaceAgent.ID,
+				AppName: codersdk.UsageAppNameVscode,
+			})
+			defer closeUsage()
+
 			rawSSH, err := agentConn.SSH(ctx)
 			if err != nil {
 				return err

diff --git a/cli/vscodessh_test.go b/cli/vscodessh_test.go
@@ -9,9 +9,16 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"cdr.dev/slog"
+	"cdr.dev/slog/sloggers/slogtest"
+
 	"github.com/coder/coder/v2/agent/agenttest"
+	agentproto "github.com/coder/coder/v2/agent/proto"
 	"github.com/coder/coder/v2/cli/clitest"
 	"github.com/coder/coder/v2/coderd/coderdtest"
+	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/coderd/database/dbfake"
+	"github.com/coder/coder/v2/coderd/workspacestats/workspacestatstest"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/pty/ptytest"
 	"github.com/coder/coder/v2/testutil"
@@ -22,7 +29,25 @@ import (
 func TestVSCodeSSH(t *testing.T) {
 	t.Parallel()
 	ctx := testutil.Context(t, testutil.WaitLong)
-	client, workspace, agentToken := setupWorkspaceForAgent(t)
+	dv := coderdtest.DeploymentValues(t)
+	dv.Experiments = []string{string(codersdk.ExperimentWorkspaceUsage)}
+	batcher := &workspacestatstest.StatsBatcher{
+		LastStats: &agentproto.Stats{},
+	}
+	admin, store := coderdtest.NewWithDatabase(t, &coderdtest.Options{
+		DeploymentValues: dv,
+		StatsBatcher:     batcher,
+	})
+	admin.SetLogger(slogtest.Make(t, nil).Named("client").Leveled(slog.LevelDebug))
+	first := coderdtest.CreateFirstUser(t, admin)
+	client, user := coderdtest.CreateAnotherUser(t, admin, first.OrganizationID)
+	r := dbfake.WorkspaceBuild(t, store, database.Workspace{
+		OrganizationID: first.OrganizationID,
+		OwnerID:        user.ID,
+	}).WithAgent().Do()
+	workspace := r.Workspace
+	agentToken := r.AgentToken
+
 	user, err := client.User(ctx, codersdk.Me)
 	require.NoError(t, err)
 
@@ -65,4 +90,7 @@ func TestVSCodeSSH(t *testing.T) {
 	if err := waiter.Wait(); err != nil {
 		waiter.RequireIs(context.Canceled)
 	}
+
+	require.EqualValues(t, 1, batcher.Called)
+	require.EqualValues(t, 1, batcher.LastStats.SessionCountVscode)
 }
diff --git a/coderd/agentapi/api.go b/coderd/agentapi/api.go
@@ -25,6 +25,7 @@ import (
 	"github.com/coder/coder/v2/coderd/schedule"
 	"github.com/coder/coder/v2/coderd/tracing"
 	"github.com/coder/coder/v2/coderd/workspacestats"
+	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/codersdk/agentsdk"
 	"github.com/coder/coder/v2/tailnet"
 	tailnetproto "github.com/coder/coder/v2/tailnet/proto"
@@ -72,6 +73,7 @@ type Options struct {
 	DerpForceWebSockets       bool
 	DerpMapUpdateFrequency    time.Duration
 	ExternalAuthConfigs       []*externalauth.Config
+	Experiments               codersdk.Experiments
 
 	// Optional:
 	// WorkspaceID avoids a future lookup to find the workspace ID by setting
@@ -118,6 +120,7 @@ func New(opts Options) *API {
 		Log:                       opts.Log,
 		StatsReporter:             opts.StatsReporter,
 		AgentStatsRefreshInterval: opts.AgentStatsRefreshInterval,
+		Experiments:               opts.Experiments,
 	}
 
 	api.LifecycleAPI = &LifecycleAPI{

diff --git a/coderd/agentapi/stats.go b/coderd/agentapi/stats.go
@@ -12,6 +12,7 @@ import (
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbtime"
 	"github.com/coder/coder/v2/coderd/workspacestats"
+	"github.com/coder/coder/v2/codersdk"
 )
 
 type StatsAPI struct {
@@ -20,6 +21,7 @@ type StatsAPI struct {
 	Log                       slog.Logger
 	StatsReporter             *workspacestats.Reporter
 	AgentStatsRefreshInterval time.Duration
+	Experiments               codersdk.Experiments
 
 	TimeNowFn func() time.Time // defaults to dbtime.Now()
 }
@@ -55,6 +57,16 @@ func (a *StatsAPI) UpdateStats(ctx context.Context, req *agentproto.UpdateStatsR
 		slog.F("payload", req),
 	)
 
+	if a.Experiments.Enabled(codersdk.ExperimentWorkspaceUsage) {
+		// while the experiment is enabled we will not report
+		// session stats from the agent. This is because it is
+		// being handled by the CLI and the postWorkspaceUsage route.
+		req.Stats.SessionCountSsh = 0
+		req.Stats.SessionCountJetbrains = 0
+		req.Stats.SessionCountVscode = 0
+		req.Stats.SessionCountReconnectingPty = 0
+	}
+
 	err = a.StatsReporter.ReportAgentStats(
 		ctx,
 		a.now(),