Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 27 additions & 2 deletions coderd/aitasks_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,10 @@ import (
"io"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"unicode/utf8"

"github.com/google/uuid"
"github.com/stretchr/testify/assert"
Expand Down Expand Up @@ -977,6 +979,7 @@ func TestTasksNotification(t *testing.T) {
isAITask bool
isNotificationSent bool
notificationTemplate uuid.UUID
taskPrompt string
}{
// Should not send a notification when the agent app is not an AI task.
{
Expand All @@ -985,6 +988,7 @@ func TestTasksNotification(t *testing.T) {
newAppStatus: codersdk.WorkspaceAppStatusStateWorking,
isAITask: false,
isNotificationSent: false,
taskPrompt: "NoAITask",
},
// Should not send a notification when the new app status is neither 'Working' nor 'Idle'.
{
Expand All @@ -993,6 +997,7 @@ func TestTasksNotification(t *testing.T) {
newAppStatus: codersdk.WorkspaceAppStatusStateComplete,
isAITask: true,
isNotificationSent: false,
taskPrompt: "NonNotifiedState",
},
// Should not send a notification when the new app status equals the latest status (Working).
{
Expand All @@ -1001,6 +1006,7 @@ func TestTasksNotification(t *testing.T) {
newAppStatus: codersdk.WorkspaceAppStatusStateWorking,
isAITask: true,
isNotificationSent: false,
taskPrompt: "NonNotifiedTransition",
},
// Should send TemplateTaskWorking when the AI task transitions to 'Working'.
{
Expand All @@ -1010,6 +1016,7 @@ func TestTasksNotification(t *testing.T) {
isAITask: true,
isNotificationSent: true,
notificationTemplate: notifications.TemplateTaskWorking,
taskPrompt: "TemplateTaskWorking",
},
// Should send TemplateTaskWorking when the AI task transitions to 'Working' from 'Idle'.
{
Expand All @@ -1022,6 +1029,7 @@ func TestTasksNotification(t *testing.T) {
isAITask: true,
isNotificationSent: true,
notificationTemplate: notifications.TemplateTaskWorking,
taskPrompt: "TemplateTaskWorkingFromIdle",
},
// Should send TemplateTaskIdle when the AI task transitions to 'Idle'.
{
Expand All @@ -1031,6 +1039,17 @@ func TestTasksNotification(t *testing.T) {
isAITask: true,
isNotificationSent: true,
notificationTemplate: notifications.TemplateTaskIdle,
taskPrompt: "TemplateTaskIdle",
},
// Long task prompts should be truncated to 160 characters.
{
name: "LongTaskPrompt",
latestAppStatuses: []codersdk.WorkspaceAppStatusState{codersdk.WorkspaceAppStatusStateWorking},
newAppStatus: codersdk.WorkspaceAppStatusStateIdle,
isAITask: true,
isNotificationSent: true,
notificationTemplate: notifications.TemplateTaskIdle,
taskPrompt: "This is a very long task prompt that should be truncated to 160 characters. Lorem ipsum dolor sit amet, consectetur adipiscing elit. Sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.",
},
} {
t.Run(tc.name, func(t *testing.T) {
Expand Down Expand Up @@ -1067,7 +1086,7 @@ func TestTasksNotification(t *testing.T) {
}).Seed(workspaceBuildSeed).Params(database.WorkspaceBuildParameter{
WorkspaceBuildID: workspaceBuildID,
Name: codersdk.AITaskPromptParameterName,
Value: "task prompt",
Value: tc.taskPrompt,
}).WithAgent(func(agent []*proto.Agent) []*proto.Agent {
agent[0].Apps = []*proto.App{{
Id: workspaceAgentAppID.String(),
Expand Down Expand Up @@ -1115,7 +1134,13 @@ 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)
require.Equal(t, "task prompt", sent[0].Labels["task"])
// 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, workspace.Name, sent[0].Labels["workspace"])
} else {
// Then: No notification is sent
Expand Down
55 changes: 52 additions & 3 deletions coderd/util/strings/strings.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,15 +23,64 @@ func JoinWithConjunction(s []string) string {
)
}

// Truncate returns the first n characters of s.
func Truncate(s string, n int) string {
type TruncateOption int

func (o TruncateOption) String() string {
switch o {
case TruncateWithEllipsis:
return "TruncateWithEllipsis"
case TruncateWithFullWords:
return "TruncateWithFullWords"
default:
return fmt.Sprintf("TruncateOption(%d)", o)
}
}

const (
// TruncateWithEllipsis adds a Unicode ellipsis character to the end of the string.
TruncateWithEllipsis TruncateOption = 1 << 0
// TruncateWithFullWords ensures that words are not split in the middle.
// As a special case, if there is no word boundary, the string is truncated.
TruncateWithFullWords TruncateOption = 1 << 1
)

// Truncate truncates s to n characters.
// Additional behaviors can be specified using TruncateOptions.
func Truncate(s string, n int, opts ...TruncateOption) string {
var options TruncateOption
for _, opt := range opts {
options |= opt
}
if n < 1 {
return ""
}
if len(s) <= n {
return s
}
return s[:n]

maxLen := n
if options&TruncateWithEllipsis != 0 {
maxLen--
}
var sb strings.Builder
// If we need to truncate to full words, find the last word boundary before n.
if options&TruncateWithFullWords != 0 {
lastWordBoundary := strings.LastIndexFunc(s[:maxLen], unicode.IsSpace)
if lastWordBoundary < 0 {
// We cannot find a word boundary. At this point, we'll truncate the string.
// It's better than nothing.
_, _ = sb.WriteString(s[:maxLen])
} else { // lastWordBoundary <= maxLen
_, _ = sb.WriteString(s[:lastWordBoundary])
}
} else {
_, _ = sb.WriteString(s[:maxLen])
}

if options&TruncateWithEllipsis != 0 {
_, _ = sb.WriteString("…")
}
return sb.String()
}

var bmPolicy = bluemonday.StrictPolicy()
Expand Down
47 changes: 39 additions & 8 deletions coderd/util/strings/strings_test.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
package strings_test

import (
"fmt"
"testing"

"github.com/stretchr/testify/assert"
Expand All @@ -23,17 +24,47 @@ func TestTruncate(t *testing.T) {
s string
n int
expected string
options []strings.TruncateOption
}{
{"foo", 4, "foo"},
{"foo", 3, "foo"},
{"foo", 2, "fo"},
{"foo", 1, "f"},
{"foo", 0, ""},
{"foo", -1, ""},
{"foo", 4, "foo", nil},
{"foo", 3, "foo", nil},
{"foo", 2, "fo", nil},
{"foo", 1, "f", nil},
{"foo", 0, "", nil},
{"foo", -1, "", nil},
{"foo bar", 7, "foo bar", []strings.TruncateOption{strings.TruncateWithEllipsis}},
{"foo bar", 6, "foo b…", []strings.TruncateOption{strings.TruncateWithEllipsis}},
{"foo bar", 5, "foo …", []strings.TruncateOption{strings.TruncateWithEllipsis}},
{"foo bar", 4, "foo…", []strings.TruncateOption{strings.TruncateWithEllipsis}},
{"foo bar", 3, "fo…", []strings.TruncateOption{strings.TruncateWithEllipsis}},
{"foo bar", 2, "f…", []strings.TruncateOption{strings.TruncateWithEllipsis}},
{"foo bar", 1, "…", []strings.TruncateOption{strings.TruncateWithEllipsis}},
{"foo bar", 0, "", []strings.TruncateOption{strings.TruncateWithEllipsis}},
{"foo bar", 7, "foo bar", []strings.TruncateOption{strings.TruncateWithFullWords}},
{"foo bar", 6, "foo", []strings.TruncateOption{strings.TruncateWithFullWords}},
{"foo bar", 5, "foo", []strings.TruncateOption{strings.TruncateWithFullWords}},
{"foo bar", 4, "foo", []strings.TruncateOption{strings.TruncateWithFullWords}},
{"foo bar", 3, "foo", []strings.TruncateOption{strings.TruncateWithFullWords}},
{"foo bar", 2, "fo", []strings.TruncateOption{strings.TruncateWithFullWords}},
{"foo bar", 1, "f", []strings.TruncateOption{strings.TruncateWithFullWords}},
{"foo bar", 0, "", []strings.TruncateOption{strings.TruncateWithFullWords}},
{"foo bar", 7, "foo bar", []strings.TruncateOption{strings.TruncateWithFullWords, strings.TruncateWithEllipsis}},
{"foo bar", 6, "foo…", []strings.TruncateOption{strings.TruncateWithFullWords, strings.TruncateWithEllipsis}},
{"foo bar", 5, "foo…", []strings.TruncateOption{strings.TruncateWithFullWords, strings.TruncateWithEllipsis}},
{"foo bar", 4, "foo…", []strings.TruncateOption{strings.TruncateWithFullWords, strings.TruncateWithEllipsis}},
{"foo bar", 3, "fo…", []strings.TruncateOption{strings.TruncateWithFullWords, strings.TruncateWithEllipsis}},
{"foo bar", 2, "f…", []strings.TruncateOption{strings.TruncateWithFullWords, strings.TruncateWithEllipsis}},
{"foo bar", 1, "…", []strings.TruncateOption{strings.TruncateWithFullWords, strings.TruncateWithEllipsis}},
{"foo bar", 0, "", []strings.TruncateOption{strings.TruncateWithFullWords, strings.TruncateWithEllipsis}},
{"This is a very long task prompt that should be truncated to 160 characters. Lorem ipsum dolor sit amet, consectetur adipiscing elit. Sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.", 160, "This is a very long task prompt that should be truncated to 160 characters. Lorem ipsum dolor sit amet, consectetur adipiscing elit. Sed do eiusmod tempor…", []strings.TruncateOption{strings.TruncateWithFullWords, strings.TruncateWithEllipsis}},
} {
t.Run(tt.expected, func(t *testing.T) {
tName := fmt.Sprintf("%s_%d", tt.s, tt.n)
for _, opt := range tt.options {
tName += fmt.Sprintf("_%v", opt)
}
t.Run(tName, func(t *testing.T) {
t.Parallel()
actual := strings.Truncate(tt.s, tt.n)
actual := strings.Truncate(tt.s, tt.n, tt.options...)
require.Equal(t, tt.expected, actual)
})
}
Expand Down
5 changes: 5 additions & 0 deletions coderd/workspaceagents.go
Original file line number Diff line number Diff line change
Expand Up @@ -484,6 +484,11 @@ func (api *API) enqueueAITaskStateNotification(
}
}

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

if _, err := api.NotificationsEnqueuer.EnqueueWithData(
// nolint:gocritic // Need notifier actor to enqueue notifications
dbauthz.AsNotifier(ctx),
Expand Down
Loading