fix: bound skill invocations per message

coderd/exp_chats.go

@@ -7491,6 +7491,7 @@ func createChatInputFromParts(
 
 	content := make([]codersdk.ChatMessagePart, 0, len(parts))
 	var pasteData map[uuid.UUID][]byte
+	skillInvocations := 0
 	for i, part := range parts {
 		switch strings.ToLower(strings.TrimSpace(string(part.Type))) {
 		case string(codersdk.ChatInputPartTypeText):
@@ -7638,6 +7639,13 @@ func createChatInputFromParts(
 					Detail:  fmt.Sprintf("%s[%d].skill_name must be a kebab-case skill name.", fieldName, i),
 				}
 			}
+			skillInvocations++
+			if skillInvocations > chatd.MaxSkillInvocationsPerMessage {
+				return nil, nil, &codersdk.Response{
+					Message: "Too many skill invocations.",
+					Detail:  fmt.Sprintf("%s can invoke at most %d skills.", fieldName, chatd.MaxSkillInvocationsPerMessage),
+				}
+			}
 			content = append(content, codersdk.ChatMessageSkillInvocation(part.SkillSource, part.SkillName))
 		default:
 			return nil, nil, &codersdk.Response{
@@ -7686,9 +7694,17 @@ func (api *API) writeChatFileError(ctx context.Context, rw http.ResponseWriter,
 	return true
 }
 
-// writeChatSkillInvocationError writes a 400 for a message that invokes a
-// skill the chat owner cannot use and reports whether it wrote a response.
+// writeChatSkillInvocationError writes a 400 for a message that invokes too
+// many skills or a skill the chat owner cannot use, and reports whether it
+// wrote a response.
 func writeChatSkillInvocationError(ctx context.Context, rw http.ResponseWriter, err error) bool {
+	if errors.Is(err, chatd.ErrTooManySkillInvocations) {
+		httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
+			Message: "Too many skill invocations.",
+			Detail:  err.Error(),
+		})
+		return true
+	}
 	var unavailable *chatd.SkillInvocationUnavailableError
 	if !errors.As(err, &unavailable) {
 		return false

coderd/exp_chats_skill_invocation_test.go

@@ -2,6 +2,7 @@ package coderd_test
 
 import (
 	"net/http"
+	"slices"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -9,6 +10,7 @@ import (
 	"github.com/coder/coder/v2/coderd/coderdtest"
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbauthz"
+	"github.com/coder/coder/v2/coderd/x/chatd"
 	"github.com/coder/coder/v2/coderd/x/chatd/chatprompt"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/testutil"
@@ -51,6 +53,14 @@ func TestChatSkillInvocation(t *testing.T) {
 		require.Equal(t, tt.detail, sdkErr.Detail, tt.name)
 	}
 
+	tooMany := slices.Repeat([]codersdk.ChatInputPart{
+		{Type: codersdk.ChatInputPartTypeSkillInvocation, SkillSource: codersdk.ChatSkillSourcePersonal, SkillName: "notes"},
+	}, chatd.MaxSkillInvocationsPerMessage+1)
+	_, err = client.CreateChat(ctx, codersdk.CreateChatRequest{OrganizationID: firstUser.OrganizationID, Content: tooMany})
+	tooManyErr := requireSDKError(t, err, http.StatusBadRequest)
+	require.Equal(t, "Too many skill invocations.", tooManyErr.Message)
+	require.Equal(t, "content can invoke at most 8 skills.", tooManyErr.Detail)
+
 	// The body is stored with the message but never returned by the API.
 	chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{
 		OrganizationID: firstUser.OrganizationID,

coderd/x/chatd/skill_invocation.go

@@ -17,6 +17,14 @@ import (
 	"github.com/coder/coder/v2/codersdk"
 )
 
+// MaxSkillInvocationsPerMessage bounds the skill-invocation parts in one
+// message, and with them the skill bodies a message stores.
+const MaxSkillInvocationsPerMessage = 8
+
+// ErrTooManySkillInvocations reports a message with more than
+// MaxSkillInvocationsPerMessage skill-invocation parts.
+var ErrTooManySkillInvocations = xerrors.Errorf("a message can invoke at most %d skills", MaxSkillInvocationsPerMessage)
+
 // SkillInvocationUnavailableError reports a skill-invocation part whose
 // skill the chat owner cannot use: it is missing, disabled, or not readable.
 type SkillInvocationUnavailableError struct {
@@ -57,9 +65,13 @@ func (p *Server) resolveSkillInvocations(
 	chat database.Chat,
 	parts []codersdk.ChatMessagePart,
 ) ([]codersdk.ChatMessagePart, error) {
-	if !slices.ContainsFunc(parts, isSkillInvocation) {
+	invocations := countSkillInvocations(parts)
+	if invocations == 0 {
 		return parts, nil
 	}
+	if invocations > MaxSkillInvocationsPerMessage {
+		return nil, ErrTooManySkillInvocations
+	}
 	loadOwner := sync.OnceValues(func() (rbac.Subject, error) {
 		return p.chatOwnerSubject(ctx, chat.OwnerID)
 	})
@@ -104,11 +116,22 @@ func (p *Server) resolveSkillInvocations(
 		return "", skillspkg.ErrSkillNotFound
 	}
 
+	// A message that repeats a skill loads it once.
+	type skillKey struct {
+		source codersdk.ChatSkillSource
+		name   string
+	}
+	bodies := make(map[skillKey]string, invocations)
 	resolved := slices.Clone(parts)
 	for i, part := range resolved {
 		if !isSkillInvocation(part) {
 			continue
 		}
+		key := skillKey{source: part.SkillSource, name: part.SkillName}
+		if body, ok := bodies[key]; ok {
+			resolved[i].SkillBody = body
+			continue
+		}
 		body, err := loadBody(part)
 		if errors.Is(err, skillspkg.ErrSkillNotFound) {
 			return nil, &SkillInvocationUnavailableError{Source: part.SkillSource, Name: part.SkillName}
@@ -116,11 +139,22 @@ func (p *Server) resolveSkillInvocations(
 		if err != nil {
 			return nil, xerrors.Errorf("load invoked skill %s/%s: %w", part.SkillSource, part.SkillName, err)
 		}
+		bodies[key] = body
 		resolved[i].SkillBody = body
 	}
 	return resolved, nil
 }
 
+func countSkillInvocations(parts []codersdk.ChatMessagePart) int {
+	count := 0
+	for _, part := range parts {
+		if isSkillInvocation(part) {
+			count++
+		}
+	}
+	return count
+}
+
 func isSkillInvocation(part codersdk.ChatMessagePart) bool {
 	return part.Type == codersdk.ChatMessagePartTypeSkillInvocation
 }

coderd/x/chatd/skill_invocation_internal_test.go

@@ -3,6 +3,7 @@ package chatd
 import (
 	"context"
 	"errors"
+	"slices"
 	"testing"
 
 	"github.com/google/uuid"
@@ -140,4 +141,29 @@ func TestResolveSkillInvocations(t *testing.T) {
 		})
 		requireUnavailable(t, err, "Skill workspace/missing is not available.")
 	})
+
+	t.Run("RepeatedSkillLoadsOnce", func(t *testing.T) {
+		t.Parallel()
+
+		ctrl := gomock.NewController(t)
+		db := dbmock.NewMockStore(ctrl)
+		chat := database.Chat{ID: uuid.New(), OwnerID: uuid.New()}
+		db.EXPECT().GetUserSkillByUserIDAndName(gomock.Any(), database.GetUserSkillByUserIDAndNameParams{UserID: chat.OwnerID, Name: "notes"}).
+			Return(database.Skill{Enabled: true, Content: "---\nname: notes\ndescription: Notes\n---\n\nTake notes.\n"}, nil).
+			Times(1)
+		server := newPinServer(t, db)
+		invocation := codersdk.ChatMessageSkillInvocation(codersdk.ChatSkillSourcePersonal, "notes")
+
+		resolved, err := server.resolveSkillInvocations(context.Background(), chat,
+			slices.Repeat([]codersdk.ChatMessagePart{invocation}, MaxSkillInvocationsPerMessage))
+		require.NoError(t, err)
+		for _, part := range resolved {
+			require.Equal(t, "Take notes.", part.SkillBody)
+		}
+
+		// One invocation over the cap is rejected before any skill loads.
+		_, err = server.resolveSkillInvocations(context.Background(), chat,
+			slices.Repeat([]codersdk.ChatMessagePart{invocation}, MaxSkillInvocationsPerMessage+1))
+		require.ErrorIs(t, err, ErrTooManySkillInvocations)
+	})
 }