From ce51155ecf70dfde47e552fa70bb62ceb8e75546 Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Thu, 13 Aug 2026 18:18:48 +0300 Subject: [PATCH] feat: report arbitrary workspace agent session types --- agent/agent.go | 35 +++++++------------- agent/agent_test.go | 20 +++++------ agent/agenttest/client.go | 20 +++++++++++ cli/ssh.go | 13 +++----- cli/ssh_test.go | 70 ++++++++++++++++++++------------------- cli/vscodessh_test.go | 2 +- coderd/workspaces.go | 27 ++------------- coderd/workspaces_test.go | 6 ++-- codersdk/workspaces.go | 5 +++ 9 files changed, 93 insertions(+), 105 deletions(-) diff --git a/agent/agent.go b/agent/agent.go index c0ef8c7b571fe..a46e32dde126c 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -128,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 @@ -1175,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 } @@ -2161,14 +2152,12 @@ func (a *agent) Collect(ctx context.Context, networkStats map[netlogtype.Connect stats.TxPackets += int64(counts.TxPackets) } - // The count of active sessions. Only the four canonical names the protocol - // has fields for are reported; other app names are dropped, as before. - sessionCounts := a.sshServer.SessionCounts() - stats.SessionCountSsh = sessionCounts[idemetadata.AppNameSSH] - stats.SessionCountVscode = sessionCounts[idemetadata.AppNameVSCode] - stats.SessionCountJetbrains = sessionCounts[idemetadata.AppNameJetBrains] - - 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") diff --git a/agent/agent_test.go b/agent/agent_test.go index 908ce0ed3d8c6..a51c842168b47 100644 --- a/agent/agent_test.go +++ b/agent/agent_test.go @@ -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 } @@ -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 @@ -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 @@ -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 @@ -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", ) @@ -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", ) diff --git a/agent/agenttest/client.go b/agent/agenttest/client.go index 0f5d83a98f982..788653b0a36c4 100644 --- a/agent/agenttest/client.go +++ b/agent/agenttest/client.go @@ -182,6 +182,26 @@ func (c *Client) ConnectRPC210WithRole(ctx context.Context, _ string) ( return c.ConnectRPC210(ctx) } +func (c *Client) ConnectRPC211(ctx context.Context) ( + agentproto.DRPCAgentClient211, proto.DRPCTailnetClient28, error, +) { + aAPI, tAPI, err := c.ConnectRPC210(ctx) + if err != nil { + return nil, nil, err + } + client, ok := aAPI.(agentproto.DRPCAgentClient211) + if !ok { + return nil, nil, xerrors.Errorf("agenttest: connection does not implement DRPCAgentClient211; got %T", aAPI) + } + return client, tAPI, nil +} + +func (c *Client) ConnectRPC211WithRole(ctx context.Context, _ string) ( + agentproto.DRPCAgentClient211, proto.DRPCTailnetClient28, error, +) { + return c.ConnectRPC211(ctx) +} + func (c *Client) ConnectRPC29(ctx context.Context) ( agentproto.DRPCAgentClient29, proto.DRPCTailnetClient28, error, ) { diff --git a/cli/ssh.go b/cli/ssh.go index d18ac8909f575..007a5af2ea57c 100644 --- a/cli/ssh.go +++ b/cli/ssh.go @@ -1532,17 +1532,12 @@ func getUsageAppName(usageApp string) codersdk.UsageAppName { if usageApp == disableUsageApp { return "" } - - allowedUsageApps := []string{ - string(codersdk.UsageAppNameSSH), - string(codersdk.UsageAppNameVscode), - string(codersdk.UsageAppNameJetbrains), - } - if slices.Contains(allowedUsageApps, usageApp) { - return codersdk.UsageAppName(usageApp) + if usageApp == "" { + return codersdk.UsageAppNameSSH } - return codersdk.UsageAppNameSSH + // The server accepts arbitrary app names. + return codersdk.UsageAppName(usageApp) } func setStatsCallback( diff --git a/cli/ssh_test.go b/cli/ssh_test.go index 2221a23e7bf9c..b4b5d1387a5ba 100644 --- a/cli/ssh_test.go +++ b/cli/ssh_test.go @@ -1630,51 +1630,51 @@ func TestSSH(t *testing.T) { t.Parallel() type testCase struct { - name string - experiment bool - usageAppName string - expectedCalls int - expectedCountSSH int - expectedCountJetbrains int - expectedCountVscode int + name string + experiment bool + usageAppName string + expectedCalls int + expectedCounts map[string]int64 } tcs := []testCase{ { name: "NoExperiment", }, { - name: "Empty", - experiment: true, - expectedCalls: 1, - expectedCountSSH: 1, + name: "Empty", + experiment: true, + expectedCalls: 1, + expectedCounts: map[string]int64{"ssh": 1}, }, { - name: "SSH", - experiment: true, - usageAppName: "ssh", - expectedCalls: 1, - expectedCountSSH: 1, + name: "SSH", + experiment: true, + usageAppName: "ssh", + expectedCalls: 1, + expectedCounts: map[string]int64{"ssh": 1}, }, { - name: "Jetbrains", - experiment: true, - usageAppName: "jetbrains", - expectedCalls: 1, - expectedCountJetbrains: 1, + name: "Jetbrains", + experiment: true, + usageAppName: "jetbrains", + expectedCalls: 1, + expectedCounts: map[string]int64{"jetbrains": 1}, }, { - name: "Vscode", - experiment: true, - usageAppName: "vscode", - expectedCalls: 1, - expectedCountVscode: 1, + name: "Vscode", + experiment: true, + usageAppName: "vscode", + expectedCalls: 1, + expectedCounts: map[string]int64{"vscode": 1}, }, { - name: "InvalidDefaultsToSSH", - experiment: true, - usageAppName: "invalid", - expectedCalls: 1, - expectedCountSSH: 1, + // Arbitrary app names pass through raw (normalized at + // ingestion), so new IDEs need no CLI changes. + name: "ArbitraryNamePassthrough", + experiment: true, + usageAppName: "SomeFutureIDE", + expectedCalls: 1, + expectedCounts: map[string]int64{"SomeFutureIDE": 1}, }, { name: "Disable", @@ -1730,9 +1730,11 @@ func TestSSH(t *testing.T) { <-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) + if len(tc.expectedCounts) == 0 { + require.Empty(t, batcher.LastStats.GetSessionCounts()) + } else { + require.EqualValues(t, tc.expectedCounts, batcher.LastStats.GetSessionCounts()) + } }) } }) diff --git a/cli/vscodessh_test.go b/cli/vscodessh_test.go index 32afb52ca1da2..beb08e9cc841d 100644 --- a/cli/vscodessh_test.go +++ b/cli/vscodessh_test.go @@ -87,5 +87,5 @@ func TestVSCodeSSH(t *testing.T) { } require.EqualValues(t, 1, batcher.Called) - require.EqualValues(t, 1, batcher.LastStats.SessionCountVscode) + require.EqualValues(t, 1, batcher.LastStats.GetSessionCounts()["vscode"]) } diff --git a/coderd/workspaces.go b/coderd/workspaces.go index 21d11d88b7636..20f0efc83dbc4 100644 --- a/coderd/workspaces.go +++ b/coderd/workspaces.go @@ -1783,34 +1783,11 @@ func (api *API) postWorkspaceUsage(rw http.ResponseWriter, r *http.Request) { }) return } - if !slices.Contains(codersdk.AllowedAppNames, req.AppName) { - httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ - Message: "Invalid request", - Validations: []codersdk.ValidationError{{ - Field: "app_name", - Detail: fmt.Sprintf("must be one of %v", codersdk.AllowedAppNames), - }}, - }) - return - } + // Any app name is accepted, normalized at ingestion. stat := &proto.Stats{ ConnectionCount: 1, - } - switch req.AppName { - case codersdk.UsageAppNameVscode: - stat.SessionCountVscode = 1 - case codersdk.UsageAppNameJetbrains: - stat.SessionCountJetbrains = 1 - case codersdk.UsageAppNameReconnectingPty: - stat.SessionCountReconnectingPty = 1 - case codersdk.UsageAppNameSSH: - stat.SessionCountSsh = 1 - default: - // This means the app_name is in the codersdk.AllowedAppNames but not being - // handled by this switch statement. - httpapi.InternalServerError(rw, xerrors.Errorf("unknown app_name %q", req.AppName)) - return + SessionCounts: map[string]int64{string(req.AppName): 1}, } agent, err := api.Database.GetWorkspaceAgentByID(ctx, req.AgentID) diff --git a/coderd/workspaces_test.go b/coderd/workspaces_test.go index fdf33461af1c9..232c75fac0abd 100644 --- a/coderd/workspaces_test.go +++ b/coderd/workspaces_test.go @@ -5242,12 +5242,12 @@ func TestWorkspaceUsageTracking(t *testing.T) { AppName: "ssh", }) require.ErrorContains(t, err, "app_name") - // unknown app name fails + // unknown app names are accepted err = client.PostWorkspaceUsageWithBody(ctx, r.Workspace.ID, codersdk.PostWorkspaceUsageRequest{ AgentID: workspace.LatestBuild.Resources[0].Agents[0].ID, - AppName: "unknown", + AppName: "SomeFutureIDE", }) - require.ErrorContains(t, err, "app_name") + require.NoError(t, err) // vscode works err = client.PostWorkspaceUsageWithBody(ctx, r.Workspace.ID, codersdk.PostWorkspaceUsageRequest{ diff --git a/codersdk/workspaces.go b/codersdk/workspaces.go index 6a78ecd7b364d..c77fca10a3a06 100644 --- a/codersdk/workspaces.go +++ b/codersdk/workspaces.go @@ -383,6 +383,9 @@ type PostWorkspaceUsageRequest struct { type UsageAppName string +// Well-known usage app names. The API accepts any name, normalized at +// ingestion. These are wire format and cannot change: the hyphen in +// reconnecting-pty folds to the canonical reconnecting_pty. const ( UsageAppNameVscode UsageAppName = "vscode" UsageAppNameJetbrains UsageAppName = "jetbrains" @@ -390,6 +393,8 @@ const ( UsageAppNameSSH UsageAppName = "ssh" ) +// Deprecated: the workspace usage API no longer restricts app_name to this +// list. Use it only to recognize the well-known names. var AllowedAppNames = []UsageAppName{ UsageAppNameVscode, UsageAppNameJetbrains,