Skip to content

Commit a2a2e10

Browse files
johnstcnjohnstcn
committed
feat(telemetry): track AI task usage
1 parent 9e5c83a commit a2a2e10

File tree

2 files changed

+67
-9
lines changed

2 files changed

+67
-9
lines changed

‎coderd/telemetry/telemetry.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -776,6 +776,7 @@ func ConvertWorkspaceBuild(build database.WorkspaceBuild) WorkspaceBuild {
776776
TemplateVersionID: build.TemplateVersionID,
777777
// #nosec G115 - Safe conversion as build numbers are expected to be positive and within uint32 range
778778
BuildNumber: uint32(build.BuildNumber),
779+
HasAITask: build.HasAITask.Valid && build.HasAITask.Bool,
779780
}
780781
}
781782

@@ -1105,6 +1106,7 @@ func ConvertTemplateVersion(version database.TemplateVersion) TemplateVersion {
11051106
if version.SourceExampleID.Valid {
11061107
snapVersion.SourceExampleID = &version.SourceExampleID.String
11071108
}
1109+
snapVersion.HasAITask = version.HasAITask.Valid && version.HasAITask.Bool
11081110
return snapVersion
11091111
}
11101112

@@ -1357,6 +1359,7 @@ type WorkspaceBuild struct {
13571359
TemplateVersionID uuid.UUID `json:"template_version_id"`
13581360
JobID uuid.UUID `json:"job_id"`
13591361
BuildNumber uint32 `json:"build_number"`
1362+
HasAITask bool `json:"has_ai_task"`
13601363
}
13611364

13621365
type Workspace struct {
@@ -1404,6 +1407,7 @@ type TemplateVersion struct {
14041407
OrganizationID uuid.UUID `json:"organization_id"`
14051408
JobID uuid.UUID `json:"job_id"`
14061409
SourceExampleID *string `json:"source_example_id,omitempty"`
1410+
HasAITask bool `json:"has_ai_task"`
14071411
}
14081412

14091413
type ProvisionerJob struct {

‎coderd/telemetry/telemetry_test.go

Lines changed: 63 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77
"net/http"
88
"net/http/httptest"
99
"net/url"
10+
"slices"
1011
"sort"
1112
"testing"
1213
"time"
@@ -105,6 +106,52 @@ func TestTelemetry(t *testing.T) {
105106
OpenIn: database.WorkspaceAppOpenInSlimWindow,
106107
AgentID: wsagent.ID,
107108
})
109+
110+
taskJob := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
111+
Provisioner: database.ProvisionerTypeTerraform,
112+
StorageMethod: database.ProvisionerStorageMethodFile,
113+
Type: database.ProvisionerJobTypeTemplateVersionDryRun,
114+
OrganizationID: org.ID,
115+
})
116+
taskTpl := dbgen.Template(t, db, database.Template{
117+
Provisioner: database.ProvisionerTypeTerraform,
118+
OrganizationID: org.ID,
119+
CreatedBy: user.ID,
120+
})
121+
_ = dbgen.TemplateVersion(t, db, database.TemplateVersion{
122+
OrganizationID: org.ID,
123+
TemplateID: uuid.NullUUID{UUID: taskTpl.ID, Valid: true},
124+
CreatedBy: user.ID,
125+
JobID: taskJob.ID,
126+
HasAITask: sql.NullBool{Bool: true, Valid: true},
127+
})
128+
taskWs := dbgen.Workspace(t, db, database.WorkspaceTable{
129+
OwnerID: user.ID,
130+
OrganizationID: org.ID,
131+
TemplateID: taskTpl.ID,
132+
})
133+
taskWsResource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
134+
JobID: taskJob.ID,
135+
})
136+
taskWsAgent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
137+
ResourceID: taskWsResource.ID,
138+
})
139+
taskWsApp := dbgen.WorkspaceApp(t, db, database.WorkspaceApp{
140+
SharingLevel: database.AppSharingLevelOwner,
141+
Health: database.WorkspaceAppHealthDisabled,
142+
OpenIn: database.WorkspaceAppOpenInSlimWindow,
143+
AgentID: taskWsAgent.ID,
144+
})
145+
_ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
146+
Transition: database.WorkspaceTransitionStart,
147+
Reason: database.BuildReasonAutostart,
148+
WorkspaceID: taskWs.ID,
149+
TemplateVersionID: tv.ID,
150+
JobID: taskJob.ID,
151+
HasAITask: sql.NullBool{Valid: true, Bool: true},
152+
AITaskSidebarAppID: uuid.NullUUID{Valid: true, UUID: taskWsApp.ID},
153+
})
154+
108155
group := dbgen.Group(t, db, database.Group{
109156
OrganizationID: org.ID,
110157
})
@@ -148,19 +195,19 @@ func TestTelemetry(t *testing.T) {
148195
})
149196

150197
_, snapshot := collectSnapshot(ctx, t, db, nil)
151-
require.Len(t, snapshot.ProvisionerJobs, 1)
198+
require.Len(t, snapshot.ProvisionerJobs, 2)
152199
require.Len(t, snapshot.Licenses, 1)
153-
require.Len(t, snapshot.Templates, 1)
154-
require.Len(t, snapshot.TemplateVersions, 2)
200+
require.Len(t, snapshot.Templates, 2)
201+
require.Len(t, snapshot.TemplateVersions, 3)
155202
require.Len(t, snapshot.Users, 1)
156203
require.Len(t, snapshot.Groups, 2)
157204
// 1 member in the everyone group + 1 member in the custom group
158205
require.Len(t, snapshot.GroupMembers, 2)
159-
require.Len(t, snapshot.Workspaces, 1)
160-
require.Len(t, snapshot.WorkspaceApps, 1)
161-
require.Len(t, snapshot.WorkspaceAgents, 1)
162-
require.Len(t, snapshot.WorkspaceBuilds, 1)
163-
require.Len(t, snapshot.WorkspaceResources, 1)
206+
require.Len(t, snapshot.Workspaces, 2)
207+
require.Len(t, snapshot.WorkspaceApps, 2)
208+
require.Len(t, snapshot.WorkspaceAgents, 2)
209+
require.Len(t, snapshot.WorkspaceBuilds, 2)
210+
require.Len(t, snapshot.WorkspaceResources, 2)
164211
require.Len(t, snapshot.WorkspaceAgentStats, 1)
165212
require.Len(t, snapshot.WorkspaceProxies, 1)
166213
require.Len(t, snapshot.WorkspaceModules, 1)
@@ -169,11 +216,18 @@ func TestTelemetry(t *testing.T) {
169216
require.Len(t, snapshot.TelemetryItems, 2)
170217
require.Len(t, snapshot.WorkspaceAgentMemoryResourceMonitors, 1)
171218
require.Len(t, snapshot.WorkspaceAgentVolumeResourceMonitors, 1)
172-
wsa := snapshot.WorkspaceAgents[0]
219+
wsa := snapshot.WorkspaceAgents[1]
173220
require.Len(t, wsa.Subsystems, 2)
174221
require.Equal(t, string(database.WorkspaceAgentSubsystemEnvbox), wsa.Subsystems[0])
175222
require.Equal(t, string(database.WorkspaceAgentSubsystemExectrace), wsa.Subsystems[1])
176223

224+
require.True(t, slices.ContainsFunc(snapshot.TemplateVersions, func(ttv telemetry.TemplateVersion) bool {
225+
return ttv.HasAITask
226+
}))
227+
require.True(t, slices.ContainsFunc(snapshot.WorkspaceBuilds, func(twb telemetry.WorkspaceBuild) bool {
228+
return twb.HasAITask
229+
}))
230+
177231
tvs := snapshot.TemplateVersions
178232
sort.Slice(tvs, func(i, j int) bool {
179233
// Sort by SourceExampleID presence (non-nil comes before nil)

0 commit comments

Comments
 (0)