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

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

Author: mafredri

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{
+			workspaceBuilder := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
+				OrganizationID: ownerUser.OrganizationID,
+				OwnerID:        memberUser.ID,
+			}).Seed(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},
-				}
+				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 := 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()
+			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

@@ -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/database/queries/workspaceagents.sql

@@ -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/httpmw/workspaceagent.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/rbac/scopes.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
 	}
 }