Skip to content
Open
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
26 changes: 22 additions & 4 deletions runtime/ai/router_agent.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import (
"regexp"
"slices"
"strings"
"unicode/utf8"

"github.com/google/jsonschema-go/jsonschema"
"github.com/modelcontextprotocol/go-sdk/mcp"
Expand Down Expand Up @@ -234,10 +235,10 @@ func promptToTitle(message string) string {
title := whitespaceRegexp.ReplaceAllString(message, " ")
title = strings.TrimSpace(title)

// Truncate to 50 characters.
if len(title) > 50 {
title = title[:47] + "..."
}
// Truncate to 50 bytes on a UTF-8 rune boundary. Byte-wise title[:47]
// splits multi-byte runes (e.g. Chinese) and stores invalid UTF-8, which
// then fails protobuf marshaling in ListConversations.
title = truncateUTF8(title, 50)

// Fallback title if empty.
if title == "" {
Expand All @@ -246,6 +247,23 @@ func promptToTitle(message string) string {
return title
}

// truncateUTF8 shortens s to at most maxBytes, appending "..." when truncated.
// The cut is always on a UTF-8 rune boundary so the result stays valid UTF-8.
func truncateUTF8(s string, maxBytes int) string {

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

truncateUTF8 is called exactly once, with a constant, which also makes the maxBytes <= len(ellipsis) branch unreachable. Per the repo's convention against single-use utility functions, the back-off could be inlined in promptToTitle as a walk from index 47 while !utf8.RuneStart(title[i]), which avoids re-validating the whole prefix on every iteration. Either form relies on the prompt being valid UTF-8 on entry, which holds for all current entry points.

if len(s) <= maxBytes {
return s
}
const ellipsis = "..."
if maxBytes <= len(ellipsis) {
return ellipsis[:maxBytes]
}
s = s[:maxBytes-len(ellipsis)]
for s != "" && !utf8.ValidString(s) {
s = s[:len(s)-1]
}
return s + ellipsis
}

// mapAgentErr maps common agent errors to more user-friendly messages.
//
// NOTE: For context errors, it does not include the underlying error to keep messages clean.
Expand Down
34 changes: 34 additions & 0 deletions runtime/ai/router_agent_title_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
package ai

import (
"strings"
"testing"
"unicode/utf8"

"github.com/stretchr/testify/require"
)

func TestPromptToTitle_multibyteTruncationStaysValidUTF8(t *testing.T) {
// Live Rill_UO prompt whose previous byte-wise cut (title[:47]+"...") split 有.
prompt := "今年乘用车增长快于大盘的区域都有哪些"
require.Greater(t, len(prompt), 50)

oldCut := prompt[:47] + "..."
require.False(t, utf8.ValidString(oldCut), "the previous cut point must be invalid UTF-8 (the bug)")

title := promptToTitle(prompt)
require.True(t, utf8.ValidString(title))
require.LessOrEqual(t, len(title), 50)
require.True(t, strings.HasSuffix(title, "..."))
require.Equal(t, "今年乘用车增长快于大盘的区域都...", title)
}

func TestPromptToTitle_asciiStillTruncatesAt50Bytes(t *testing.T) {
title := promptToTitle(strings.Repeat("a", 80))
require.Equal(t, strings.Repeat("a", 47)+"...", title)
require.True(t, utf8.ValidString(title))
}

func TestPromptToTitle_emptyFallsBack(t *testing.T) {
require.Equal(t, "New Conversation", promptToTitle(" \n\t "))
}
3 changes: 2 additions & 1 deletion runtime/server/chat.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import (
"fmt"
"io"
"net/http"
"strings"
"time"

aiv1 "github.com/rilldata/rill/proto/gen/rill/ai/v1"
Expand Down Expand Up @@ -601,7 +602,7 @@ func sessionToPB(s *drivers.AISession, messages []*runtimev1.Message) *runtimev1
return &runtimev1.Conversation{
Id: s.ID,
OwnerId: s.OwnerID,
Title: s.Title,
Title: strings.ToValidUTF8(s.Title, "\uFFFD"),
UserAgent: s.UserAgent,
CreatedOn: timestamppb.New(s.CreatedOn),
UpdatedOn: timestamppb.New(s.UpdatedOn),
Expand Down
53 changes: 53 additions & 0 deletions runtime/server/chat_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,11 +4,13 @@ import (
"context"
"testing"
"time"
"unicode/utf8"

aiv1 "github.com/rilldata/rill/proto/gen/rill/ai/v1"
runtimev1 "github.com/rilldata/rill/proto/gen/rill/runtime/v1"
"github.com/rilldata/rill/runtime"
"github.com/rilldata/rill/runtime/ai"
"github.com/rilldata/rill/runtime/drivers"
"github.com/rilldata/rill/runtime/pkg/activity"
"github.com/rilldata/rill/runtime/pkg/ratelimit"
"github.com/rilldata/rill/runtime/server"
Expand All @@ -17,6 +19,8 @@ import (
"github.com/rilldata/rill/runtime/testruntime/testmode"
"github.com/stretchr/testify/require"
"go.uber.org/zap"
"google.golang.org/protobuf/encoding/protojson"
"google.golang.org/protobuf/proto"
"google.golang.org/protobuf/types/known/timestamppb"
)

Expand Down Expand Up @@ -478,3 +482,52 @@ func TestListTools(t *testing.T) {
require.NotEmpty(t, tool.OutputSchema)
}
}

func TestListConversationsInvalidUTF8Title(t *testing.T) {
rt, instanceID := testruntime.NewInstance(t)
srv, err := server.NewServer(context.Background(), &server.Options{}, rt, zap.NewNop(), ratelimit.NewNoop(), activity.NewNoopClient())
require.NoError(t, err)

// Bytes stored by the old promptToTitle byte-wise cut of
// "今年乘用车增长快于大盘的区域都有哪些" (splits 有, then appends "...").
invalidTitle := string([]byte{
0xe4, 0xbb, 0x8a, 0xe5, 0xb9, 0xb4, 0xe4, 0xb9, 0x98, 0xe7, 0x94, 0xa8,
0xe8, 0xbd, 0xa6, 0xe5, 0xa2, 0x9e, 0xe9, 0x95, 0xbf, 0xe5, 0xbf, 0xab,
0xe4, 0xba, 0x8e, 0xe5, 0xa4, 0xa7, 0xe7, 0x9b, 0x98, 0xe7, 0x9a, 0x84,
0xe5, 0x8c, 0xba, 0xe5, 0x9f, 0x9f, 0xe9, 0x83, 0xbd, 0xe6, 0x9c, 0x2e,
0x2e, 0x2e,
})

catalog, release, err := rt.Catalog(t.Context(), instanceID)
require.NoError(t, err)
defer release()

now := time.Now()
err = catalog.InsertAISession(t.Context(), &drivers.AISession{
ID: "invalid-utf8-title",
InstanceID: instanceID,
OwnerID: "foo",
Title: invalidTitle,
UserAgent: "rill/test",
CreatedOn: now,
UpdatedOn: now,
})
require.NoError(t, err)

ctx := auth.WithClaims(t.Context(), &runtime.SecurityClaims{
UserID: "foo",
Permissions: []runtime.Permission{runtime.UseAI},
})
list, err := srv.ListConversations(ctx, &runtimev1.ListConversationsRequest{
InstanceId: instanceID,
})
require.NoError(t, err)
require.Len(t, list.Conversations, 1)
require.Equal(t, "invalid-utf8-title", list.Conversations[0].Id)
require.True(t, utf8.ValidString(list.Conversations[0].Title))

_, err = proto.Marshal(list)
require.NoError(t, err)
_, err = protojson.Marshal(list)
require.NoError(t, err)
}
Loading