Files
cursor-byok/internal/backend/forwarder/prompt_guard.go
T

359 lines
12 KiB
Go

package forwarder
import (
"fmt"
"sort"
"strings"
"unicode/utf8"
"google.golang.org/protobuf/proto"
"cursor/gen/agentv1"
)
const (
promptGuardUserTextChars = 32000
promptGuardUserRichTextChars = 32000
promptGuardSubagentReminderChars = 12000
promptGuardSelectedFileChars = 16000
promptGuardSelectedFilesTotalChars = 64000
promptGuardSelectedFilesMaxCount = 12
promptGuardCursorCommandNameChars = 256
promptGuardCursorCommandChars = 12000
promptGuardCursorCommandsTotalChars = 32000
promptGuardCursorCommandsMaxCount = 8
promptGuardRequestFileChars = 16000
promptGuardRequestFilesTotalChars = 64000
promptGuardRequestFilesMaxCount = 12
promptGuardRuleChars = 6000
promptGuardRulesTotalChars = 24000
promptGuardRulesMaxCount = 40
promptGuardSkillDescriptionChars = 800
promptGuardSkillDescriptionsTotal = 16000
promptGuardSkillDescriptorsMaxCount = 32
promptGuardAgentSkillContentChars = 6000
promptGuardAgentSkillsMaxCount = 16
promptGuardRealtimeTextChars = 12000
promptGuardCompiledMessageChars = 120000
)
func normalizeUserMessageForStorage(userMessage *agentv1.UserMessage) *agentv1.UserMessage {
if userMessage == nil {
return nil
}
cloned, ok := proto.Clone(userMessage).(*agentv1.UserMessage)
if !ok || cloned == nil {
return userMessage
}
cloned.Text = truncatePromptGuardText("user_message.text", cloned.GetText(), promptGuardUserTextChars)
if richText := strings.TrimSpace(cloned.GetRichText()); richText != "" {
cloned.RichText = stringPtr(truncatePromptGuardText("user_message.rich_text", richText, promptGuardUserRichTextChars))
}
if reminder := strings.TrimSpace(cloned.GetSubagentSystemReminder()); reminder != "" {
cloned.SubagentSystemReminder = stringPtr(truncatePromptGuardText("user_message.subagent_system_reminder", reminder, promptGuardSubagentReminderChars))
}
cloned.SelectedContext = guardSelectedContext(cloned.GetSelectedContext())
return cloned
}
func guardRequestContextForStorage(requestContext *agentv1.RequestContext) {
if requestContext == nil {
return
}
requestContext.Rules = guardCursorRules(requestContext.GetRules())
requestContext.FileContents = guardStringMap(
requestContext.GetFileContents(),
"request_context.file_contents",
promptGuardRequestFileChars,
promptGuardRequestFilesTotalChars,
promptGuardRequestFilesMaxCount,
)
if summary := strings.TrimSpace(requestContext.GetUserIntentSummary()); summary != "" {
requestContext.UserIntentSummary = stringPtr(truncatePromptGuardText("request_context.user_intent_summary", summary, promptGuardRealtimeTextChars))
}
if hooks := strings.TrimSpace(requestContext.GetHooksAdditionalContext()); hooks != "" {
requestContext.HooksAdditionalContext = stringPtr(truncatePromptGuardText("request_context.hooks_additional_context", hooks, promptGuardRealtimeTextChars))
}
if commit := strings.TrimSpace(requestContext.GetCommitAttributionMessage()); commit != "" {
requestContext.CommitAttributionMessage = stringPtr(truncatePromptGuardText("request_context.commit_attribution_message", commit, promptGuardRealtimeTextChars))
}
if pr := strings.TrimSpace(requestContext.GetPrAttributionMessage()); pr != "" {
requestContext.PrAttributionMessage = stringPtr(truncatePromptGuardText("request_context.pr_attribution_message", pr, promptGuardRealtimeTextChars))
}
}
func guardCompiledConversationForProvider(compiled CompiledConversation) CompiledConversation {
for index := range compiled.Messages {
message := &compiled.Messages[index]
if strings.TrimSpace(message.Role) == "system" {
continue
}
if strings.TrimSpace(message.Content) != "" {
message.Content = truncatePromptGuardText("compiled."+firstNonEmpty(strings.TrimSpace(message.Role), "message"), message.Content, promptGuardCompiledMessageChars)
}
for partIndex := range message.ContentParts {
if strings.TrimSpace(message.ContentParts[partIndex].Text) == "" {
continue
}
message.ContentParts[partIndex].Text = truncatePromptGuardText("compiled.content_part", message.ContentParts[partIndex].Text, promptGuardCompiledMessageChars)
}
}
return compiled
}
func guardSelectedContext(selectedContext *agentv1.SelectedContext) *agentv1.SelectedContext {
if selectedContext == nil {
return nil
}
cloned, ok := proto.Clone(selectedContext).(*agentv1.SelectedContext)
if !ok || cloned == nil {
return selectedContext
}
cloned.Files = guardSelectedFiles(cloned.GetFiles())
cloned.CursorCommands = guardSelectedCursorCommands(cloned.GetCursorCommands())
cloned.SelectedSkills = guardAgentSkills(cloned.GetSelectedSkills())
cloned.ExtraContext = guardStringSlice(cloned.GetExtraContext(), "selected_context.extra_context", promptGuardRealtimeTextChars, promptGuardRealtimeTextChars, promptGuardAgentSkillsMaxCount)
return cloned
}
func guardSelectedCursorCommands(commands []*agentv1.SelectedCursorCommand) []*agentv1.SelectedCursorCommand {
if len(commands) == 0 {
return nil
}
result := make([]*agentv1.SelectedCursorCommand, 0, minInt(len(commands), promptGuardCursorCommandsMaxCount))
remaining := promptGuardCursorCommandsTotalChars
for _, command := range commands {
if command == nil || len(result) >= promptGuardCursorCommandsMaxCount {
continue
}
content := strings.TrimSpace(command.GetContent())
if content == "" {
continue
}
limit := minInt(promptGuardCursorCommandChars, remaining)
if limit <= 0 {
break
}
cloned, ok := proto.Clone(command).(*agentv1.SelectedCursorCommand)
if !ok || cloned == nil {
continue
}
cloned.Name = truncatePromptGuardText("selected_context.cursor_commands.name", strings.TrimSpace(cloned.GetName()), promptGuardCursorCommandNameChars)
cloned.Content = truncatePromptGuardText("selected_context.cursor_commands.content", content, limit)
remaining -= promptGuardRuneCount(cloned.GetContent())
result = append(result, cloned)
}
return result
}
func guardSelectedFiles(files []*agentv1.SelectedFile) []*agentv1.SelectedFile {
if len(files) == 0 {
return nil
}
result := make([]*agentv1.SelectedFile, 0, minInt(len(files), promptGuardSelectedFilesMaxCount))
remaining := promptGuardSelectedFilesTotalChars
for _, file := range files {
if file == nil || len(result) >= promptGuardSelectedFilesMaxCount {
continue
}
content := strings.TrimSpace(file.GetContent())
if content == "" {
continue
}
limit := minInt(promptGuardSelectedFileChars, remaining)
if limit <= 0 {
break
}
cloned, ok := proto.Clone(file).(*agentv1.SelectedFile)
if !ok || cloned == nil {
continue
}
cloned.Content = truncatePromptGuardText("selected_context.files.content", content, limit)
remaining -= promptGuardRuneCount(cloned.GetContent())
result = append(result, cloned)
}
return result
}
func guardCursorRules(rules []*agentv1.CursorRule) []*agentv1.CursorRule {
if len(rules) == 0 {
return nil
}
result := make([]*agentv1.CursorRule, 0, minInt(len(rules), promptGuardRulesMaxCount))
remaining := promptGuardRulesTotalChars
for _, rule := range rules {
if rule == nil || len(result) >= promptGuardRulesMaxCount {
continue
}
content := strings.TrimSpace(rule.GetContent())
if content == "" {
continue
}
limit := minInt(promptGuardRuleChars, remaining)
if limit <= 0 {
break
}
cloned, ok := proto.Clone(rule).(*agentv1.CursorRule)
if !ok || cloned == nil {
continue
}
cloned.Content = truncatePromptGuardText("request_context.rules.content", content, limit)
remaining -= promptGuardRuneCount(cloned.GetContent())
result = append(result, cloned)
}
return result
}
func guardSkillDescriptors(descriptors []*agentv1.SkillDescriptor) []*agentv1.SkillDescriptor {
if len(descriptors) == 0 {
return nil
}
result := make([]*agentv1.SkillDescriptor, 0, minInt(len(descriptors), promptGuardSkillDescriptorsMaxCount))
remaining := promptGuardSkillDescriptionsTotal
for _, descriptor := range descriptors {
if descriptor == nil || len(result) >= promptGuardSkillDescriptorsMaxCount {
continue
}
description := strings.TrimSpace(descriptor.GetDescription())
if description == "" {
continue
}
limit := minInt(promptGuardSkillDescriptionChars, remaining)
if limit <= 0 {
break
}
cloned, ok := proto.Clone(descriptor).(*agentv1.SkillDescriptor)
if !ok || cloned == nil {
continue
}
cloned.Description = truncatePromptGuardText("skill.description", description, limit)
remaining -= promptGuardRuneCount(cloned.GetDescription())
result = append(result, cloned)
}
return result
}
func guardAgentSkills(skills []*agentv1.AgentSkill) []*agentv1.AgentSkill {
if len(skills) == 0 {
return nil
}
result := make([]*agentv1.AgentSkill, 0, minInt(len(skills), promptGuardAgentSkillsMaxCount))
for _, skill := range skills {
if skill == nil || len(result) >= promptGuardAgentSkillsMaxCount {
continue
}
cloned, ok := proto.Clone(skill).(*agentv1.AgentSkill)
if !ok || cloned == nil {
continue
}
if content := strings.TrimSpace(cloned.GetContent()); content != "" {
cloned.Content = truncatePromptGuardText("agent_skill.content", content, promptGuardAgentSkillContentChars)
}
if description := strings.TrimSpace(cloned.GetDescription()); description != "" {
cloned.Description = truncatePromptGuardText("agent_skill.description", description, promptGuardSkillDescriptionChars)
}
result = append(result, cloned)
}
return result
}
func guardStringMap(input map[string]string, label string, itemLimit int, totalLimit int, maxItems int) map[string]string {
if len(input) == 0 {
return nil
}
keys := make([]string, 0, len(input))
values := make(map[string]string, len(input))
for key, value := range input {
trimmedKey := strings.TrimSpace(key)
trimmedValue := strings.TrimSpace(value)
if trimmedKey == "" || trimmedValue == "" {
continue
}
if _, exists := values[trimmedKey]; exists {
continue
}
keys = append(keys, trimmedKey)
values[trimmedKey] = trimmedValue
}
sort.Strings(keys)
result := make(map[string]string, minInt(len(keys), maxItems))
remaining := totalLimit
for _, key := range keys {
if len(result) >= maxItems {
break
}
content := values[key]
if content == "" {
continue
}
limit := minInt(itemLimit, remaining)
if limit <= 0 {
break
}
result[key] = truncatePromptGuardText(label, content, limit)
remaining -= promptGuardRuneCount(result[key])
}
if len(result) == 0 {
return nil
}
return result
}
func guardStringSlice(input []string, label string, itemLimit int, totalLimit int, maxItems int) []string {
if len(input) == 0 {
return nil
}
result := make([]string, 0, minInt(len(input), maxItems))
remaining := totalLimit
for _, item := range input {
if len(result) >= maxItems {
break
}
content := strings.TrimSpace(item)
if content == "" {
continue
}
limit := minInt(itemLimit, remaining)
if limit <= 0 {
break
}
truncated := truncatePromptGuardText(label, content, limit)
remaining -= promptGuardRuneCount(truncated)
result = append(result, truncated)
}
return result
}
func truncatePromptGuardText(label string, text string, limit int) string {
if limit <= 0 {
return ""
}
if promptGuardRuneCount(text) <= limit {
return text
}
runes := []rune(text)
notice := fmt.Sprintf("\n\n[truncated: %s exceeded %d chars; kept head and tail from %d chars]\n\n", strings.TrimSpace(label), limit, len(runes))
noticeRunes := []rune(notice)
keep := limit - len(noticeRunes)
if keep <= 0 {
return string(runes[:limit])
}
head := keep * 2 / 3
tail := keep - head
if tail <= 0 {
return string(runes[:head]) + notice
}
return string(runes[:head]) + notice + string(runes[len(runes)-tail:])
}
func promptGuardRuneCount(text string) int {
return utf8.RuneCountInString(text)
}
func minInt(a int, b int) int {
if a < b {
return a
}
return b
}