refactor: use task data model for notifications (#20590)

Updates coder/internal#973
Updates coder/internal#974

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

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
 }

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

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,
 				}),
 			)

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...),
 	}
 }
 

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
 	}
 }