From 7ae3fdc7496376291a86d684695460c69a206b2e Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Fri, 31 Oct 2025 15:53:27 +0200 Subject: [PATCH] refactor: use task data model for notifications (#20590) Updates coder/internal#973 Updates coder/internal#974 --- coderd/aitasks_test.go | 57 ++++++------ coderd/database/queries.sql.go | 9 +- coderd/database/queries/workspaceagents.sql | 7 +- coderd/httpmw/workspaceagent.go | 1 + coderd/rbac/scopes.go | 14 ++- coderd/workspaceagents.go | 96 +++++++++------------ 6 files changed, 93 insertions(+), 91 deletions(-) diff --git a/coderd/aitasks_test.go b/coderd/aitasks_test.go index 4f645d34cf..34f6dd4a07 100644 --- a/coderd/aitasks_test.go +++ b/coderd/aitasks_test.go @@ -2,15 +2,12 @@ package coderd_test import ( "context" - "database/sql" "encoding/json" "io" "net/http" "net/http/httptest" - "strings" "testing" "time" - "unicode/utf8" "github.com/google/uuid" "github.com/stretchr/testify/assert" @@ -1325,31 +1322,31 @@ func TestTasksNotification(t *testing.T) { // Given: a workspace build with an agent containing an App workspaceAgentAppID := uuid.New() workspaceBuildID := uuid.New() - workspaceBuildSeed := database.WorkspaceBuild{ - ID: workspaceBuildID, - } - if tc.isAITask { - workspaceBuildSeed = database.WorkspaceBuild{ - ID: workspaceBuildID, - // AI Task configuration - HasAITask: sql.NullBool{Bool: true, Valid: true}, - AITaskSidebarAppID: uuid.NullUUID{UUID: workspaceAgentAppID, Valid: true}, - } - } - workspaceBuild := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ + workspaceBuilder := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ OrganizationID: ownerUser.OrganizationID, OwnerID: memberUser.ID, - }).Seed(workspaceBuildSeed).Params(database.WorkspaceBuildParameter{ - WorkspaceBuildID: workspaceBuildID, - Name: codersdk.AITaskPromptParameterName, - Value: tc.taskPrompt, - }).WithAgent(func(agent []*proto.Agent) []*proto.Agent { - agent[0].Apps = []*proto.App{{ - Id: workspaceAgentAppID.String(), - Slug: "ccw", - }} - return agent - }).Do() + }).Seed(database.WorkspaceBuild{ + ID: workspaceBuildID, + }) + if tc.isAITask { + workspaceBuilder = workspaceBuilder. + WithTask(database.TaskTable{ + Prompt: tc.taskPrompt, + }, &proto.App{ + Id: workspaceAgentAppID.String(), + Slug: "ccw", + }) + } else { + workspaceBuilder = workspaceBuilder. + WithAgent(func(agent []*proto.Agent) []*proto.Agent { + agent[0].Apps = []*proto.App{{ + Id: workspaceAgentAppID.String(), + Slug: "ccw", + }} + return agent + }) + } + workspaceBuild := workspaceBuilder.Do() // Given: the workspace agent app has previous statuses agentClient := agentsdk.New(client.URL, agentsdk.WithFixedToken(workspaceBuild.AgentToken)) @@ -1390,13 +1387,7 @@ func TestTasksNotification(t *testing.T) { require.Len(t, sent, 1) require.Equal(t, memberUser.ID, sent[0].UserID) require.Len(t, sent[0].Labels, 2) - // NOTE: len(string) is the number of bytes in the string, not the number of runes. - require.LessOrEqual(t, utf8.RuneCountInString(sent[0].Labels["task"]), 160) - if len(tc.taskPrompt) > 160 { - require.Contains(t, tc.taskPrompt, strings.TrimSuffix(sent[0].Labels["task"], "…")) - } else { - require.Equal(t, tc.taskPrompt, sent[0].Labels["task"]) - } + require.Equal(t, workspaceBuild.Task.Name, sent[0].Labels["task"]) require.Equal(t, workspace.Name, sent[0].Labels["workspace"]) } else { // Then: No notification is sent diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index e9108751cf..65fac4733b 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -17569,7 +17569,8 @@ const getWorkspaceAgentAndLatestBuildByAuthToken = `-- name: GetWorkspaceAgentAn SELECT workspaces.id, workspaces.created_at, workspaces.updated_at, workspaces.owner_id, workspaces.organization_id, workspaces.template_id, workspaces.deleted, workspaces.name, workspaces.autostart_schedule, workspaces.ttl, workspaces.last_used_at, workspaces.dormant_at, workspaces.deleting_at, workspaces.automatic_updates, workspaces.favorite, workspaces.next_start_at, workspaces.group_acl, workspaces.user_acl, workspace_agents.id, workspace_agents.created_at, workspace_agents.updated_at, workspace_agents.name, workspace_agents.first_connected_at, workspace_agents.last_connected_at, workspace_agents.disconnected_at, workspace_agents.resource_id, workspace_agents.auth_token, workspace_agents.auth_instance_id, workspace_agents.architecture, workspace_agents.environment_variables, workspace_agents.operating_system, workspace_agents.instance_metadata, workspace_agents.resource_metadata, workspace_agents.directory, workspace_agents.version, workspace_agents.last_connected_replica_id, workspace_agents.connection_timeout_seconds, workspace_agents.troubleshooting_url, workspace_agents.motd_file, workspace_agents.lifecycle_state, workspace_agents.expanded_directory, workspace_agents.logs_length, workspace_agents.logs_overflowed, workspace_agents.started_at, workspace_agents.ready_at, workspace_agents.subsystems, workspace_agents.display_apps, workspace_agents.api_version, workspace_agents.display_order, workspace_agents.parent_id, workspace_agents.api_key_scope, workspace_agents.deleted, - workspace_build_with_user.id, workspace_build_with_user.created_at, workspace_build_with_user.updated_at, workspace_build_with_user.workspace_id, workspace_build_with_user.template_version_id, workspace_build_with_user.build_number, workspace_build_with_user.transition, workspace_build_with_user.initiator_id, workspace_build_with_user.provisioner_state, workspace_build_with_user.job_id, workspace_build_with_user.deadline, workspace_build_with_user.reason, workspace_build_with_user.daily_cost, workspace_build_with_user.max_deadline, workspace_build_with_user.template_version_preset_id, workspace_build_with_user.has_ai_task, workspace_build_with_user.ai_task_sidebar_app_id, workspace_build_with_user.has_external_agent, workspace_build_with_user.initiator_by_avatar_url, workspace_build_with_user.initiator_by_username, workspace_build_with_user.initiator_by_name + workspace_build_with_user.id, workspace_build_with_user.created_at, workspace_build_with_user.updated_at, workspace_build_with_user.workspace_id, workspace_build_with_user.template_version_id, workspace_build_with_user.build_number, workspace_build_with_user.transition, workspace_build_with_user.initiator_id, workspace_build_with_user.provisioner_state, workspace_build_with_user.job_id, workspace_build_with_user.deadline, workspace_build_with_user.reason, workspace_build_with_user.daily_cost, workspace_build_with_user.max_deadline, workspace_build_with_user.template_version_preset_id, workspace_build_with_user.has_ai_task, workspace_build_with_user.ai_task_sidebar_app_id, workspace_build_with_user.has_external_agent, workspace_build_with_user.initiator_by_avatar_url, workspace_build_with_user.initiator_by_username, workspace_build_with_user.initiator_by_name, + tasks.id AS task_id FROM workspace_agents JOIN @@ -17584,6 +17585,10 @@ JOIN workspaces ON workspace_build_with_user.workspace_id = workspaces.id +LEFT JOIN + tasks +ON + tasks.workspace_id = workspaces.id WHERE -- This should only match 1 agent, so 1 returned row or 0. workspace_agents.auth_token = $1::uuid @@ -17607,6 +17612,7 @@ type GetWorkspaceAgentAndLatestBuildByAuthTokenRow struct { WorkspaceTable WorkspaceTable `db:"workspace_table" json:"workspace_table"` WorkspaceAgent WorkspaceAgent `db:"workspace_agent" json:"workspace_agent"` WorkspaceBuild WorkspaceBuild `db:"workspace_build" json:"workspace_build"` + TaskID uuid.NullUUID `db:"task_id" json:"task_id"` } func (q *sqlQuerier) GetWorkspaceAgentAndLatestBuildByAuthToken(ctx context.Context, authToken uuid.UUID) (GetWorkspaceAgentAndLatestBuildByAuthTokenRow, error) { @@ -17686,6 +17692,7 @@ func (q *sqlQuerier) GetWorkspaceAgentAndLatestBuildByAuthToken(ctx context.Cont &i.WorkspaceBuild.InitiatorByAvatarUrl, &i.WorkspaceBuild.InitiatorByUsername, &i.WorkspaceBuild.InitiatorByName, + &i.TaskID, ) return i, err } diff --git a/coderd/database/queries/workspaceagents.sql b/coderd/database/queries/workspaceagents.sql index cc59e96544..b60d1f2c88 100644 --- a/coderd/database/queries/workspaceagents.sql +++ b/coderd/database/queries/workspaceagents.sql @@ -285,7 +285,8 @@ WHERE SELECT sqlc.embed(workspaces), sqlc.embed(workspace_agents), - sqlc.embed(workspace_build_with_user) + sqlc.embed(workspace_build_with_user), + tasks.id AS task_id FROM workspace_agents JOIN @@ -300,6 +301,10 @@ JOIN workspaces ON workspace_build_with_user.workspace_id = workspaces.id +LEFT JOIN + tasks +ON + tasks.workspace_id = workspaces.id WHERE -- This should only match 1 agent, so 1 returned row or 0. workspace_agents.auth_token = @auth_token::uuid diff --git a/coderd/httpmw/workspaceagent.go b/coderd/httpmw/workspaceagent.go index 0ee231b2f5..d5f4e6fef2 100644 --- a/coderd/httpmw/workspaceagent.go +++ b/coderd/httpmw/workspaceagent.go @@ -118,6 +118,7 @@ func ExtractWorkspaceAgentAndLatestBuild(opts ExtractWorkspaceAgentAndLatestBuil OwnerID: row.WorkspaceTable.OwnerID, TemplateID: row.WorkspaceTable.TemplateID, VersionID: row.WorkspaceBuild.TemplateVersionID, + TaskID: row.TaskID, BlockUserData: row.WorkspaceAgent.APIKeyScope == database.AgentKeyScopeEnumNoUserData, }), ) diff --git a/coderd/rbac/scopes.go b/coderd/rbac/scopes.go index 5c8c803056..4e5babba29 100644 --- a/coderd/rbac/scopes.go +++ b/coderd/rbac/scopes.go @@ -18,6 +18,7 @@ type WorkspaceAgentScopeParams struct { OwnerID uuid.UUID TemplateID uuid.UUID VersionID uuid.UUID + TaskID uuid.NullUUID BlockUserData bool } @@ -42,6 +43,15 @@ func WorkspaceAgentScope(params WorkspaceAgentScopeParams) Scope { panic("failed to expand scope, this should never happen") } + // Include task in the allow list if the workspace has an associated task. + var extraAllowList []AllowListElement + if params.TaskID.Valid { + extraAllowList = append(extraAllowList, AllowListElement{ + Type: ResourceTask.Type, + ID: params.TaskID.UUID.String(), + }) + } + return Scope{ // TODO: We want to limit the role too to be extra safe. // Even though the allowlist blocks anything else, it is still good @@ -52,12 +62,12 @@ func WorkspaceAgentScope(params WorkspaceAgentScopeParams) Scope { // Limit the agent to only be able to access the singular workspace and // the template/version it was created from. Add additional resources here // as needed, but do not add more workspace or template resource ids. - AllowIDList: []AllowListElement{ + AllowIDList: append([]AllowListElement{ {Type: ResourceWorkspace.Type, ID: params.WorkspaceID.String()}, {Type: ResourceTemplate.Type, ID: params.TemplateID.String()}, {Type: ResourceTemplate.Type, ID: params.VersionID.String()}, {Type: ResourceUser.Type, ID: params.OwnerID.String()}, - }, + }, extraAllowList...), } } diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go index 23046dab28..1374d92dc4 100644 --- a/coderd/workspaceagents.go +++ b/coderd/workspaceagents.go @@ -461,67 +461,55 @@ func (api *API) enqueueAITaskStateNotification( return } - workspaceBuild, err := api.Database.GetLatestWorkspaceBuildByWorkspaceID(ctx, workspace.ID) - if err != nil { - api.Logger.Warn(ctx, "failed to get workspace build", slog.Error(err)) + if !workspace.TaskID.Valid { + // Workspace has no task ID, do nothing. return } - // Confirm Workspace Agent App is an AI Task - if workspaceBuild.HasAITask.Valid && workspaceBuild.HasAITask.Bool && - workspaceBuild.AITaskSidebarAppID.Valid && workspaceBuild.AITaskSidebarAppID.UUID == appID { - // Skip if the latest persisted state equals the new state (no new transition) - if len(latestAppStatus) > 0 && latestAppStatus[0].State == database.WorkspaceAppStatusState(newAppStatus) { - return - } + task, err := api.Database.GetTaskByID(ctx, workspace.TaskID.UUID) + if err != nil { + api.Logger.Warn(ctx, "failed to get task", slog.Error(err)) + return + } - // Skip the initial "Working" notification when task first starts. - // This is obvious to the user since they just created the task. - // We still notify on first "Idle" status and all subsequent transitions. - if len(latestAppStatus) == 0 && newAppStatus == codersdk.WorkspaceAppStatusStateWorking { - return - } + if !task.WorkspaceAppID.Valid || task.WorkspaceAppID.UUID != appID { + // Non-task app, do nothing. + return + } - // Use the task prompt as the "task" label, fallback to workspace name - parameters, err := api.Database.GetWorkspaceBuildParameters(ctx, workspaceBuild.ID) - if err != nil { - api.Logger.Warn(ctx, "failed to get workspace build parameters", slog.Error(err)) - return - } - taskName := workspace.Name - for _, param := range parameters { - if param.Name == codersdk.AITaskPromptParameterName { - taskName = param.Value - } - } + // Skip if the latest persisted state equals the new state (no new transition) + if len(latestAppStatus) > 0 && latestAppStatus[0].State == database.WorkspaceAppStatusState(newAppStatus) { + return + } - // As task prompt may be particularly long, truncate it to 160 characters for notifications. - if len(taskName) > 160 { - taskName = strutil.Truncate(taskName, 160, strutil.TruncateWithEllipsis, strutil.TruncateWithFullWords) - } + // Skip the initial "Working" notification when task first starts. + // This is obvious to the user since they just created the task. + // We still notify on first "Idle" status and all subsequent transitions. + if len(latestAppStatus) == 0 && newAppStatus == codersdk.WorkspaceAppStatusStateWorking { + return + } - if _, err := api.NotificationsEnqueuer.EnqueueWithData( - // nolint:gocritic // Need notifier actor to enqueue notifications - dbauthz.AsNotifier(ctx), - workspace.OwnerID, - notificationTemplate, - map[string]string{ - "task": taskName, - "workspace": workspace.Name, - }, - map[string]any{ - // Use a 1-minute bucketed timestamp to bypass per-day dedupe, - // allowing identical content to resend within the same day - // (but not more than once every 10s). - "dedupe_bypass_ts": api.Clock.Now().UTC().Truncate(time.Minute), - }, - "api-workspace-agent-app-status", - // Associate this notification with related entities - workspace.ID, workspace.OwnerID, workspace.OrganizationID, appID, - ); err != nil { - api.Logger.Warn(ctx, "failed to notify of task state", slog.Error(err)) - return - } + if _, err := api.NotificationsEnqueuer.EnqueueWithData( + // nolint:gocritic // Need notifier actor to enqueue notifications + dbauthz.AsNotifier(ctx), + workspace.OwnerID, + notificationTemplate, + map[string]string{ + "task": task.Name, + "workspace": workspace.Name, + }, + map[string]any{ + // Use a 1-minute bucketed timestamp to bypass per-day dedupe, + // allowing identical content to resend within the same day + // (but not more than once every 10s). + "dedupe_bypass_ts": api.Clock.Now().UTC().Truncate(time.Minute), + }, + "api-workspace-agent-app-status", + // Associate this notification with related entities + workspace.ID, workspace.OwnerID, workspace.OrganizationID, appID, + ); err != nil { + api.Logger.Warn(ctx, "failed to notify of task state", slog.Error(err)) + return } }