Skip to content

Commit 841cb33

Browse files
mafredrimafredri
committed
refactor: use task data model for notifications
Updates coder/internal#973 Updates coder/internal#974
1 parent 94f6e83 commit 841cb33

6 files changed

Lines changed: 92 additions & 89 deletions

File tree

‎coderd/aitasks_test.go‎

Lines changed: 23 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -2,15 +2,12 @@ package coderd_test
22

33
import (
44
"context"
5-
"database/sql"
65
"encoding/json"
76
"io"
87
"net/http"
98
"net/http/httptest"
10-
"strings"
119
"testing"
1210
"time"
13-
"unicode/utf8"
1411

1512
"github.com/google/uuid"
1613
"github.com/stretchr/testify/assert"
@@ -1285,31 +1282,31 @@ func TestTasksNotification(t *testing.T) {
12851282
// Given: a workspace build with an agent containing an App
12861283
workspaceAgentAppID := uuid.New()
12871284
workspaceBuildID := uuid.New()
1288-
workspaceBuildSeed := database.WorkspaceBuild{
1285+
workspaceBuilder := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
1286+
OrganizationID: ownerUser.OrganizationID,
1287+
OwnerID: memberUser.ID,
1288+
}).Seed(database.WorkspaceBuild{
12891289
ID: workspaceBuildID,
1290-
}
1290+
})
12911291
if tc.isAITask {
1292-
workspaceBuildSeed = database.WorkspaceBuild{
1293-
ID: workspaceBuildID,
1294-
// AI Task configuration
1295-
HasAITask: sql.NullBool{Bool: true, Valid: true},
1296-
AITaskSidebarAppID: uuid.NullUUID{UUID: workspaceAgentAppID, Valid: true},
1297-
}
1292+
workspaceBuilder = workspaceBuilder.
1293+
WithTask(database.TaskTable{
1294+
Prompt: tc.taskPrompt,
1295+
}, &proto.App{
1296+
Id: workspaceAgentAppID.String(),
1297+
Slug: "ccw",
1298+
})
1299+
} else {
1300+
workspaceBuilder = workspaceBuilder.
1301+
WithAgent(func(agent []*proto.Agent) []*proto.Agent {
1302+
agent[0].Apps = []*proto.App{{
1303+
Id: workspaceAgentAppID.String(),
1304+
Slug: "ccw",
1305+
}}
1306+
return agent
1307+
})
12981308
}
1299-
workspaceBuild := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
1300-
OrganizationID: ownerUser.OrganizationID,
1301-
OwnerID: memberUser.ID,
1302-
}).Seed(workspaceBuildSeed).Params(database.WorkspaceBuildParameter{
1303-
WorkspaceBuildID: workspaceBuildID,
1304-
Name: codersdk.AITaskPromptParameterName,
1305-
Value: tc.taskPrompt,
1306-
}).WithAgent(func(agent []*proto.Agent) []*proto.Agent {
1307-
agent[0].Apps = []*proto.App{{
1308-
Id: workspaceAgentAppID.String(),
1309-
Slug: "ccw",
1310-
}}
1311-
return agent
1312-
}).Do()
1309+
workspaceBuild := workspaceBuilder.Do()
13131310

13141311
// Given: the workspace agent app has previous statuses
13151312
agentClient := agentsdk.New(client.URL, agentsdk.WithFixedToken(workspaceBuild.AgentToken))
@@ -1350,13 +1347,7 @@ func TestTasksNotification(t *testing.T) {
13501347
require.Len(t, sent, 1)
13511348
require.Equal(t, memberUser.ID, sent[0].UserID)
13521349
require.Len(t, sent[0].Labels, 2)
1353-
// NOTE: len(string) is the number of bytes in the string, not the number of runes.
1354-
require.LessOrEqual(t, utf8.RuneCountInString(sent[0].Labels["task"]), 160)
1355-
if len(tc.taskPrompt) > 160 {
1356-
require.Contains(t, tc.taskPrompt, strings.TrimSuffix(sent[0].Labels["task"], "…"))
1357-
} else {
1358-
require.Equal(t, tc.taskPrompt, sent[0].Labels["task"])
1359-
}
1350+
require.Equal(t, workspaceBuild.Task.Name, sent[0].Labels["task"])
13601351
require.Equal(t, workspace.Name, sent[0].Labels["workspace"])
13611352
} else {
13621353
// Then: No notification is sent

‎coderd/database/queries.sql.go‎

Lines changed: 8 additions & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/queries/workspaceagents.sql‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -285,7 +285,8 @@ WHERE
285285
SELECT
286286
sqlc.embed(workspaces),
287287
sqlc.embed(workspace_agents),
288-
sqlc.embed(workspace_build_with_user)
288+
sqlc.embed(workspace_build_with_user),
289+
tasks.id AS task_id
289290
FROM
290291
workspace_agents
291292
JOIN
@@ -300,6 +301,10 @@ JOIN
300301
workspaces
301302
ON
302303
workspace_build_with_user.workspace_id = workspaces.id
304+
LEFT JOIN
305+
tasks
306+
ON
307+
tasks.workspace_id = workspaces.id
303308
WHERE
304309
-- This should only match 1 agent, so 1 returned row or 0.
305310
workspace_agents.auth_token = @auth_token::uuid

‎coderd/httpmw/workspaceagent.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,7 @@ func ExtractWorkspaceAgentAndLatestBuild(opts ExtractWorkspaceAgentAndLatestBuil
118118
OwnerID: row.WorkspaceTable.OwnerID,
119119
TemplateID: row.WorkspaceTable.TemplateID,
120120
VersionID: row.WorkspaceBuild.TemplateVersionID,
121+
TaskID: row.TaskID,
121122
BlockUserData: row.WorkspaceAgent.APIKeyScope == database.AgentKeyScopeEnumNoUserData,
122123
}),
123124
)

‎coderd/rbac/scopes.go‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ type WorkspaceAgentScopeParams struct {
1818
OwnerID uuid.UUID
1919
TemplateID uuid.UUID
2020
VersionID uuid.UUID
21+
TaskID uuid.NullUUID
2122
BlockUserData bool
2223
}
2324

@@ -42,6 +43,15 @@ func WorkspaceAgentScope(params WorkspaceAgentScopeParams) Scope {
4243
panic("failed to expand scope, this should never happen")
4344
}
4445

46+
// Include task in the allow list if the workspace has an associated task.
47+
var extraAllowList []AllowListElement
48+
if params.TaskID.Valid {
49+
extraAllowList = append(extraAllowList, AllowListElement{
50+
Type: ResourceTask.Type,
51+
ID: params.TaskID.UUID.String(),
52+
})
53+
}
54+
4555
return Scope{
4656
// TODO: We want to limit the role too to be extra safe.
4757
// Even though the allowlist blocks anything else, it is still good
@@ -52,12 +62,12 @@ func WorkspaceAgentScope(params WorkspaceAgentScopeParams) Scope {
5262
// Limit the agent to only be able to access the singular workspace and
5363
// the template/version it was created from. Add additional resources here
5464
// as needed, but do not add more workspace or template resource ids.
55-
AllowIDList: []AllowListElement{
65+
AllowIDList: append([]AllowListElement{
5666
{Type: ResourceWorkspace.Type, ID: params.WorkspaceID.String()},
5767
{Type: ResourceTemplate.Type, ID: params.TemplateID.String()},
5868
{Type: ResourceTemplate.Type, ID: params.VersionID.String()},
5969
{Type: ResourceUser.Type, ID: params.OwnerID.String()},
60-
},
70+
}, extraAllowList...),
6171
}
6272
}
6373

‎coderd/workspaceagents.go‎

Lines changed: 42 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -461,67 +461,56 @@ func (api *API) enqueueAITaskStateNotification(
461461
return
462462
}
463463

464-
workspaceBuild, err := api.Database.GetLatestWorkspaceBuildByWorkspaceID(ctx, workspace.ID)
464+
if !workspace.TaskID.Valid {
465+
api.Logger.Warn(ctx, "workspace has no task ID")
466+
return
467+
}
468+
469+
task, err := api.Database.GetTaskByID(ctx, workspace.TaskID.UUID)
465470
if err != nil {
466-
api.Logger.Warn(ctx, "failed to get workspace build", slog.Error(err))
471+
api.Logger.Warn(ctx, "failed to get task", slog.Error(err))
467472
return
468473
}
469474

470475
// Confirm Workspace Agent App is an AI Task
471-
if workspaceBuild.HasAITask.Valid && workspaceBuild.HasAITask.Bool &&
472-
workspaceBuild.AITaskSidebarAppID.Valid && workspaceBuild.AITaskSidebarAppID.UUID == appID {
473-
// Skip if the latest persisted state equals the new state (no new transition)
474-
if len(latestAppStatus) > 0 && latestAppStatus[0].State == database.WorkspaceAppStatusState(newAppStatus) {
475-
return
476-
}
477-
478-
// Skip the initial "Working" notification when task first starts.
479-
// This is obvious to the user since they just created the task.
480-
// We still notify on first "Idle" status and all subsequent transitions.
481-
if len(latestAppStatus) == 0 && newAppStatus == codersdk.WorkspaceAppStatusStateWorking {
482-
return
483-
}
476+
if !task.WorkspaceAppID.Valid || task.WorkspaceAppID.UUID != appID {
477+
api.Logger.Warn(ctx, "workspace agent app is not an AI task")
478+
return
479+
}
484480

485-
// Use the task prompt as the "task" label, fallback to workspace name
486-
parameters, err := api.Database.GetWorkspaceBuildParameters(ctx, workspaceBuild.ID)
487-
if err != nil {
488-
api.Logger.Warn(ctx, "failed to get workspace build parameters", slog.Error(err))
489-
return
490-
}
491-
taskName := workspace.Name
492-
for _, param := range parameters {
493-
if param.Name == codersdk.AITaskPromptParameterName {
494-
taskName = param.Value
495-
}
496-
}
481+
// Skip if the latest persisted state equals the new state (no new transition)
482+
if len(latestAppStatus) > 0 && latestAppStatus[0].State == database.WorkspaceAppStatusState(newAppStatus) {
483+
return
484+
}
497485

498-
// As task prompt may be particularly long, truncate it to 160 characters for notifications.
499-
if len(taskName) > 160 {
500-
taskName = strutil.Truncate(taskName, 160, strutil.TruncateWithEllipsis, strutil.TruncateWithFullWords)
501-
}
486+
// Skip the initial "Working" notification when task first starts.
487+
// This is obvious to the user since they just created the task.
488+
// We still notify on first "Idle" status and all subsequent transitions.
489+
if len(latestAppStatus) == 0 && newAppStatus == codersdk.WorkspaceAppStatusStateWorking {
490+
return
491+
}
502492

503-
if _, err := api.NotificationsEnqueuer.EnqueueWithData(
504-
// nolint:gocritic // Need notifier actor to enqueue notifications
505-
dbauthz.AsNotifier(ctx),
506-
workspace.OwnerID,
507-
notificationTemplate,
508-
map[string]string{
509-
"task": taskName,
510-
"workspace": workspace.Name,
511-
},
512-
map[string]any{
513-
// Use a 1-minute bucketed timestamp to bypass per-day dedupe,
514-
// allowing identical content to resend within the same day
515-
// (but not more than once every 10s).
516-
"dedupe_bypass_ts": api.Clock.Now().UTC().Truncate(time.Minute),
517-
},
518-
"api-workspace-agent-app-status",
519-
// Associate this notification with related entities
520-
workspace.ID, workspace.OwnerID, workspace.OrganizationID, appID,
521-
); err != nil {
522-
api.Logger.Warn(ctx, "failed to notify of task state", slog.Error(err))
523-
return
524-
}
493+
if _, err := api.NotificationsEnqueuer.EnqueueWithData(
494+
// nolint:gocritic // Need notifier actor to enqueue notifications
495+
dbauthz.AsNotifier(ctx),
496+
workspace.OwnerID,
497+
notificationTemplate,
498+
map[string]string{
499+
"task": task.Name,
500+
"workspace": workspace.Name,
501+
},
502+
map[string]any{
503+
// Use a 1-minute bucketed timestamp to bypass per-day dedupe,
504+
// allowing identical content to resend within the same day
505+
// (but not more than once every 10s).
506+
"dedupe_bypass_ts": api.Clock.Now().UTC().Truncate(time.Minute),
507+
},
508+
"api-workspace-agent-app-status",
509+
// Associate this notification with related entities
510+
workspace.ID, workspace.OwnerID, workspace.OrganizationID, appID,
511+
); err != nil {
512+
api.Logger.Warn(ctx, "failed to notify of task state", slog.Error(err))
513+
return
525514
}
526515
}
527516

0 commit comments

Comments
 (0)