From efab2498ea3966b21bef99a4ac18a90a1a4c7f95 Mon Sep 17 00:00:00 2001
From: Garrett Delfosse <garrett@coder.com>
Date: Tue, 18 Jun 2024 19:42:18 +0000
Subject: [PATCH 1/2] chore: report ssh stats by getting  session type from env

---
 cli/ssh.go | 35 +++++++++++++++++++++++++++++------
 1 file changed, 29 insertions(+), 6 deletions(-)

diff --git a/cli/ssh.go b/cli/ssh.go
index a55cd437ffe3d..38fde2c8a155d 100644
--- a/cli/ssh.go
+++ b/cli/ssh.go
@@ -28,6 +28,7 @@ import (
 
 	"cdr.dev/slog"
 	"cdr.dev/slog/sloggers/sloghuman"
+	"github.com/coder/coder/v2/agent/agentssh"
 	"github.com/coder/coder/v2/cli/cliui"
 	"github.com/coder/coder/v2/cli/cliutil"
 	"github.com/coder/coder/v2/coderd/autobuild/notify"
@@ -58,6 +59,7 @@ func (r *RootCmd) ssh() *serpent.Command {
 		remoteForwards   []string
 		env              []string
 		disableAutostart bool
+		noUsageTracking  bool
 	)
 	client := new(codersdk.Client)
 	cmd := &serpent.Command{
@@ -408,12 +410,27 @@ func (r *RootCmd) ssh() *serpent.Command {
 				return xerrors.Errorf("start shell: %w", err)
 			}
 
-			// track workspace usage while connection is open
-			closeUsage := client.UpdateWorkspaceUsageWithBodyContext(ctx, workspace.ID, codersdk.PostWorkspaceUsageRequest{
-				AgentID: workspaceAgent.ID,
-				AppName: codersdk.UsageAppNameSSH,
-			})
-			defer closeUsage()
+			if !noUsageTracking {
+				// check env for magic variable to determine the session type
+				appName := codersdk.UsageAppNameSSH
+				for _, kv := range parsedEnv {
+					if kv[0] == agentssh.MagicSessionTypeEnvironmentVariable {
+						switch kv[1] {
+						case agentssh.MagicSessionTypeJetBrains:
+							appName = codersdk.UsageAppNameJetbrains
+						case agentssh.MagicSessionTypeVSCode:
+							appName = codersdk.UsageAppNameVscode
+						}
+					}
+				}
+
+				// track workspace usage while connection is open
+				closeUsage := client.UpdateWorkspaceUsageWithBodyContext(ctx, workspace.ID, codersdk.PostWorkspaceUsageRequest{
+					AgentID: workspaceAgent.ID,
+					AppName: appName,
+				})
+				defer closeUsage()
+			}
 
 			// Put cancel at the top of the defer stack to initiate
 			// shutdown of services.
@@ -517,6 +534,12 @@ func (r *RootCmd) ssh() *serpent.Command {
 			Value:         serpent.StringArrayOf(&env),
 		},
 		sshDisableAutostartOption(serpent.BoolOf(&disableAutostart)),
+		{
+			Flag:        "no-usage-tracking",
+			Description: "Disables tracking of workspace usage.",
+			Env:         "CODER_SSH_NO_USAGE_TRACKING",
+			Value:       serpent.BoolOf(&noUsageTracking),
+		},
 	}
 	return cmd
 }

From b161219566dbb67824e0cd89021b336dccce4765 Mon Sep 17 00:00:00 2001
From: Garrett Delfosse <garrett@coder.com>
Date: Tue, 18 Jun 2024 20:13:05 +0000
Subject: [PATCH 2/2] lint

---
 cli/testdata/coder_ssh_--help.golden | 3 +++
 docs/cli/ssh.md                      | 9 +++++++++
 2 files changed, 12 insertions(+)

diff --git a/cli/testdata/coder_ssh_--help.golden b/cli/testdata/coder_ssh_--help.golden
index 80aaa3c204fda..153ca20911173 100644
--- a/cli/testdata/coder_ssh_--help.golden
+++ b/cli/testdata/coder_ssh_--help.golden
@@ -30,6 +30,9 @@ OPTIONS:
   -l, --log-dir string, $CODER_SSH_LOG_DIR
           Specify the directory containing SSH diagnostic log files.
 
+      --no-usage-tracking bool, $CODER_SSH_NO_USAGE_TRACKING
+          Disables tracking of workspace usage.
+
       --no-wait bool, $CODER_SSH_NO_WAIT
           Enter workspace immediately after the agent has connected. This is the
           default if the template has configured the agent startup script
diff --git a/docs/cli/ssh.md b/docs/cli/ssh.md
index d2110628feec4..6e6f040b9119c 100644
--- a/docs/cli/ssh.md
+++ b/docs/cli/ssh.md
@@ -113,3 +113,12 @@ Set environment variable(s) for session (key1=value1,key2=value2,...).
 | Default     | <code>false</code>                        |
 
 Disable starting the workspace automatically when connecting via SSH.
+
+### --no-usage-tracking
+
+|             |                                           |
+| ----------- | ----------------------------------------- |
+| Type        | <code>bool</code>                         |
+| Environment | <code>$CODER_SSH_NO_USAGE_TRACKING</code> |
+
+Disables tracking of workspace usage.