feat: add workspace agent session count protocol (#28125)

Phase 2 of the Scalable Approach for Adding New IDE Session Types RFC,
tracked in #27410.

Adds a session_counts map to Stats on the Agent API, so an agent can report
a count per app name instead of the four fixed fields, which are now marked
deprecated. Bumps the API to v2.11.

- coderd normalizes names on ingest, caps distinct entries, and sums the
  overflow under unknown, so an arbitrary client name cannot blow up
  cardinality.
- Agents predating v2.11 set the fixed fields instead, never both, and the
  ingest path falls back to them.

Nothing sends session_counts yet. The agent starts populating it in #28338.

.gitattributes

@@ -23,6 +23,7 @@ coderd/database/unique_constraint.go linguist-generated=true
 coderd/database/foreign_key_constraint.go linguist-generated=true
 coderd/database/check_constraint.go linguist-generated=true
 
+agent/proto/*.pb.go linguist-generated=true
 peerbroker/proto/*.go linguist-generated=true
 provisionerd/proto/*.go linguist-generated=true
 provisionerd/proto/version.go linguist-generated=false

agent/proto/agent.proto

@@ -151,17 +151,18 @@ message Stats {
 	int64 tx_bytes = 7;
 
 	// SessionCountVSCode is the number of connections received by an agent
-	// that are from our VS Code extension.
-	int64 session_count_vscode = 8;
+	// that are from our VS Code extension. Superseded by session_counts.
+	int64 session_count_vscode = 8 [deprecated = true];
 	// SessionCountJetBrains is the number of connections received by an agent
-	// that are from our JetBrains extension.
-	int64 session_count_jetbrains = 9;
+	// that are from our JetBrains extension. Superseded by session_counts.
+	int64 session_count_jetbrains = 9 [deprecated = true];
 	// SessionCountReconnectingPTY is the number of connections received by an agent
-	// that are from the reconnecting web terminal.
-	int64 session_count_reconnecting_pty = 10;
+	// that are from the reconnecting web terminal. Superseded by
+	// session_counts.
+	int64 session_count_reconnecting_pty = 10 [deprecated = true];
 	// SessionCountSSH is the number of connections received by an agent
-	// that are normal, non-tagged SSH sessions.
-	int64 session_count_ssh = 11;
+	// that are normal, non-tagged SSH sessions. Superseded by session_counts.
+	int64 session_count_ssh = 11 [deprecated = true];
 
 	message Metric {
 		string name = 1;
@@ -182,6 +183,10 @@ message Stats {
 		repeated Label labels = 4;
 	}
 	repeated Metric metrics = 12;
+
+	// SessionCounts is the number of active sessions per app name, such as
+	// "vscode" or "ssh". Coderd normalizes the names on ingest.
+	map<string, int64> session_counts = 13;
 }
 
 message UpdateStatsRequest{

agent/proto/agent_drpc_old.go

@@ -99,3 +99,9 @@ type DRPCAgentClient210 interface {
 	DRPCAgentClient29
 	PushContextState(ctx context.Context, in *PushContextStateRequest) (*PushContextStateResponse, error)
 }
+
+// DRPCAgentClient211 is the Agent API at v2.11. It adds the session_counts
+// map to Stats. No new RPCs.
+type DRPCAgentClient211 interface {
+	DRPCAgentClient210
+}

coderd/agentapi/stats.go

@@ -66,10 +66,7 @@ func (a *StatsAPI) UpdateStats(ctx context.Context, req *agentproto.UpdateStatsR
 		// 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
+		workspacestats.ClearSessionCounts(req.Stats)
 	}
 
 	err := a.StatsReporter.ReportAgentStats(

coderd/agentapi/stats_test.go

@@ -525,6 +525,7 @@ func TestUpdateStats(t *testing.T) {
 		batcher.Mu.Lock()
 		defer batcher.Mu.Unlock()
 		require.EqualValues(t, 1, batcher.Called)
+		require.Empty(t, batcher.LastStats.GetSessionCounts())
 		require.EqualValues(t, 0, batcher.LastStats.SessionCountSsh)
 		require.EqualValues(t, 0, batcher.LastStats.SessionCountJetbrains)
 		require.EqualValues(t, 0, batcher.LastStats.SessionCountVscode)

coderd/workspaceagents_test.go

@@ -3382,8 +3382,8 @@ func requireGetManifest(ctx context.Context, t testing.TB, aAPI agentproto.DRPCA
 	return manifest
 }
 
-func postStartup(ctx context.Context, t testing.TB, client agent.Client, startup *agentproto.Startup) error {
-	aAPI, _, err := client.ConnectRPC210(ctx)
+func postStartup(ctx context.Context, t testing.TB, client *agentsdk.Client, startup *agentproto.Startup) error {
+	aAPI, _, err := client.ConnectRPC211WithRole(ctx, "")
 	require.NoError(t, err)
 	defer func() {
 		cErr := aAPI.DRPCConn().Close()

coderd/workspacestats/reporter.go

@@ -154,10 +154,7 @@ func (r *Reporter) ReportAgentStats(ctx context.Context, now time.Time, workspac
 	}
 
 	// workspace activity: if no sessions we do not bump activity
-	if usage && stats.SessionCountVscode == 0 &&
-		stats.SessionCountJetbrains == 0 &&
-		stats.SessionCountReconnectingPty == 0 &&
-		stats.SessionCountSsh == 0 {
+	if usage && !HasSessionCounts(stats) {
 		return nil
 	}
 

coderd/workspacestats/sessioncounts.go

@@ -0,0 +1,87 @@
+package workspacestats
+
+import (
+	"cmp"
+	"maps"
+	"slices"
+
+	agentproto "github.com/coder/coder/v2/agent/proto"
+	"github.com/coder/coder/v2/codersdk"
+)
+
+// deprecatedSessionCounts returns the fixed session_count_* fields keyed by
+// normalized app name. Agents predating API v2.11 set these instead of
+// session_counts, never both.
+func deprecatedSessionCounts(st *agentproto.Stats) map[string]int64 {
+	return map[string]int64{
+		string(codersdk.AppFamilyVSCode):          st.SessionCountVscode,
+		string(codersdk.AppFamilyJetBrains):       st.SessionCountJetbrains,
+		string(codersdk.AppFamilyReconnectingPTY): st.SessionCountReconnectingPty,
+		string(codersdk.AppFamilySSH):             st.SessionCountSsh,
+	}
+}
+
+// normalizedSessionCounts returns per-app session counts, names normalized and
+// non-positive counts dropped.
+func normalizedSessionCounts(st *agentproto.Stats) map[string]int64 {
+	reported := st.GetSessionCounts()
+	if len(reported) == 0 {
+		reported = deprecatedSessionCounts(st)
+	}
+	counts := make(map[string]int64, len(reported))
+	for app, count := range reported {
+		if count > 0 {
+			counts[codersdk.NormalizeAppName(app)] += count
+		}
+	}
+	return counts
+}
+
+// maxSessionCountEntries bounds distinct app names per stats report. Overflow
+// aggregates under AppFamilyUnknown.
+const maxSessionCountEntries = 64
+
+// capSessionCounts keeps the busiest maxSessionCountEntries names, preferring
+// known apps, and sums the rest into AppFamilyUnknown, so the result can hold
+// one name past the cap.
+func capSessionCounts(counts map[string]int64) map[string]int64 {
+	if len(counts) <= maxSessionCountEntries {
+		return counts
+	}
+	// Known apps rank 0, unknown 1, so known apps win the cap.
+	rank := func(name string) int {
+		if codersdk.AppNameFamily(name) == codersdk.AppFamilyUnknown {
+			return 1
+		}
+		return 0
+	}
+	ranked := slices.SortedFunc(maps.Keys(counts), func(a, b string) int {
+		return cmp.Or(
+			cmp.Compare(rank(a), rank(b)),     // known apps first
+			cmp.Compare(counts[b], counts[a]), // then busiest
+			cmp.Compare(a, b),                 // then name, so the cut is stable
+		)
+	})
+	capped := make(map[string]int64, maxSessionCountEntries+1)
+	for _, name := range ranked[:maxSessionCountEntries] {
+		capped[name] = counts[name]
+	}
+	for _, name := range ranked[maxSessionCountEntries:] {
+		capped[string(codersdk.AppFamilyUnknown)] += counts[name]
+	}
+	return capped
+}
+
+// HasSessionCounts reports whether the stats contain any active session.
+func HasSessionCounts(st *agentproto.Stats) bool {
+	return len(normalizedSessionCounts(st)) > 0
+}
+
+// ClearSessionCounts zeroes every session count on the given stats.
+func ClearSessionCounts(st *agentproto.Stats) {
+	st.SessionCounts = nil
+	st.SessionCountVscode = 0
+	st.SessionCountJetbrains = 0
+	st.SessionCountReconnectingPty = 0
+	st.SessionCountSsh = 0
+}

coderd/workspacestats/sessioncounts_internal_test.go

@@ -0,0 +1,122 @@
+package workspacestats
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	agentproto "github.com/coder/coder/v2/agent/proto"
+)
+
+// sessionCountsFromProto is what the batcher applies on ingest.
+func sessionCountsFromProto(st *agentproto.Stats) map[string]int64 {
+	return capSessionCounts(normalizedSessionCounts(st))
+}
+
+func TestSessionCountsFromProto(t *testing.T) {
+	t.Parallel()
+
+	// The deprecated fields simulate agents older than API v2.11.
+	for _, tc := range []struct {
+		name  string
+		stats *agentproto.Stats
+		want  map[string]int64
+	}{
+		{
+			"Empty",
+			&agentproto.Stats{},
+			map[string]int64{},
+		},
+		{
+			"NormalizesAndMergesNames",
+			&agentproto.Stats{SessionCounts: map[string]int64{"VSCode": 1, "vscode": 2, "Reconnecting-PTY": 1}},
+			map[string]int64{"vscode": 3, "reconnecting_pty": 1},
+		},
+		{
+			"DropsNonPositiveEntries",
+			&agentproto.Stats{SessionCounts: map[string]int64{"vscode": 1, "reconnecting_pty": 0, "bogus": -1}},
+			map[string]int64{"vscode": 1},
+		},
+		{
+			"MapTakesPrecedenceOverDeprecatedFields",
+			&agentproto.Stats{SessionCounts: map[string]int64{"cursor": 2, "ssh": 1}, SessionCountVscode: 9},
+			map[string]int64{"cursor": 2, "ssh": 1},
+		},
+		{
+			"OldAgentFallback",
+			&agentproto.Stats{SessionCountVscode: 3, SessionCountJetbrains: 1, SessionCountSsh: 2},
+			map[string]int64{"vscode": 3, "jetbrains": 1, "ssh": 2},
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			require.Equal(t, tc.want, sessionCountsFromProto(tc.stats))
+		})
+	}
+}
+
+func TestCapSessionCounts(t *testing.T) {
+	t.Parallel()
+
+	// Known names carry the lowest counts and the busiest name sorts last, so
+	// only the ordering rules can keep them under the cap.
+	counts := map[string]int64{"ssh": 1, "vscode": 1, "zzz-busy-ide": 5}
+	const overflowing = 200
+	for i := range overflowing {
+		counts[fmt.Sprintf("0ide-%03d", i)] = 2
+	}
+
+	got := sessionCountsFromProto(&agentproto.Stats{SessionCounts: counts})
+
+	require.Len(t, got, maxSessionCountEntries+1, "cap plus the unknown bucket")
+	require.EqualValues(t, 1, got["ssh"])
+	require.EqualValues(t, 1, got["vscode"])
+	require.EqualValues(t, 5, got["zzz_busy_ide"])
+	// Those three plus 61 more fill the cap; the other 139 sum into unknown.
+	require.EqualValues(t, 2, got["0ide_000"])
+	require.NotContains(t, got, "0ide_199")
+	require.EqualValues(t, (overflowing-61)*2, got["unknown"])
+
+	var reported, stored int64
+	for _, count := range counts {
+		reported += count
+	}
+	for _, count := range got {
+		stored += count
+	}
+	require.Equal(t, reported, stored, "capping must not lose counts")
+}
+
+func TestHasSessionCounts(t *testing.T) {
+	t.Parallel()
+
+	for _, tc := range []struct {
+		name  string
+		stats *agentproto.Stats
+		want  bool
+	}{
+		{"Empty", &agentproto.Stats{}, false},
+		{"MapPositive", &agentproto.Stats{SessionCounts: map[string]int64{"cursor": 1}}, true},
+		{"MapZeroOnly", &agentproto.Stats{SessionCounts: map[string]int64{"ssh": 0}}, false},
+		{"DeprecatedFieldPositive", &agentproto.Stats{SessionCountSsh: 2}, true},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			require.Equal(t, tc.want, HasSessionCounts(tc.stats))
+		})
+	}
+}
+
+func TestClearSessionCounts(t *testing.T) {
+	t.Parallel()
+
+	st := &agentproto.Stats{
+		SessionCounts:      map[string]int64{"vscode": 1},
+		SessionCountVscode: 1,
+		SessionCountSsh:    2,
+	}
+	ClearSessionCounts(st)
+	require.Empty(t, sessionCountsFromProto(st))
+	require.False(t, HasSessionCounts(st))
+}

codersdk/agentsdk/agentsdk.go

@@ -362,6 +362,19 @@ func (c *Client) ConnectRPC210WithRole(ctx context.Context, role string) (
 	return proto.NewDRPCAgentClient(conn), tailnetproto.NewDRPCTailnetClient(conn), nil
 }
 
+// ConnectRPC211WithRole returns a dRPC client to the Agent API v2.11, which
+// reports per-app session counts on Stats. Pass role "agent" for workspace
+// agents to enable connection monitoring.
+func (c *Client) ConnectRPC211WithRole(ctx context.Context, role string) (
+	proto.DRPCAgentClient211, tailnetproto.DRPCTailnetClient28, error,
+) {
+	conn, err := c.connectRPCVersion(ctx, apiversion.New(2, 11), role)
+	if err != nil {
+		return nil, nil, err
+	}
+	return proto.NewDRPCAgentClient(conn), tailnetproto.NewDRPCTailnetClient(conn), nil
+}
+
 // ConnectRPC connects to the workspace agent API and tailnet API.
 // It does not send a role query parameter, so the server will apply
 // its default behavior (currently: enable connection monitoring for

tailnet/proto/version.go

@@ -81,9 +81,13 @@ import (
 //     deployments remain interoperable. Real persistence,
 //     KindMCPServer provider, and chatd hydration land in
 //     CODAGT-569.
+//
+// API v2.11:
+//   - Added the session_counts map to Stats on the Agent API, deprecating
+//     the fixed session_count_* fields.
 const (
 	CurrentMajor = 2
-	CurrentMinor = 10
+	CurrentMinor = 11
 )
 
 var CurrentVersion = apiversion.New(CurrentMajor, CurrentMinor)