coder · f0ssel · Jun 14, 2024 · Jun 11, 2024 · Jun 11, 2024 · Jun 12, 2024
diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go
diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json
diff --git a/coderd/workspaces.go b/coderd/workspaces.go
@@ -7,6 +7,7 @@ import (
 	"errors"
 	"fmt"
 	"net/http"
+	"slices"
 	"strconv"
 	"time"
 
@@ -15,6 +16,7 @@ import (
 	"golang.org/x/xerrors"
 
 	"cdr.dev/slog"
+	"github.com/coder/coder/v2/agent/proto"
 	"github.com/coder/coder/v2/coderd/audit"
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/db2sdk"
@@ -1105,7 +1107,9 @@ func (api *API) putExtendWorkspace(rw http.ResponseWriter, r *http.Request) {
 // @ID post-workspace-usage-by-id
 // @Security CoderSessionToken
 // @Tags Workspaces
+// @Accept json
 // @Param workspace path string true "Workspace ID" format(uuid)
+// @Param request body codersdk.PostWorkspaceUsageRequest false "Post workspace usage request"
 // @Success 204
 // @Router /workspaces/{workspace}/usage [post]
 func (api *API) postWorkspaceUsage(rw http.ResponseWriter, r *http.Request) {
@@ -1116,6 +1120,102 @@ func (api *API) postWorkspaceUsage(rw http.ResponseWriter, r *http.Request) {
 	}
 
 	api.statsReporter.TrackUsage(workspace.ID)
+
+	if !api.Experiments.Enabled(codersdk.ExperimentWorkspaceUsage) {
+		// Continue previous behavior if the experiment is not enabled.
+		rw.WriteHeader(http.StatusNoContent)
+		return
+	}
+
+	if r.Body == http.NoBody {
+		// Continue previous behavior if no body is present.
+		rw.WriteHeader(http.StatusNoContent)
+		return
+	}
+
+	ctx := r.Context()
+	var req codersdk.PostWorkspaceUsageRequest
+	if !httpapi.Read(ctx, rw, r, &req) {
+		return
+	}
+
+	if req.AgentID == uuid.Nil && req.AppName == "" {
+		// Continue previous behavior if body is empty.
+		rw.WriteHeader(http.StatusNoContent)
+		return
+	}
+	if req.AgentID == uuid.Nil {
+		httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
+			Message: "Invalid request",
+			Validations: []codersdk.ValidationError{{
+				Field:  "agent_id",
+				Detail: "must be set when app_name is set",
+			}},
+		})
+		return
+	}
+	if req.AppName == "" {
+		httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
+			Message: "Invalid request",
+			Validations: []codersdk.ValidationError{{
+				Field:  "app_name",
+				Detail: "must be set when agent_id is set",
+			}},
+		})
+		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
+	}
+
+	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
+	}
+
+	agent, err := api.Database.GetWorkspaceAgentByID(ctx, req.AgentID)
+	if err != nil {
+		if httpapi.Is404Error(err) {
+			httpapi.ResourceNotFound(rw)
+			return
+		}
+		httpapi.InternalServerError(rw, err)
+		return
+	}
+
+	template, err := api.Database.GetTemplateByID(ctx, workspace.TemplateID)
+	if err != nil {
+		httpapi.InternalServerError(rw, err)
+		return
+	}
+
+	err = api.statsReporter.ReportAgentStats(ctx, dbtime.Now(), workspace, agent, template.Name, stat)
+	if err != nil {
+		httpapi.InternalServerError(rw, err)
+		return
+	}
+
 	rw.WriteHeader(http.StatusNoContent)
 }