diff --git a/cli/ssh.go b/cli/ssh.go index ac849649b9184..e4e9fadf5e8e8 100644 --- a/cli/ssh.go +++ b/cli/ssh.go @@ -12,6 +12,7 @@ import ( "os" "os/exec" "path/filepath" + "slices" "strings" "sync" "time" @@ -40,6 +41,10 @@ import ( "github.com/coder/serpent" ) +const ( + disableUsageApp = "disable" +) + var ( workspacePollInterval = time.Minute autostopNotifyCountdown = []time.Duration{30 * time.Minute} @@ -57,6 +62,7 @@ func (r *RootCmd) ssh() *serpent.Command { logDirPath string remoteForwards []string env []string + usageApp string disableAutostart bool ) client := new(codersdk.Client) @@ -251,6 +257,15 @@ func (r *RootCmd) ssh() *serpent.Command { stopPolling := tryPollWorkspaceAutostop(ctx, client, workspace) defer stopPolling() + usageAppName := getUsageAppName(usageApp) + if usageAppName != "" { + closeUsage := client.UpdateWorkspaceUsageWithBodyContext(ctx, workspace.ID, codersdk.PostWorkspaceUsageRequest{ + AgentID: workspaceAgent.ID, + AppName: usageAppName, + }) + defer closeUsage() + } + if stdio { rawSSH, err := conn.SSH(ctx) if err != nil { @@ -509,6 +524,13 @@ func (r *RootCmd) ssh() *serpent.Command { 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), + Hidden: true, + }, sshDisableAutostartOption(serpent.BoolOf(&disableAutostart)), } return cmd @@ -1044,3 +1066,20 @@ func (r stdioErrLogReader) Read(_ []byte) (int, error) { r.l.Error(context.Background(), "reading from stdin in stdio mode is not allowed") return 0, io.EOF } + +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) + } + + return codersdk.UsageAppNameSSH +} diff --git a/cli/ssh_test.go b/cli/ssh_test.go index 8c3c1a4e40fd1..ae93c4b0cea05 100644 --- a/cli/ssh_test.go +++ b/cli/ssh_test.go @@ -36,6 +36,7 @@ import ( "github.com/coder/coder/v2/agent" "github.com/coder/coder/v2/agent/agentssh" "github.com/coder/coder/v2/agent/agenttest" + agentproto "github.com/coder/coder/v2/agent/proto" "github.com/coder/coder/v2/cli/clitest" "github.com/coder/coder/v2/cli/cliui" "github.com/coder/coder/v2/coderd/coderdtest" @@ -43,6 +44,7 @@ import ( "github.com/coder/coder/v2/coderd/database/dbfake" "github.com/coder/coder/v2/coderd/database/dbtestutil" "github.com/coder/coder/v2/coderd/rbac" + "github.com/coder/coder/v2/coderd/workspacestats/workspacestatstest" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/provisioner/echo" "github.com/coder/coder/v2/provisionersdk/proto" @@ -1292,6 +1294,115 @@ func TestSSH(t *testing.T) { require.NoError(t, err) require.Len(t, ents, 1, "expected one file in logdir %s", logDir) }) + t.Run("UpdateUsage", func(t *testing.T) { + t.Parallel() + + type testCase struct { + name string + experiment bool + usageAppName string + expectedCalls int + expectedCountSSH int + expectedCountJetbrains int + expectedCountVscode int + } + tcs := []testCase{ + { + name: "NoExperiment", + }, + { + name: "Empty", + experiment: true, + expectedCalls: 1, + expectedCountSSH: 1, + }, + { + name: "SSH", + experiment: true, + usageAppName: "ssh", + expectedCalls: 1, + expectedCountSSH: 1, + }, + { + name: "Jetbrains", + experiment: true, + usageAppName: "jetbrains", + expectedCalls: 1, + expectedCountJetbrains: 1, + }, + { + name: "Vscode", + experiment: true, + usageAppName: "vscode", + expectedCalls: 1, + expectedCountVscode: 1, + }, + { + name: "InvalidDefaultsToSSH", + experiment: true, + usageAppName: "invalid", + expectedCalls: 1, + expectedCountSSH: 1, + }, + { + name: "Disable", + experiment: true, + usageAppName: "disable", + }, + } + + for _, tc := range tcs { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + dv := coderdtest.DeploymentValues(t) + if tc.experiment { + dv.Experiments = []string{string(codersdk.ExperimentWorkspaceUsage)} + } + batcher := &workspacestatstest.StatsBatcher{ + LastStats: &agentproto.Stats{}, + } + admin, store := coderdtest.NewWithDatabase(t, &coderdtest.Options{ + DeploymentValues: dv, + StatsBatcher: batcher, + }) + admin.SetLogger(slogtest.Make(t, nil).Named("client").Leveled(slog.LevelDebug)) + first := coderdtest.CreateFirstUser(t, admin) + client, user := coderdtest.CreateAnotherUser(t, admin, first.OrganizationID) + r := dbfake.WorkspaceBuild(t, store, database.Workspace{ + OrganizationID: first.OrganizationID, + OwnerID: user.ID, + }).WithAgent().Do() + workspace := r.Workspace + agentToken := r.AgentToken + inv, root := clitest.New(t, "ssh", workspace.Name, fmt.Sprintf("--usage-app=%s", tc.usageAppName)) + clitest.SetupConfig(t, client, root) + pty := ptytest.New(t).Attach(inv) + + ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) + defer cancel() + + cmdDone := tGo(t, func() { + err := inv.WithContext(ctx).Run() + assert.NoError(t, err) + }) + pty.ExpectMatch("Waiting") + + _ = agenttest.New(t, client.URL, agentToken) + coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) + + // Shells on Mac, Windows, and Linux all exit shells with the "exit" command. + pty.WriteLine("exit") + <-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) + }) + } + }) } //nolint:paralleltest // This test uses t.Setenv, parent test MUST NOT be parallel. diff --git a/cli/vscodessh.go b/cli/vscodessh.go index 147436374b1f6..558b50c00fe95 100644 --- a/cli/vscodessh.go +++ b/cli/vscodessh.go @@ -110,7 +110,7 @@ func (r *RootCmd) vscodeSSH() *serpent.Command { // will call this command after the workspace is started. autostart := false - _, workspaceAgent, err := getWorkspaceAndAgent(ctx, inv, client, autostart, fmt.Sprintf("%s/%s", owner, name)) + workspace, workspaceAgent, err := getWorkspaceAndAgent(ctx, inv, client, autostart, fmt.Sprintf("%s/%s", owner, name)) if err != nil { return xerrors.Errorf("find workspace and agent: %w", err) } @@ -176,6 +176,13 @@ func (r *RootCmd) vscodeSSH() *serpent.Command { defer agentConn.Close() agentConn.AwaitReachable(ctx) + + closeUsage := client.UpdateWorkspaceUsageWithBodyContext(ctx, workspace.ID, codersdk.PostWorkspaceUsageRequest{ + AgentID: workspaceAgent.ID, + AppName: codersdk.UsageAppNameVscode, + }) + defer closeUsage() + rawSSH, err := agentConn.SSH(ctx) if err != nil { return err diff --git a/cli/vscodessh_test.go b/cli/vscodessh_test.go index a4f6ca19132c6..f80b6b0b6029e 100644 --- a/cli/vscodessh_test.go +++ b/cli/vscodessh_test.go @@ -9,9 +9,16 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "cdr.dev/slog" + "cdr.dev/slog/sloggers/slogtest" + "github.com/coder/coder/v2/agent/agenttest" + agentproto "github.com/coder/coder/v2/agent/proto" "github.com/coder/coder/v2/cli/clitest" "github.com/coder/coder/v2/coderd/coderdtest" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbfake" + "github.com/coder/coder/v2/coderd/workspacestats/workspacestatstest" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/pty/ptytest" "github.com/coder/coder/v2/testutil" @@ -22,7 +29,25 @@ import ( func TestVSCodeSSH(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) - client, workspace, agentToken := setupWorkspaceForAgent(t) + dv := coderdtest.DeploymentValues(t) + dv.Experiments = []string{string(codersdk.ExperimentWorkspaceUsage)} + batcher := &workspacestatstest.StatsBatcher{ + LastStats: &agentproto.Stats{}, + } + admin, store := coderdtest.NewWithDatabase(t, &coderdtest.Options{ + DeploymentValues: dv, + StatsBatcher: batcher, + }) + admin.SetLogger(slogtest.Make(t, nil).Named("client").Leveled(slog.LevelDebug)) + first := coderdtest.CreateFirstUser(t, admin) + client, user := coderdtest.CreateAnotherUser(t, admin, first.OrganizationID) + r := dbfake.WorkspaceBuild(t, store, database.Workspace{ + OrganizationID: first.OrganizationID, + OwnerID: user.ID, + }).WithAgent().Do() + workspace := r.Workspace + agentToken := r.AgentToken + user, err := client.User(ctx, codersdk.Me) require.NoError(t, err) @@ -65,4 +90,7 @@ func TestVSCodeSSH(t *testing.T) { if err := waiter.Wait(); err != nil { waiter.RequireIs(context.Canceled) } + + require.EqualValues(t, 1, batcher.Called) + require.EqualValues(t, 1, batcher.LastStats.SessionCountVscode) } diff --git a/coderd/agentapi/api.go b/coderd/agentapi/api.go index ae0d594314e66..4e5e30ad9c761 100644 --- a/coderd/agentapi/api.go +++ b/coderd/agentapi/api.go @@ -25,6 +25,7 @@ import ( "github.com/coder/coder/v2/coderd/schedule" "github.com/coder/coder/v2/coderd/tracing" "github.com/coder/coder/v2/coderd/workspacestats" + "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/codersdk/agentsdk" "github.com/coder/coder/v2/tailnet" tailnetproto "github.com/coder/coder/v2/tailnet/proto" @@ -72,6 +73,7 @@ type Options struct { DerpForceWebSockets bool DerpMapUpdateFrequency time.Duration ExternalAuthConfigs []*externalauth.Config + Experiments codersdk.Experiments // Optional: // WorkspaceID avoids a future lookup to find the workspace ID by setting @@ -118,6 +120,7 @@ func New(opts Options) *API { Log: opts.Log, StatsReporter: opts.StatsReporter, AgentStatsRefreshInterval: opts.AgentStatsRefreshInterval, + Experiments: opts.Experiments, } api.LifecycleAPI = &LifecycleAPI{ diff --git a/coderd/agentapi/stats.go b/coderd/agentapi/stats.go index a167fb5d6f275..4f6a6da1c8c66 100644 --- a/coderd/agentapi/stats.go +++ b/coderd/agentapi/stats.go @@ -12,6 +12,7 @@ import ( "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbtime" "github.com/coder/coder/v2/coderd/workspacestats" + "github.com/coder/coder/v2/codersdk" ) type StatsAPI struct { @@ -20,6 +21,7 @@ type StatsAPI struct { Log slog.Logger StatsReporter *workspacestats.Reporter AgentStatsRefreshInterval time.Duration + Experiments codersdk.Experiments TimeNowFn func() time.Time // defaults to dbtime.Now() } @@ -55,6 +57,16 @@ func (a *StatsAPI) UpdateStats(ctx context.Context, req *agentproto.UpdateStatsR slog.F("payload", req), ) + if a.Experiments.Enabled(codersdk.ExperimentWorkspaceUsage) { + // 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 + } + err = a.StatsReporter.ReportAgentStats( ctx, a.now(), diff --git a/coderd/agentapi/stats_test.go b/coderd/agentapi/stats_test.go index 8b4d72fc1d579..57534208be110 100644 --- a/coderd/agentapi/stats_test.go +++ b/coderd/agentapi/stats_test.go @@ -3,7 +3,6 @@ package agentapi_test import ( "context" "database/sql" - "sync" "sync/atomic" "testing" "time" @@ -23,37 +22,11 @@ import ( "github.com/coder/coder/v2/coderd/prometheusmetrics" "github.com/coder/coder/v2/coderd/schedule" "github.com/coder/coder/v2/coderd/workspacestats" + "github.com/coder/coder/v2/coderd/workspacestats/workspacestatstest" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/testutil" ) -type statsBatcher struct { - mu sync.Mutex - - called int64 - lastTime time.Time - lastAgentID uuid.UUID - lastTemplateID uuid.UUID - lastUserID uuid.UUID - lastWorkspaceID uuid.UUID - lastStats *agentproto.Stats -} - -var _ workspacestats.Batcher = &statsBatcher{} - -func (b *statsBatcher) Add(now time.Time, agentID uuid.UUID, templateID uuid.UUID, userID uuid.UUID, workspaceID uuid.UUID, st *agentproto.Stats) error { - b.mu.Lock() - defer b.mu.Unlock() - b.called++ - b.lastTime = now - b.lastAgentID = agentID - b.lastTemplateID = templateID - b.lastUserID = userID - b.lastWorkspaceID = workspaceID - b.lastStats = st - return nil -} - func TestUpdateStates(t *testing.T) { t.Parallel() @@ -94,7 +67,7 @@ func TestUpdateStates(t *testing.T) { panic("not implemented") }, } - batcher = &statsBatcher{} + batcher = &workspacestatstest.StatsBatcher{} updateAgentMetricsFnCalled = false req = &agentproto.UpdateStatsRequest{ @@ -188,15 +161,15 @@ func TestUpdateStates(t *testing.T) { ReportInterval: durationpb.New(10 * time.Second), }, resp) - batcher.mu.Lock() - defer batcher.mu.Unlock() - require.Equal(t, int64(1), batcher.called) - require.Equal(t, now, batcher.lastTime) - require.Equal(t, agent.ID, batcher.lastAgentID) - require.Equal(t, template.ID, batcher.lastTemplateID) - require.Equal(t, user.ID, batcher.lastUserID) - require.Equal(t, workspace.ID, batcher.lastWorkspaceID) - require.Equal(t, req.Stats, batcher.lastStats) + batcher.Mu.Lock() + defer batcher.Mu.Unlock() + require.Equal(t, int64(1), batcher.Called) + require.Equal(t, now, batcher.LastTime) + require.Equal(t, agent.ID, batcher.LastAgentID) + require.Equal(t, template.ID, batcher.LastTemplateID) + require.Equal(t, user.ID, batcher.LastUserID) + require.Equal(t, workspace.ID, batcher.LastWorkspaceID) + require.Equal(t, req.Stats, batcher.LastStats) ctx := testutil.Context(t, testutil.WaitShort) select { case <-ctx.Done(): @@ -222,7 +195,7 @@ func TestUpdateStates(t *testing.T) { panic("not implemented") }, } - batcher = &statsBatcher{} + batcher = &workspacestatstest.StatsBatcher{} req = &agentproto.UpdateStatsRequest{ Stats: &agentproto.Stats{ @@ -336,7 +309,7 @@ func TestUpdateStates(t *testing.T) { panic("not implemented") }, } - batcher = &statsBatcher{} + batcher = &workspacestatstest.StatsBatcher{} updateAgentMetricsFnCalled = false req = &agentproto.UpdateStatsRequest{ @@ -406,6 +379,138 @@ func TestUpdateStates(t *testing.T) { require.True(t, updateAgentMetricsFnCalled) }) + + t.Run("WorkspaceUsageExperiment", func(t *testing.T) { + t.Parallel() + + var ( + now = dbtime.Now() + dbM = dbmock.NewMockStore(gomock.NewController(t)) + ps = pubsub.NewInMemory() + + templateScheduleStore = schedule.MockTemplateScheduleStore{ + GetFn: func(context.Context, database.Store, uuid.UUID) (schedule.TemplateScheduleOptions, error) { + t.Fatal("getfn should not be called") + return schedule.TemplateScheduleOptions{}, nil + }, + SetFn: func(context.Context, database.Store, database.Template, schedule.TemplateScheduleOptions) (database.Template, error) { + t.Fatal("setfn not implemented") + return database.Template{}, nil + }, + } + batcher = &workspacestatstest.StatsBatcher{} + updateAgentMetricsFnCalled = false + + req = &agentproto.UpdateStatsRequest{ + Stats: &agentproto.Stats{ + ConnectionsByProto: map[string]int64{ + "tcp": 1, + "dean": 2, + }, + ConnectionCount: 3, + ConnectionMedianLatencyMs: 23, + RxPackets: 120, + RxBytes: 1000, + TxPackets: 130, + TxBytes: 2000, + SessionCountVscode: 1, + SessionCountJetbrains: 2, + SessionCountReconnectingPty: 3, + SessionCountSsh: 4, + Metrics: []*agentproto.Stats_Metric{ + { + Name: "awesome metric", + Value: 42, + }, + { + Name: "uncool metric", + Value: 0, + }, + }, + }, + } + ) + api := agentapi.StatsAPI{ + AgentFn: func(context.Context) (database.WorkspaceAgent, error) { + return agent, nil + }, + Database: dbM, + StatsReporter: workspacestats.NewReporter(workspacestats.ReporterOptions{ + Database: dbM, + Pubsub: ps, + StatsBatcher: batcher, + TemplateScheduleStore: templateScheduleStorePtr(templateScheduleStore), + UpdateAgentMetricsFn: func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric) { + updateAgentMetricsFnCalled = true + assert.Equal(t, prometheusmetrics.AgentMetricLabels{ + Username: user.Username, + WorkspaceName: workspace.Name, + AgentName: agent.Name, + TemplateName: template.Name, + }, labels) + assert.Equal(t, req.Stats.Metrics, metrics) + }, + }), + AgentStatsRefreshInterval: 10 * time.Second, + TimeNowFn: func() time.Time { + return now + }, + Experiments: codersdk.Experiments{ + codersdk.ExperimentWorkspaceUsage, + }, + } + + // Workspace gets fetched. + dbM.EXPECT().GetWorkspaceByAgentID(gomock.Any(), agent.ID).Return(database.GetWorkspaceByAgentIDRow{ + Workspace: workspace, + TemplateName: template.Name, + }, nil) + + // We expect an activity bump because ConnectionCount > 0. + dbM.EXPECT().ActivityBumpWorkspace(gomock.Any(), database.ActivityBumpWorkspaceParams{ + WorkspaceID: workspace.ID, + NextAutostart: time.Time{}.UTC(), + }).Return(nil) + + // Workspace last used at gets bumped. + dbM.EXPECT().UpdateWorkspaceLastUsedAt(gomock.Any(), database.UpdateWorkspaceLastUsedAtParams{ + ID: workspace.ID, + LastUsedAt: now, + }).Return(nil) + + // User gets fetched to hit the UpdateAgentMetricsFn. + dbM.EXPECT().GetUserByID(gomock.Any(), user.ID).Return(user, nil) + + // Ensure that pubsub notifications are sent. + notifyDescription := make(chan []byte) + ps.Subscribe(codersdk.WorkspaceNotifyChannel(workspace.ID), func(_ context.Context, description []byte) { + go func() { + notifyDescription <- description + }() + }) + + resp, err := api.UpdateStats(context.Background(), req) + require.NoError(t, err) + require.Equal(t, &agentproto.UpdateStatsResponse{ + ReportInterval: durationpb.New(10 * time.Second), + }, resp) + + batcher.Mu.Lock() + defer batcher.Mu.Unlock() + require.EqualValues(t, 1, batcher.Called) + require.EqualValues(t, 0, batcher.LastStats.SessionCountSsh) + require.EqualValues(t, 0, batcher.LastStats.SessionCountJetbrains) + require.EqualValues(t, 0, batcher.LastStats.SessionCountVscode) + require.EqualValues(t, 0, batcher.LastStats.SessionCountReconnectingPty) + ctx := testutil.Context(t, testutil.WaitShort) + select { + case <-ctx.Done(): + t.Error("timed out while waiting for pubsub notification") + case description := <-notifyDescription: + require.Equal(t, description, []byte{}) + } + require.True(t, updateAgentMetricsFnCalled) + }) } func templateScheduleStorePtr(store schedule.TemplateScheduleStore) *atomic.Pointer[schedule.TemplateScheduleStore] { diff --git a/coderd/coderd.go b/coderd/coderd.go index cca4faf36a203..f92e3008604c3 100644 --- a/coderd/coderd.go +++ b/coderd/coderd.go @@ -187,7 +187,7 @@ type Options struct { HTTPClient *http.Client UpdateAgentMetrics func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric) - StatsBatcher *workspacestats.DBBatcher + StatsBatcher workspacestats.Batcher WorkspaceAppsStatsCollectorOptions workspaceapps.StatsCollectorOptions diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go index 49388aa3537a5..fbeaed43bd8be 100644 --- a/coderd/coderdtest/coderdtest.go +++ b/coderd/coderdtest/coderdtest.go @@ -145,7 +145,7 @@ type Options struct { // Logger should only be overridden if you expect errors // as part of your test. Logger *slog.Logger - StatsBatcher *workspacestats.DBBatcher + StatsBatcher workspacestats.Batcher WorkspaceAppsStatsCollectorOptions workspaceapps.StatsCollectorOptions AllowWorkspaceRenames bool diff --git a/coderd/workspaceagentsrpc.go b/coderd/workspaceagentsrpc.go index ec8dcd8a0e3fc..b413db264feac 100644 --- a/coderd/workspaceagentsrpc.go +++ b/coderd/workspaceagentsrpc.go @@ -143,6 +143,7 @@ func (api *API) workspaceAgentRPC(rw http.ResponseWriter, r *http.Request) { DerpForceWebSockets: api.DeploymentValues.DERP.Config.ForceWebSockets.Value(), DerpMapUpdateFrequency: api.Options.DERPMapUpdateFrequency, ExternalAuthConfigs: api.ExternalAuthConfigs, + Experiments: api.Experiments, // Optional: WorkspaceID: build.WorkspaceID, // saves the extra lookup later diff --git a/coderd/workspacestats/workspacestatstest/batcher.go b/coderd/workspacestats/workspacestatstest/batcher.go new file mode 100644 index 0000000000000..ad5ba60ad16d0 --- /dev/null +++ b/coderd/workspacestats/workspacestatstest/batcher.go @@ -0,0 +1,38 @@ +package workspacestatstest + +import ( + "sync" + "time" + + "github.com/google/uuid" + + agentproto "github.com/coder/coder/v2/agent/proto" + "github.com/coder/coder/v2/coderd/workspacestats" +) + +type StatsBatcher struct { + Mu sync.Mutex + + Called int64 + LastTime time.Time + LastAgentID uuid.UUID + LastTemplateID uuid.UUID + LastUserID uuid.UUID + LastWorkspaceID uuid.UUID + LastStats *agentproto.Stats +} + +var _ workspacestats.Batcher = &StatsBatcher{} + +func (b *StatsBatcher) Add(now time.Time, agentID uuid.UUID, templateID uuid.UUID, userID uuid.UUID, workspaceID uuid.UUID, st *agentproto.Stats) error { + b.Mu.Lock() + defer b.Mu.Unlock() + b.Called++ + b.LastTime = now + b.LastAgentID = agentID + b.LastTemplateID = templateID + b.LastUserID = userID + b.LastWorkspaceID = workspaceID + b.LastStats = st + return nil +} diff --git a/site/src/api/api.ts b/site/src/api/api.ts index 50dbc32a1867d..01b4ccec76b9d 100644 --- a/site/src/api/api.ts +++ b/site/src/api/api.ts @@ -24,6 +24,7 @@ import type dayjs from "dayjs"; import userAgentParser from "ua-parser-js"; import { delay } from "../utils/delay"; import * as TypesGen from "./typesGenerated"; +import type { PostWorkspaceUsageRequest } from "./typesGenerated"; const getMissingParameters = ( oldBuildParameters: TypesGen.WorkspaceBuildParameter[], @@ -1879,6 +1880,18 @@ class ApiMethods { throw error; } }; + + postWorkspaceUsage = async ( + workspaceID: string, + options: PostWorkspaceUsageRequest, + ) => { + const response = await this.axios.post( + `/api/v2/workspaces/${workspaceID}/usage`, + options, + ); + + return response.data; + }; } // This is a hard coded CSRF token/cookie pair for local development. In prod, diff --git a/site/src/api/queries/workspaces.ts b/site/src/api/queries/workspaces.ts index 809c8a4c2862f..aa5f8f29a9783 100644 --- a/site/src/api/queries/workspaces.ts +++ b/site/src/api/queries/workspaces.ts @@ -8,12 +8,14 @@ import { type DeleteWorkspaceOptions, API } from "api/api"; import type { CreateWorkspaceRequest, ProvisionerLogLevel, + UsageAppName, Workspace, WorkspaceBuild, WorkspaceBuildParameter, WorkspacesRequest, WorkspacesResponse, } from "api/typesGenerated"; +import type { ConnectionStatus } from "pages/TerminalPage/types"; import { disabledRefetchOptions } from "./util"; import { workspaceBuildsKey } from "./workspaceBuilds"; @@ -313,3 +315,41 @@ export const agentLogs = (workspaceId: string, agentId: string) => { ...disabledRefetchOptions, }; }; + +// workspace usage options +export interface WorkspaceUsageOptions { + usageApp: UsageAppName; + connectionStatus: ConnectionStatus; + workspaceId: string | undefined; + agentId: string | undefined; +} + +export const workspaceUsage = (options: WorkspaceUsageOptions) => { + return { + queryKey: [ + "workspaces", + options.workspaceId, + "agents", + options.agentId, + "usage", + options.usageApp, + ], + enabled: + options.workspaceId !== undefined && + options.agentId !== undefined && + options.connectionStatus === "connected", + queryFn: () => { + if (options.workspaceId === undefined || options.agentId === undefined) { + return Promise.reject(); + } + + return API.postWorkspaceUsage(options.workspaceId, { + agent_id: options.agentId, + app_name: options.usageApp, + }); + }, + // ...disabledRefetchOptions, + refetchInterval: 60000, + refetchIntervalInBackground: true, + }; +}; diff --git a/site/src/pages/TerminalPage/TerminalPage.tsx b/site/src/pages/TerminalPage/TerminalPage.tsx index 73ddc1ceebe06..4d99b3e4862f7 100644 --- a/site/src/pages/TerminalPage/TerminalPage.tsx +++ b/site/src/pages/TerminalPage/TerminalPage.tsx @@ -12,7 +12,10 @@ import { Unicode11Addon } from "xterm-addon-unicode11"; import { WebLinksAddon } from "xterm-addon-web-links"; import { WebglAddon } from "xterm-addon-webgl"; import { deploymentConfig } from "api/queries/deployment"; -import { workspaceByOwnerAndName } from "api/queries/workspaces"; +import { + workspaceByOwnerAndName, + workspaceUsage, +} from "api/queries/workspaces"; import { useProxy } from "contexts/ProxyContext"; import { ThemeOverride } from "contexts/ThemeProvider"; import themes from "theme"; @@ -67,6 +70,16 @@ const TerminalPage: FC = () => { const config = useQuery(deploymentConfig()); const renderer = config.data?.config.web_terminal_renderer; + // Periodically report workspace usage. + useQuery( + workspaceUsage({ + usageApp: "reconnecting-pty", + connectionStatus, + workspaceId: workspace.data?.id, + agentId: workspaceAgent?.id, + }), + ); + // handleWebLink handles opening of URLs in the terminal! const handleWebLink = useCallback( (uri: string) => {