Skip to content

Commit fdb0267

Browse files
ssncferreirassncferreira
authored
feat: add notification for task status (#19965)
## Description Send a notification to the workspace owner when an AI task’s app state becomes `Working` or `Idle`. An AI task is identified by a workspace build with `HasAITask = true` and `AITaskSidebarAppID` matching the agent app’s ID. ## Changes * Add `TemplateTaskWorking` notification template. * Add `TemplateTaskIdle` notification template. * Add `GetLatestWorkspaceAppStatusesByAppID` SQL query to get the workspace app statuses ordered by latest first. * Update `PATCH /workspaceagents/me/app-status` to enqueue: * `TemplateTaskWorking` when state transitions to `working` * `TemplateTaskIdle` when state transitions to `idle` * Notification labels include: * `task`: task initial prompt * `workspace`: workspace name * Notification dedupe: include a minute-bucketed timestamp (UTC truncated to the minute) in the enqueue data to allow identical content to resend within the same day (but not more than once per minute). Closes: #19776
1 parent abdea72 commit fdb0267

18 files changed

Lines changed: 689 additions & 1 deletion

‎coderd/aitasks_test.go‎

Lines changed: 166 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package coderd_test
22

33
import (
4+
"database/sql"
45
"fmt"
56
"io"
67
"net/http"
@@ -17,8 +18,12 @@ import (
1718
"github.com/coder/coder/v2/coderd/coderdtest"
1819
"github.com/coder/coder/v2/coderd/database"
1920
"github.com/coder/coder/v2/coderd/database/dbauthz"
21+
"github.com/coder/coder/v2/coderd/database/dbfake"
22+
"github.com/coder/coder/v2/coderd/database/dbgen"
2023
"github.com/coder/coder/v2/coderd/database/dbtestutil"
2124
"github.com/coder/coder/v2/coderd/database/dbtime"
25+
"github.com/coder/coder/v2/coderd/notifications"
26+
"github.com/coder/coder/v2/coderd/notifications/notificationstest"
2227
"github.com/coder/coder/v2/coderd/util/slice"
2328
"github.com/coder/coder/v2/codersdk"
2429
"github.com/coder/coder/v2/codersdk/agentsdk"
@@ -961,3 +966,164 @@ func TestTasksCreate(t *testing.T) {
961966
assert.Equal(t, http.StatusNotFound, sdkErr.StatusCode())
962967
})
963968
}
969+
970+
func TestTasksNotification(t *testing.T) {
971+
t.Parallel()
972+
973+
for _, tc := range []struct {
974+
name string
975+
latestAppStatuses []codersdk.WorkspaceAppStatusState
976+
newAppStatus codersdk.WorkspaceAppStatusState
977+
isAITask bool
978+
isNotificationSent bool
979+
notificationTemplate uuid.UUID
980+
}{
981+
// Should not send a notification when the agent app is not an AI task.
982+
{
983+
name: "NoAITask",
984+
latestAppStatuses: nil,
985+
newAppStatus: codersdk.WorkspaceAppStatusStateWorking,
986+
isAITask: false,
987+
isNotificationSent: false,
988+
},
989+
// Should not send a notification when the new app status is neither 'Working' nor 'Idle'.
990+
{
991+
name: "NonNotifiedState",
992+
latestAppStatuses: nil,
993+
newAppStatus: codersdk.WorkspaceAppStatusStateComplete,
994+
isAITask: true,
995+
isNotificationSent: false,
996+
},
997+
// Should not send a notification when the new app status equals the latest status (Working).
998+
{
999+
name: "NonNotifiedTransition",
1000+
latestAppStatuses: []codersdk.WorkspaceAppStatusState{codersdk.WorkspaceAppStatusStateWorking},
1001+
newAppStatus: codersdk.WorkspaceAppStatusStateWorking,
1002+
isAITask: true,
1003+
isNotificationSent: false,
1004+
},
1005+
// Should send TemplateTaskWorking when the AI task transitions to 'Working'.
1006+
{
1007+
name: "TemplateTaskWorking",
1008+
latestAppStatuses: nil,
1009+
newAppStatus: codersdk.WorkspaceAppStatusStateWorking,
1010+
isAITask: true,
1011+
isNotificationSent: true,
1012+
notificationTemplate: notifications.TemplateTaskWorking,
1013+
},
1014+
// Should send TemplateTaskWorking when the AI task transitions to 'Working' from 'Idle'.
1015+
{
1016+
name: "TemplateTaskWorkingFromIdle",
1017+
latestAppStatuses: []codersdk.WorkspaceAppStatusState{
1018+
codersdk.WorkspaceAppStatusStateWorking,
1019+
codersdk.WorkspaceAppStatusStateIdle,
1020+
}, // latest
1021+
newAppStatus: codersdk.WorkspaceAppStatusStateWorking,
1022+
isAITask: true,
1023+
isNotificationSent: true,
1024+
notificationTemplate: notifications.TemplateTaskWorking,
1025+
},
1026+
// Should send TemplateTaskIdle when the AI task transitions to 'Idle'.
1027+
{
1028+
name: "TemplateTaskIdle",
1029+
latestAppStatuses: []codersdk.WorkspaceAppStatusState{codersdk.WorkspaceAppStatusStateWorking},
1030+
newAppStatus: codersdk.WorkspaceAppStatusStateIdle,
1031+
isAITask: true,
1032+
isNotificationSent: true,
1033+
notificationTemplate: notifications.TemplateTaskIdle,
1034+
},
1035+
} {
1036+
t.Run(tc.name, func(t *testing.T) {
1037+
t.Parallel()
1038+
1039+
ctx := testutil.Context(t, testutil.WaitShort)
1040+
notifyEnq := &notificationstest.FakeEnqueuer{}
1041+
client, db := coderdtest.NewWithDatabase(t, &coderdtest.Options{
1042+
DeploymentValues: coderdtest.DeploymentValues(t),
1043+
NotificationsEnqueuer: notifyEnq,
1044+
})
1045+
1046+
// Given: a member user
1047+
ownerUser := coderdtest.CreateFirstUser(t, client)
1048+
client, memberUser := coderdtest.CreateAnotherUser(t, client, ownerUser.OrganizationID)
1049+
1050+
// Given: a workspace build with an agent containing an App
1051+
workspaceAgentAppID := uuid.New()
1052+
workspaceBuildID := uuid.New()
1053+
workspaceBuildSeed := database.WorkspaceBuild{
1054+
ID: workspaceBuildID,
1055+
}
1056+
if tc.isAITask {
1057+
workspaceBuildSeed = database.WorkspaceBuild{
1058+
ID: workspaceBuildID,
1059+
// AI Task configuration
1060+
HasAITask: sql.NullBool{Bool: true, Valid: true},
1061+
AITaskSidebarAppID: uuid.NullUUID{UUID: workspaceAgentAppID, Valid: true},
1062+
}
1063+
}
1064+
workspaceBuild := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
1065+
OrganizationID: ownerUser.OrganizationID,
1066+
OwnerID: memberUser.ID,
1067+
}).Seed(workspaceBuildSeed).Params(database.WorkspaceBuildParameter{
1068+
WorkspaceBuildID: workspaceBuildID,
1069+
Name: codersdk.AITaskPromptParameterName,
1070+
Value: "task prompt",
1071+
}).WithAgent(func(agent []*proto.Agent) []*proto.Agent {
1072+
agent[0].Apps = []*proto.App{{
1073+
Id: workspaceAgentAppID.String(),
1074+
Slug: "ccw",
1075+
}}
1076+
return agent
1077+
}).Do()
1078+
1079+
// Given: the workspace agent app has previous statuses
1080+
agentClient := agentsdk.New(client.URL, agentsdk.WithFixedToken(workspaceBuild.AgentToken))
1081+
if len(tc.latestAppStatuses) > 0 {
1082+
workspace := coderdtest.MustWorkspace(t, client, workspaceBuild.Workspace.ID)
1083+
for _, appStatus := range tc.latestAppStatuses {
1084+
dbgen.WorkspaceAppStatus(t, db, database.WorkspaceAppStatus{
1085+
WorkspaceID: workspaceBuild.Workspace.ID,
1086+
AgentID: workspace.LatestBuild.Resources[0].Agents[0].ID,
1087+
AppID: workspaceAgentAppID,
1088+
State: database.WorkspaceAppStatusState(appStatus),
1089+
})
1090+
}
1091+
}
1092+
1093+
// When: the agent updates the app status
1094+
err := agentClient.PatchAppStatus(ctx, agentsdk.PatchAppStatus{
1095+
AppSlug: "ccw",
1096+
Message: "testing",
1097+
URI: "https://example.com",
1098+
State: tc.newAppStatus,
1099+
})
1100+
require.NoError(t, err)
1101+
1102+
// Then: The workspace app status transitions successfully
1103+
workspace, err := client.Workspace(ctx, workspaceBuild.Workspace.ID)
1104+
require.NoError(t, err)
1105+
workspaceAgent, err := client.WorkspaceAgent(ctx, workspace.LatestBuild.Resources[0].Agents[0].ID)
1106+
require.NoError(t, err)
1107+
require.Len(t, workspaceAgent.Apps, 1)
1108+
require.GreaterOrEqual(t, len(workspaceAgent.Apps[0].Statuses), 1)
1109+
latestStatusIndex := len(workspaceAgent.Apps[0].Statuses) - 1
1110+
require.Equal(t, tc.newAppStatus, workspaceAgent.Apps[0].Statuses[latestStatusIndex].State)
1111+
1112+
if tc.isNotificationSent {
1113+
// Then: A notification is sent to the workspace owner (memberUser)
1114+
sent := notifyEnq.Sent(notificationstest.WithTemplateID(tc.notificationTemplate))
1115+
require.Len(t, sent, 1)
1116+
require.Equal(t, memberUser.ID, sent[0].UserID)
1117+
require.Len(t, sent[0].Labels, 2)
1118+
require.Equal(t, "task prompt", sent[0].Labels["task"])
1119+
require.Equal(t, workspace.Name, sent[0].Labels["workspace"])
1120+
} else {
1121+
// Then: No notification is sent
1122+
sentWorking := notifyEnq.Sent(notificationstest.WithTemplateID(notifications.TemplateTaskWorking))
1123+
sentIdle := notifyEnq.Sent(notificationstest.WithTemplateID(notifications.TemplateTaskIdle))
1124+
require.Len(t, sentWorking, 0)
1125+
require.Len(t, sentIdle, 0)
1126+
}
1127+
})
1128+
}
1129+
}

‎coderd/database/dbauthz/dbauthz.go‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2313,6 +2313,13 @@ func (q *querier) GetLatestCryptoKeyByFeature(ctx context.Context, feature datab
23132313
return q.db.GetLatestCryptoKeyByFeature(ctx, feature)
23142314
}
23152315

2316+
func (q *querier) GetLatestWorkspaceAppStatusesByAppID(ctx context.Context, appID uuid.UUID) ([]database.WorkspaceAppStatus, error) {
2317+
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceSystem); err != nil {
2318+
return nil, err
2319+
}
2320+
return q.db.GetLatestWorkspaceAppStatusesByAppID(ctx, appID)
2321+
}
2322+
23162323
func (q *querier) GetLatestWorkspaceAppStatusesByWorkspaceIDs(ctx context.Context, ids []uuid.UUID) ([]database.WorkspaceAppStatus, error) {
23172324
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceSystem); err != nil {
23182325
return nil, err

‎coderd/database/dbauthz/dbauthz_test.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2683,6 +2683,11 @@ func (s *MethodTestSuite) TestSystemFunctions() {
26832683
dbm.EXPECT().UpdateUserLinkedID(gomock.Any(), arg).Return(l, nil).AnyTimes()
26842684
check.Args(arg).Asserts(rbac.ResourceSystem, policy.ActionUpdate).Returns(l)
26852685
}))
2686+
s.Run("GetLatestWorkspaceAppStatusesByAppID", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
2687+
appID := uuid.New()
2688+
dbm.EXPECT().GetLatestWorkspaceAppStatusesByAppID(gomock.Any(), appID).Return([]database.WorkspaceAppStatus{}, nil).AnyTimes()
2689+
check.Args(appID).Asserts(rbac.ResourceSystem, policy.ActionRead)
2690+
}))
26862691
s.Run("GetLatestWorkspaceAppStatusesByWorkspaceIDs", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
26872692
ids := []uuid.UUID{uuid.New()}
26882693
dbm.EXPECT().GetLatestWorkspaceAppStatusesByWorkspaceIDs(gomock.Any(), ids).Return([]database.WorkspaceAppStatus{}, nil).AnyTimes()

‎coderd/database/dbgen/dbgen.go‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -905,6 +905,21 @@ func WorkspaceAppStat(t testing.TB, db database.Store, orig database.WorkspaceAp
905905
return scheme
906906
}
907907

908+
func WorkspaceAppStatus(t testing.TB, db database.Store, orig database.WorkspaceAppStatus) database.WorkspaceAppStatus {
909+
appStatus, err := db.InsertWorkspaceAppStatus(genCtx, database.InsertWorkspaceAppStatusParams{
910+
ID: takeFirst(orig.ID, uuid.New()),
911+
CreatedAt: takeFirst(orig.CreatedAt, dbtime.Now()),
912+
WorkspaceID: takeFirst(orig.WorkspaceID, uuid.New()),
913+
AgentID: takeFirst(orig.AgentID, uuid.New()),
914+
AppID: takeFirst(orig.AppID, uuid.New()),
915+
State: takeFirst(orig.State, database.WorkspaceAppStatusStateWorking),
916+
Message: takeFirst(orig.Message, ""),
917+
Uri: takeFirst(orig.Uri, sql.NullString{}),
918+
})
919+
require.NoError(t, err, "insert workspace agent status")
920+
return appStatus
921+
}
922+
908923
func WorkspaceResource(t testing.TB, db database.Store, orig database.WorkspaceResource) database.WorkspaceResource {
909924
resource, err := db.InsertWorkspaceResource(genCtx, database.InsertWorkspaceResourceParams{
910925
ID: takeFirst(orig.ID, uuid.New()),

‎coderd/database/dbmetrics/querymetrics.go‎

Lines changed: 7 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/dbmock/dbmock.go‎

Lines changed: 15 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/migrations/000376_task_status_notifications.down.sql‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
-- Remove Task 'working' transition template notification
2+
DELETE FROM notification_templates WHERE id = 'bd4b7168-d05e-4e19-ad0f-3593b77aa90f';
3+
-- Remove Task 'idle' transition template notification
4+
DELETE FROM notification_templates WHERE id = 'd4a6271c-cced-4ed0-84ad-afd02a9c7799';

‎coderd/database/migrations/000376_task_status_notifications.up.sql‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
-- Task transition to 'working' status
2+
INSERT INTO notification_templates (
3+
id,
4+
name,
5+
title_template,
6+
body_template,
7+
actions,
8+
"group",
9+
method,
10+
kind,
11+
enabled_by_default
12+
) VALUES (
13+
'bd4b7168-d05e-4e19-ad0f-3593b77aa90f',
14+
'Task Working',
15+
E'Task ''{{.Labels.workspace}}'' is working',
16+
E'The task ''{{.Labels.task}}'' transitioned to a working state.',
17+
'[
18+
{
19+
"label": "View task",
20+
"url": "{{base_url}}/tasks/{{.UserUsername}}/{{.Labels.workspace}}"
21+
},
22+
{
23+
"label": "View workspace",
24+
"url": "{{base_url}}/@{{.UserUsername}}/{{.Labels.workspace}}"
25+
}
26+
]'::jsonb,
27+
'Task Events',
28+
NULL,
29+
'system'::notification_template_kind,
30+
true
31+
);
32+
33+
-- Task transition to 'idle' status
34+
INSERT INTO notification_templates (
35+
id,
36+
name,
37+
title_template,
38+
body_template,
39+
actions,
40+
"group",
41+
method,
42+
kind,
43+
enabled_by_default
44+
) VALUES (
45+
'd4a6271c-cced-4ed0-84ad-afd02a9c7799',
46+
'Task Idle',
47+
E'Task ''{{.Labels.workspace}}'' is idle',
48+
E'The task ''{{.Labels.task}}'' is idle and ready for input.',
49+
'[
50+
{
51+
"label": "View task",
52+
"url": "{{base_url}}/tasks/{{.UserUsername}}/{{.Labels.workspace}}"
53+
},
54+
{
55+
"label": "View workspace",
56+
"url": "{{base_url}}/@{{.UserUsername}}/{{.Labels.workspace}}"
57+
}
58+
]'::jsonb,
59+
'Task Events',
60+
NULL,
61+
'system'::notification_template_kind,
62+
true
63+
);

‎coderd/database/querier.go‎

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/queries.sql.go‎

Lines changed: 39 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)