Thanks to visit codestin.com
Credit goes to github.com

Skip to content

feat: normalize workspace agent session counts into a child table #27952

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
EhabY wants to merge 2 commits into from
Closed

feat: normalize workspace agent session counts into a child table #27952

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
07b995b
feat: normalize workspace agent session counts into a child table
EhabY Aug 9, 2026
1d263c2
perf: optimize workspace agent session counts
EhabY Aug 13, 2026
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
3 changes: 3 additions & 0 deletions .github/workflows/typos.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,9 @@ AIBridge = "AI Bridge"

[default.extend-words]
AKS = "AKS"
# BRIN is a postgres index type
brin = "brin"
BRIN = "BRIN"
# do as sudo replacement
doas = "doas"
darcula = "darcula"
Expand Down
47 changes: 18 additions & 29 deletions agent/agent.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@ import (
"github.com/coder/coder/v2/buildinfo"
"github.com/coder/coder/v2/cli/gitauth"
"github.com/coder/coder/v2/coderd/database/dbtime"
"github.com/coder/coder/v2/coderd/idemetadata"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/codersdk/agentsdk"
"github.com/coder/coder/v2/codersdk/workspacesdk"
Expand Down Expand Up @@ -127,23 +128,14 @@ type Options struct {
}

type Client interface {
ConnectRPC29(ctx context.Context) (
proto.DRPCAgentClient29, tailnetproto.DRPCTailnetClient28, error,
ConnectRPC211(ctx context.Context) (
proto.DRPCAgentClient211, tailnetproto.DRPCTailnetClient28, error,
)
// ConnectRPC29WithRole is like ConnectRPC29 but sends an explicit
// ConnectRPC211WithRole is like ConnectRPC211 but sends an explicit
// role query parameter to the server. The workspace agent should
// use role "agent" to enable connection monitoring.
ConnectRPC29WithRole(ctx context.Context, role string) (
proto.DRPCAgentClient29, tailnetproto.DRPCTailnetClient28, error,
)
ConnectRPC210(ctx context.Context) (
proto.DRPCAgentClient210, tailnetproto.DRPCTailnetClient28, error,
)
// ConnectRPC210WithRole is like ConnectRPC210 but sends an explicit
// role query parameter to the server. The workspace agent should
// use role "agent" to enable connection monitoring.
ConnectRPC210WithRole(ctx context.Context, role string) (
proto.DRPCAgentClient210, tailnetproto.DRPCTailnetClient28, error,
ConnectRPC211WithRole(ctx context.Context, role string) (
proto.DRPCAgentClient211, tailnetproto.DRPCTailnetClient28, error,
)
tailnet.DERPMapRewriter
agentsdk.RefreshableSessionTokenProvider
Expand Down Expand Up @@ -426,17 +418,15 @@ func (a *agent) init() {
BlockLocalPortForwarding: a.blockLocalPortForwarding,
ReportConnection: func(id uuid.UUID, magicType agentssh.MagicSessionType, ip string) func(code int, reason string) {
var connectionType proto.Connection_Type
switch magicType {
case agentssh.MagicSessionTypeSSH:
// The enum cannot hold arbitrary types, so map by family.
switch idemetadata.Family(string(magicType)) {
case idemetadata.AppNameSSH:
connectionType = proto.Connection_SSH
case agentssh.MagicSessionTypeVSCode:
case idemetadata.AppNameVSCode:
connectionType = proto.Connection_VSCODE
case agentssh.MagicSessionTypeJetBrains:
case idemetadata.AppNameJetBrains:
connectionType = proto.Connection_JETBRAINS
case agentssh.MagicSessionTypeUnknown:
connectionType = proto.Connection_TYPE_UNSPECIFIED
default:
a.logger.Error(a.hardCtx, "unhandled magic session type when reporting connection", slog.F("magic_type", magicType))
connectionType = proto.Connection_TYPE_UNSPECIFIED
}

Expand Down Expand Up @@ -1176,7 +1166,7 @@ func (a *agent) run() (retErr error) {
// ConnectRPC returns the dRPC connection we use for the Agent and Tailnet v2+ APIs.
// We pass role "agent" to enable connection monitoring on the server, which tracks
// the agent's connectivity state (first_connected_at, last_connected_at, disconnected_at).
aAPI, tAPI, err := a.client.ConnectRPC210WithRole(a.hardCtx, "agent")
aAPI, tAPI, err := a.client.ConnectRPC211WithRole(a.hardCtx, "agent")
if err != nil {
return err
}
Expand Down Expand Up @@ -2162,13 +2152,12 @@ func (a *agent) Collect(ctx context.Context, networkStats map[netlogtype.Connect
stats.TxPackets += int64(counts.TxPackets)
}

// The count of active sessions.
sshStats := a.sshServer.ConnStats()
stats.SessionCountSsh = sshStats.Sessions
stats.SessionCountVscode = sshStats.VSCode
stats.SessionCountJetbrains = sshStats.JetBrains

stats.SessionCountReconnectingPty = a.reconnectingPTYServer.ConnCount()
// Active sessions per app; the deprecated fields stay zero. A client may
// label an ssh session "reconnecting_pty", so add, don't overwrite.
stats.SessionCounts = a.sshServer.SessionCounts()
if count := a.reconnectingPTYServer.ConnCount(); count > 0 {
stats.SessionCounts[idemetadata.AppNameReconnectingPTY] += count
}

// Compute the median connection latency!
a.logger.Debug(ctx, "starting peer latency measurement for stats")
Expand Down
20 changes: 10 additions & 10 deletions agent/agent_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -229,7 +229,7 @@ func assertSSHStats(t *testing.T, stats <-chan *proto.Stats) {
return false
}
t.Logf("got stats: ConnectionCount=%d, RxBytes=%d, TxBytes=%d, SessionCountSsh=%d",
s.ConnectionCount, s.RxBytes, s.TxBytes, s.SessionCountSsh)
s.ConnectionCount, s.RxBytes, s.TxBytes, s.SessionCounts["ssh"])
if s.ConnectionCount > 0 {
connectionCountSeen = true
}
Expand All @@ -239,7 +239,7 @@ func assertSSHStats(t *testing.T, stats <-chan *proto.Stats) {
if s.TxBytes > 0 {
txBytesSeen = true
}
if s.SessionCountSsh == 1 {
if s.SessionCounts["ssh"] == 1 {
sessionCountSSHSeen = true
}
return connectionCountSeen && rxBytesSeen && txBytesSeen && sessionCountSSHSeen
Expand Down Expand Up @@ -287,7 +287,7 @@ func TestAgent_Stats_ReconnectingPTY(t *testing.T) {
if s.TxBytes > 0 {
txBytesSeen = true
}
if s.SessionCountReconnectingPty == 1 {
if s.SessionCounts["reconnecting_pty"] == 1 {
sessionCountReconnectingPTYSeen = true
}
return connectionCountSeen && rxBytesSeen && txBytesSeen && sessionCountReconnectingPTYSeen
Expand Down Expand Up @@ -346,12 +346,12 @@ func TestAgent_Stats_Magic(t *testing.T) {
require.NoError(t, err)
require.Eventuallyf(t, func() bool {
s, ok := <-stats
t.Logf("got stats: ok=%t, ConnectionCount=%d, RxBytes=%d, TxBytes=%d, SessionCountVSCode=%d, ConnectionMedianLatencyMS=%f",
ok, s.ConnectionCount, s.RxBytes, s.TxBytes, s.SessionCountVscode, s.ConnectionMedianLatencyMs)
t.Logf("got stats: ok=%t, ConnectionCount=%d, RxBytes=%d, TxBytes=%d, SessionCounts[vscode]=%d, ConnectionMedianLatencyMS=%f",
ok, s.ConnectionCount, s.RxBytes, s.TxBytes, s.SessionCounts["vscode"], s.ConnectionMedianLatencyMs)
return ok &&
// Ensure that the connection didn't count as a "normal" SSH session.
// This was a special one, so it should be labeled specially in the stats!
s.SessionCountVscode == 1 &&
s.SessionCounts["vscode"] == 1 &&
// Ensure that connection latency is being counted!
// If it isn't, it's set to -1.
s.ConnectionMedianLatencyMs >= 0
Expand Down Expand Up @@ -417,8 +417,8 @@ func TestAgent_Stats_Magic(t *testing.T) {
require.Eventuallyf(t, func() bool {
s, ok := <-stats
t.Logf("got stats with conn open: ok=%t, ConnectionCount=%d, SessionCountJetBrains=%d",
ok, s.ConnectionCount, s.SessionCountJetbrains)
return ok && s.SessionCountJetbrains == 1
ok, s.ConnectionCount, s.SessionCounts["jetbrains"])
return ok && s.SessionCounts["jetbrains"] == 1
}, testutil.WaitLong, testutil.IntervalFast,
"never saw stats with conn open",
)
Expand All @@ -431,9 +431,9 @@ func TestAgent_Stats_Magic(t *testing.T) {
require.Eventuallyf(t, func() bool {
s, ok := <-stats
t.Logf("got stats after disconnect %t, %d",
ok, s.SessionCountJetbrains)
ok, s.SessionCounts["jetbrains"])
return ok &&
s.SessionCountJetbrains == 0
s.SessionCounts["jetbrains"] == 0
}, testutil.WaitLong, testutil.IntervalFast,
"never saw stats after conn closes",
)
Expand Down
105 changes: 58 additions & 47 deletions agent/agentssh/agentssh.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (
"errors"
"fmt"
"io"
"maps"
"net"
"os"
"os/exec"
Expand All @@ -31,6 +32,7 @@ import (
"github.com/coder/coder/v2/agent/agentexec"
"github.com/coder/coder/v2/agent/agentrsa"
"github.com/coder/coder/v2/agent/usershell"
"github.com/coder/coder/v2/coderd/idemetadata"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/pty"
)
Expand Down Expand Up @@ -71,17 +73,16 @@ const (
ContainerUserEnvironmentVariable = "CODER_CONTAINER_USER"
)

// MagicSessionType enums.
// Well-known magic session types, defined as canonical app names so the
// agent and server vocabularies cannot drift.
const (
// MagicSessionTypeUnknown means the session type could not be determined.
MagicSessionTypeUnknown MagicSessionType = "unknown"
// MagicSessionTypeSSH is the default session type.
MagicSessionTypeSSH MagicSessionType = "ssh"
MagicSessionTypeSSH MagicSessionType = idemetadata.AppNameSSH
// MagicSessionTypeVSCode is set in the SSH config by the VS Code extension to identify itself.
MagicSessionTypeVSCode MagicSessionType = "vscode"
MagicSessionTypeVSCode MagicSessionType = idemetadata.AppNameVSCode
// MagicSessionTypeJetBrains is set in the SSH config by the JetBrains
// extension to identify itself.
MagicSessionTypeJetBrains MagicSessionType = "jetbrains"
MagicSessionTypeJetBrains MagicSessionType = idemetadata.AppNameJetBrains
)

// BlockedFileTransferCommands contains a list of restricted file transfer commands.
Expand Down Expand Up @@ -148,17 +149,18 @@ type Server struct {
// a lock on mu but protected by closing.
wg sync.WaitGroup

// Active sessions per session type, zero-count entries removed. Kept off
// mu, which Close holds while closing sessions.
sessionCountsMu sync.RWMutex
sessionCounts map[string]int64

Execer agentexec.Execer
logger slog.Logger
srv *ssh.Server
x11Forwarder *x11Forwarder

config *Config

connCountVSCode atomic.Int64
connCountJetBrains atomic.Int64
connCountSSHSession atomic.Int64

metrics *sshServerMetrics
}

Expand Down Expand Up @@ -200,13 +202,14 @@ func NewServer(ctx context.Context, logger slog.Logger, prometheusRegistry *prom

metrics := newSSHServerMetrics(prometheusRegistry)
s := &Server{
Execer: execer,
listeners: make(map[net.Listener]struct{}),
fs: fs,
conns: make(map[net.Conn]struct{}),
sessions: make(map[ssh.Session]struct{}),
processes: make(map[*os.Process]struct{}),
logger: logger,
Execer: execer,
listeners: make(map[net.Listener]struct{}),
fs: fs,
conns: make(map[net.Conn]struct{}),
sessions: make(map[ssh.Session]struct{}),
processes: make(map[*os.Process]struct{}),
sessionCounts: make(map[string]int64),
logger: logger,

config: config,

Expand All @@ -232,7 +235,7 @@ func NewServer(ctx context.Context, logger slog.Logger, prometheusRegistry *prom
ChannelHandlers: map[string]ssh.ChannelHandler{
"direct-tcpip": func(srv *ssh.Server, conn *gossh.ServerConn, newChan gossh.NewChannel, ctx ssh.Context) {
// Wrapper is designed to find and track JetBrains Gateway connections.
wrapped := NewJetbrainsChannelWatcher(ctx, s.logger, s.config.ReportConnection, newChan, &s.connCountJetBrains)
wrapped := NewJetbrainsChannelWatcher(ctx, s.logger, s.config.ReportConnection, newChan, s.startSession)
ssh.DirectTCPIPHandler(srv, conn, wrapped, ctx)
},
"[email protected]": func(srv *ssh.Server, conn *gossh.ServerConn, newChan gossh.NewChannel, ctx ssh.Context) {
Expand Down Expand Up @@ -324,18 +327,30 @@ func NewServer(ctx context.Context, logger slog.Logger, prometheusRegistry *prom
return s, nil
}

type ConnStats struct {
Sessions int64
VSCode int64
JetBrains int64
// startSession increments the count for the given session type and returns a
// function that decrements it. Counters are created on demand; entries are
// bounded by concurrent sessions and capped again at ingestion.
func (s *Server) startSession(magicType MagicSessionType) (endSession func()) {
key := idemetadata.Normalize(string(magicType))
s.sessionCountsMu.Lock()
defer s.sessionCountsMu.Unlock()
s.sessionCounts[key]++
return func() {
s.sessionCountsMu.Lock()
defer s.sessionCountsMu.Unlock()
s.sessionCounts[key]--
if s.sessionCounts[key] <= 0 {
delete(s.sessionCounts, key)
}
}
}

func (s *Server) ConnStats() ConnStats {
return ConnStats{
Sessions: s.connCountSSHSession.Load(),
VSCode: s.connCountVSCode.Load(),
JetBrains: s.connCountJetBrains.Load(),
}
// SessionCounts returns a snapshot of active sessions per session type. Never
// nil, so callers can merge other session sources into it.
func (s *Server) SessionCounts() map[string]int64 {
s.sessionCountsMu.RLock()
defer s.sessionCountsMu.RUnlock()
return maps.Clone(s.sessionCounts)
}

func extractMagicSessionType(env []string) (magicType MagicSessionType, rawType string, filteredEnv []string) {
Expand All @@ -348,16 +363,11 @@ func extractMagicSessionType(env []string) (magicType MagicSessionType, rawType
// Keep going, we'll use the last instance of the env.
}

// Always force lowercase checking to be case-insensitive.
switch MagicSessionType(strings.ToLower(rawType)) {
case MagicSessionTypeVSCode:
magicType = MagicSessionTypeVSCode
case MagicSessionTypeJetBrains:
magicType = MagicSessionTypeJetBrains
case "", MagicSessionTypeSSH:
if rawType == "" {
magicType = MagicSessionTypeSSH
default:
magicType = MagicSessionTypeUnknown
} else {
// Canonicalize, don't classify: unknown names flow through.
magicType = MagicSessionType(idemetadata.Normalize(rawType))
}

return magicType, rawType, slices.DeleteFunc(env, func(kv string) bool {
Expand Down Expand Up @@ -424,6 +434,13 @@ func (s *Server) sessionHandler(session ssh.Session) {

env := session.Environ()
magicType, magicTypeRaw, env := extractMagicSessionType(env)
magicTypeFamily := idemetadata.Family(string(magicType))
if magicTypeFamily == idemetadata.AppNameUnknown {
logger.Debug(ctx, "unrecognized ssh session type",
slog.F("magic_type", magicType),
slog.F("raw_type", magicTypeRaw),
)
}

// It's not safe to assume RemoteAddr() returns a non-nil value. slog.F usage is fine because it correctly
// handles nil.
Expand All @@ -449,19 +466,13 @@ func (s *Server) sessionHandler(session ssh.Session) {

reportSession := true

switch magicType {
case MagicSessionTypeVSCode:
s.connCountVSCode.Add(1)
defer s.connCountVSCode.Add(-1)
case MagicSessionTypeJetBrains:
if magicTypeFamily == idemetadata.AppNameJetBrains {
// Do nothing here because JetBrains launches hundreds of ssh sessions.
// We instead track JetBrains in the single persistent tcp forwarding channel.
reportSession = false
case MagicSessionTypeSSH:
s.connCountSSHSession.Add(1)
defer s.connCountSSHSession.Add(-1)
case MagicSessionTypeUnknown:
logger.Warn(ctx, "invalid magic ssh session type specified", slog.F("raw_type", magicTypeRaw))
} else {
endSession := s.startSession(magicType)
defer endSession()
}

closeCause := func(_ string) {}
Expand Down
Loading
Loading