Skip to content

chore: move ssh session tracking from agent to cli #13635

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 1 commit into from
Closed

chore: move ssh session tracking from agent to cli #13635

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 32 additions & 14 deletions agent/agentssh/agentssh.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,9 @@ const (
// MagicProcessCmdlineJetBrains is a string in a process's command line that
// uniquely identifies it as JetBrains software.
MagicProcessCmdlineJetBrains = "idea.vendor.name=JetBrains"
// MagicProcessCmdlineJetBrainsGateway is used to tell the agent not to report
// session stats because it's now being handled by the CLI.
MagicDisableUsageTrackingEnvironmentVariable = "CODER_SSH_DISABLE_USAGE_TRACKING"

// BlockedFileTransferErrorCode indicates that SSH server restricted the raw command from performing
// the file transfer.
Expand Down Expand Up @@ -380,26 +383,41 @@ func (s *Server) sessionStart(logger slog.Logger, session ssh.Session, extraEnv
env := append(session.Environ(), extraEnv...)
var magicType string
for index, kv := range env {
if !strings.HasPrefix(kv, MagicSessionTypeEnvironmentVariable) {
if !strings.HasPrefix(kv, MagicSessionTypeEnvironmentVariable+"=") {
continue
}
magicType = strings.ToLower(strings.TrimPrefix(kv, MagicSessionTypeEnvironmentVariable+"="))
env = append(env[:index], env[index+1:]...)
}
disableUsageTracking := false
for index, kv := range env {
if !strings.HasPrefix(kv, MagicDisableUsageTrackingEnvironmentVariable+"=") {
continue
}
if strings.ToLower(strings.TrimPrefix(kv, MagicDisableUsageTrackingEnvironmentVariable+"=")) == "true" {
disableUsageTracking = true
}
env = append(env[:index], env[index+1:]...)
}

// Always force lowercase checking to be case-insensitive.
switch magicType {
case MagicSessionTypeVSCode:
s.connCountVSCode.Add(1)
defer s.connCountVSCode.Add(-1)
case MagicSessionTypeJetBrains:
// Do nothing here because JetBrains launches hundreds of ssh sessions.
// We instead track JetBrains in the single persistent tcp forwarding channel.
case "":
s.connCountSSHSession.Add(1)
defer s.connCountSSHSession.Add(-1)
default:
logger.Warn(ctx, "invalid magic ssh session type specified", slog.F("type", magicType))
// We only want to track the session counts if they are from
// older clients that haven't migrating to reporting these stats
// from the CLI. This ensures we don't double count sessions.
if !disableUsageTracking {
// Always force lowercase checking to be case-insensitive.
switch magicType {
case MagicSessionTypeVSCode:
s.connCountVSCode.Add(1)
defer s.connCountVSCode.Add(-1)
case MagicSessionTypeJetBrains:
// Do nothing here because JetBrains launches hundreds of ssh sessions.
// We instead track JetBrains in the single persistent tcp forwarding channel.
case "":
s.connCountSSHSession.Add(1)
defer s.connCountSSHSession.Add(-1)
default:
logger.Warn(ctx, "invalid magic ssh session type specified", slog.F("type", magicType))
}
}

magicTypeLabel := magicTypeMetricLabel(magicType)
Expand Down
59 changes: 59 additions & 0 deletions cli/ssh.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
"os"
"os/exec"
"path/filepath"
"slices"
"strings"
"sync"
"time"
Expand All @@ -28,6 +29,8 @@

"cdr.dev/slog"
"cdr.dev/slog/sloggers/sloghuman"
"github.com/coder/coder/v2/agent/agentssh"
"github.com/coder/coder/v2/apiversion"
"github.com/coder/coder/v2/cli/cliui"
"github.com/coder/coder/v2/cli/cliutil"
"github.com/coder/coder/v2/coderd/autobuild/notify"
Expand Down Expand Up @@ -57,6 +60,7 @@
logDirPath string
remoteForwards []string
env []string
usageApp string
disableAutostart bool
)
client := new(codersdk.Client)
Expand Down Expand Up @@ -196,6 +200,11 @@
wait = false
}

experiments, err := client.Experiments(ctx)
if err != nil {
return err
}

templateVersion, err := client.TemplateVersion(ctx, workspace.LatestBuild.TemplateVersionID)
if err != nil {
return err
Expand Down Expand Up @@ -251,6 +260,18 @@
stopPolling := tryPollWorkspaceAutostop(ctx, client, workspace)
defer stopPolling()

usageAppName := getUsageAppName(usageApp, stdio, experiments, workspaceAgent.APIVersion)
if usageAppName != "" {
closeUsage := client.UpdateWorkspaceUsageWithBodyContext(ctx, workspace.ID, codersdk.PostWorkspaceUsageRequest{
AgentID: workspaceAgent.ID,
AppName: usageAppName,
})
defer closeUsage()

// signal to the agent that we are handling the usage tracking
parsedEnv = append(parsedEnv, [2]string{agentssh.MagicDisableUsageTrackingEnvironmentVariable, "true"})
}

if stdio {
rawSSH, err := conn.SSH(ctx)
if err != nil {
Expand Down Expand Up @@ -509,6 +530,12 @@
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
Expand Down Expand Up @@ -1044,3 +1071,35 @@
r.l.Error(context.Background(), "reading from stdin in stdio mode is not allowed")
return 0, io.EOF
}

func getUsageAppName(usageApp string, stdio bool, experiments codersdk.Experiments, agentAPIVersion string) codersdk.UsageAppName {

Check failure on line 1075 in cli/ssh.go

View workflow job for this annotation

GitHub Actions / lint 

flag-parameter: parameter 'stdio' seems to be a control flag, avoid control coupling (revive)
// if experiment is not enabled do not report usage
if !slices.Contains(experiments, codersdk.ExperimentWorkspaceUsage) {
return ""
}

// need agent version to be at or after 2.2
major, minor, err := apiversion.Parse(agentAPIVersion)
if err != nil {
return ""
}
err = apiversion.New(major, minor).Validate("2.2")
if err != nil {
return ""
}

// if usageApp is empty and not stdio, default to ssh
if usageApp == "" && !stdio {
usageApp = string(codersdk.UsageAppNameSSH)
}
allowedUsageApps := []string{
string(codersdk.UsageAppNameJetbrains),
string(codersdk.UsageAppNameVscode),
string(codersdk.UsageAppNameSSH),
}
if slices.Contains(allowedUsageApps, usageApp) {
return codersdk.UsageAppName(usageApp)
}

return ""
}
2 changes: 1 addition & 1 deletion tailnet/proto/version.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ import (

const (
CurrentMajor = 2
CurrentMinor = 1
CurrentMinor = 2
)

var CurrentVersion = apiversion.New(CurrentMajor, CurrentMinor).WithBackwardCompat(1)
Loading