You've already forked agentic-coding-workflow
feat(chat): MessageService Send/Stream/Cancel/Retry with attachment + capability + history truncation
This commit is contained in:
@@ -0,0 +1,935 @@
|
||||
package chat
|
||||
|
||||
// NOTE: Integration-level test TestSendMessage_ConcurrentPending_RejectedByPGUniqueIndex
|
||||
// requires a real Postgres database (the partial unique index messages_one_pending_per_conv
|
||||
// is enforced at DB layer). This is covered in T19 integration tests.
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
|
||||
"github.com/yan1h/agent-coding-workflow/internal/infra/llm"
|
||||
"github.com/yan1h/agent-coding-workflow/internal/infra/storage"
|
||||
)
|
||||
|
||||
// ===== msgFakeRepo =====
|
||||
// Embeds panicRepo and overrides only the methods needed by messageService.
|
||||
|
||||
type msgFakeRepo struct {
|
||||
panicRepo
|
||||
|
||||
mu sync.Mutex
|
||||
|
||||
conversations map[uuid.UUID]*Conversation
|
||||
messages map[int64]*Message
|
||||
attachments map[uuid.UUID]*Attachment
|
||||
nextMsgID int64
|
||||
|
||||
// call tracking
|
||||
insertedMessages []*Message
|
||||
deletedMessageIDs []int64
|
||||
errorUpdates []errorUpdate
|
||||
attachToMessageCalls []attachCall
|
||||
withTxCalled bool
|
||||
}
|
||||
|
||||
type errorUpdate struct {
|
||||
id int64
|
||||
msg *string
|
||||
}
|
||||
|
||||
type attachCall struct {
|
||||
ids []uuid.UUID
|
||||
messageID int64
|
||||
}
|
||||
|
||||
func newMsgFakeRepo() *msgFakeRepo {
|
||||
return &msgFakeRepo{
|
||||
conversations: make(map[uuid.UUID]*Conversation),
|
||||
messages: make(map[int64]*Message),
|
||||
attachments: make(map[uuid.UUID]*Attachment),
|
||||
nextMsgID: 1,
|
||||
}
|
||||
}
|
||||
|
||||
func (r *msgFakeRepo) addConversation(c *Conversation) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.conversations[c.ID] = c
|
||||
}
|
||||
|
||||
func (r *msgFakeRepo) addAttachment(a *Attachment) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.attachments[a.ID] = a
|
||||
}
|
||||
|
||||
func (r *msgFakeRepo) addMessage(m *Message) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.messages[m.ID] = m
|
||||
}
|
||||
|
||||
func (r *msgFakeRepo) GetConversation(_ context.Context, id uuid.UUID) (*Conversation, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
c, ok := r.conversations[id]
|
||||
if !ok {
|
||||
return nil, errs.New(errs.CodeChatConversationNotFound, "not found")
|
||||
}
|
||||
return c, nil
|
||||
}
|
||||
|
||||
func (r *msgFakeRepo) GetLLMModel(_ context.Context, id uuid.UUID) (*LLMModel, error) {
|
||||
// Returns a model with ContextWindow=4096, MaxOutputTokens=1024 by default.
|
||||
return &LLMModel{
|
||||
ID: id,
|
||||
EndpointID: uuid.New(),
|
||||
ContextWindow: 4096,
|
||||
MaxOutputTokens: 1024,
|
||||
Capabilities: ModelCapabilities{Image: true, PDF: true},
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (r *msgFakeRepo) GetLLMEndpoint(_ context.Context, id uuid.UUID) (*LLMEndpoint, error) {
|
||||
return &LLMEndpoint{ID: id}, nil
|
||||
}
|
||||
|
||||
func (r *msgFakeRepo) GetMessage(_ context.Context, id int64) (*Message, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
m, ok := r.messages[id]
|
||||
if !ok {
|
||||
return nil, errs.New(errs.CodeChatMessageNotFound, "not found")
|
||||
}
|
||||
return m, nil
|
||||
}
|
||||
|
||||
func (r *msgFakeRepo) InsertMessage(_ context.Context, convID uuid.UUID, role MessageRole, content string, status MessageStatus) (*Message, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
id := r.nextMsgID
|
||||
r.nextMsgID++
|
||||
m := &Message{
|
||||
ID: id,
|
||||
ConversationID: convID,
|
||||
Role: role,
|
||||
Content: content,
|
||||
Status: status,
|
||||
}
|
||||
r.messages[id] = m
|
||||
r.insertedMessages = append(r.insertedMessages, m)
|
||||
return m, nil
|
||||
}
|
||||
|
||||
func (r *msgFakeRepo) DeleteMessage(_ context.Context, id int64) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
delete(r.messages, id)
|
||||
r.deletedMessageIDs = append(r.deletedMessageIDs, id)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *msgFakeRepo) UpdateMessageError(_ context.Context, id int64, msg *string) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
if m, ok := r.messages[id]; ok {
|
||||
m.Status = StatusError
|
||||
m.ErrorMessage = msg
|
||||
}
|
||||
r.errorUpdates = append(r.errorUpdates, errorUpdate{id: id, msg: msg})
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *msgFakeRepo) AttachToMessage(_ context.Context, ids []uuid.UUID, messageID int64) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.attachToMessageCalls = append(r.attachToMessageCalls, attachCall{ids: ids, messageID: messageID})
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *msgFakeRepo) ListAttachmentsForMessage(_ context.Context, _ int64) ([]Attachment, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (r *msgFakeRepo) ListOrphanAttachmentsForUser(_ context.Context, _ uuid.UUID, ids []uuid.UUID) ([]Attachment, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
var out []Attachment
|
||||
for _, id := range ids {
|
||||
if a, ok := r.attachments[id]; ok {
|
||||
out = append(out, *a)
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *msgFakeRepo) ListMessagesAscending(_ context.Context, convID uuid.UUID) ([]Message, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
var out []Message
|
||||
for _, m := range r.messages {
|
||||
if m.ConversationID == convID {
|
||||
out = append(out, *m)
|
||||
}
|
||||
}
|
||||
// Sort by ID ascending.
|
||||
for i := 1; i < len(out); i++ {
|
||||
for j := i; j > 0 && out[j].ID < out[j-1].ID; j-- {
|
||||
out[j], out[j-1] = out[j-1], out[j]
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *msgFakeRepo) ListMessagesByConversation(_ context.Context, convID uuid.UUID, beforeID *int64, limit int) ([]Message, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
var out []Message
|
||||
for _, m := range r.messages {
|
||||
if m.ConversationID == convID {
|
||||
if beforeID == nil || m.ID < *beforeID {
|
||||
out = append(out, *m)
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(out) > limit {
|
||||
out = out[len(out)-limit:]
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *msgFakeRepo) WithTx(_ context.Context) (Repository, Tx, error) {
|
||||
r.mu.Lock()
|
||||
r.withTxCalled = true
|
||||
r.mu.Unlock()
|
||||
return r, &fakeTx{}, nil
|
||||
}
|
||||
|
||||
// ===== fakeHub =====
|
||||
|
||||
type fakeHub struct {
|
||||
mu sync.Mutex
|
||||
started []StreamParams
|
||||
cancelled []int64
|
||||
startErr error
|
||||
subscribeFn func(int64) (Subscription, bool)
|
||||
}
|
||||
|
||||
func (h *fakeHub) Start(_ context.Context, p StreamParams) error {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
if h.startErr != nil {
|
||||
return h.startErr
|
||||
}
|
||||
h.started = append(h.started, p)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (h *fakeHub) Cancel(messageID int64) bool {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
h.cancelled = append(h.cancelled, messageID)
|
||||
return true
|
||||
}
|
||||
|
||||
func (h *fakeHub) Subscribe(messageID int64) (Subscription, bool) {
|
||||
if h.subscribeFn != nil {
|
||||
return h.subscribeFn(messageID)
|
||||
}
|
||||
return nil, false
|
||||
}
|
||||
|
||||
// ===== fakeStorage =====
|
||||
|
||||
type fakeStorage struct {
|
||||
blobs map[string][]byte
|
||||
}
|
||||
|
||||
func newFakeStorage() *fakeStorage {
|
||||
return &fakeStorage{blobs: make(map[string][]byte)}
|
||||
}
|
||||
|
||||
func (f *fakeStorage) Put(_ context.Context, key string, r io.Reader, _ string) (storage.PutResult, error) {
|
||||
data, err := io.ReadAll(r)
|
||||
if err != nil {
|
||||
return storage.PutResult{}, err
|
||||
}
|
||||
f.blobs[key] = data
|
||||
return storage.PutResult{StoragePath: key, Size: int64(len(data))}, nil
|
||||
}
|
||||
|
||||
func (f *fakeStorage) Get(_ context.Context, key string) (io.ReadCloser, error) {
|
||||
data, ok := f.blobs[key]
|
||||
if !ok {
|
||||
return nil, errors.New("not found: " + key)
|
||||
}
|
||||
return io.NopCloser(bytes.NewReader(data)), nil
|
||||
}
|
||||
|
||||
func (f *fakeStorage) Delete(_ context.Context, _ string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// ===== test builder =====
|
||||
|
||||
func buildMessageService(repo *msgFakeRepo, hub *fakeHub, cfg MessageServiceConfig) MessageService {
|
||||
return NewMessageService(repo, newFakeStorage(), hub, noopAuditor{}, discardLogger(), cfg)
|
||||
}
|
||||
|
||||
func defaultMsgConfig() MessageServiceConfig {
|
||||
return MessageServiceConfig{
|
||||
MimeWhitelist: []string{"image/png", "image/jpeg", "application/pdf", "text/plain"},
|
||||
MaxFileBytes: 10 * 1024 * 1024, // 10 MB
|
||||
MaxMsgBytes: 50 * 1024 * 1024, // 50 MB
|
||||
SafetyPct: 0.1,
|
||||
}
|
||||
}
|
||||
|
||||
func newConv(userID, modelID uuid.UUID) *Conversation {
|
||||
c := &Conversation{
|
||||
ID: uuid.New(),
|
||||
UserID: userID,
|
||||
ModelID: modelID,
|
||||
}
|
||||
return c
|
||||
}
|
||||
|
||||
// ===== Send tests =====
|
||||
|
||||
func TestMessageService_Send_HappyPath(t *testing.T) {
|
||||
repo := newMsgFakeRepo()
|
||||
hub := &fakeHub{}
|
||||
userID := uuid.New()
|
||||
modelID := uuid.New()
|
||||
conv := newConv(userID, modelID)
|
||||
repo.addConversation(conv)
|
||||
|
||||
svc := buildMessageService(repo, hub, defaultMsgConfig())
|
||||
uID, aID, err := svc.Send(context.Background(), userID, conv.ID, "hello", nil)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.Greater(t, uID, int64(0))
|
||||
require.Greater(t, aID, int64(0))
|
||||
require.NotEqual(t, uID, aID)
|
||||
|
||||
// Two messages inserted: user (OK) + assistant (Pending).
|
||||
require.Len(t, repo.insertedMessages, 2)
|
||||
require.Equal(t, RoleUser, repo.insertedMessages[0].Role)
|
||||
require.Equal(t, StatusOK, repo.insertedMessages[0].Status)
|
||||
require.Equal(t, RoleAssistant, repo.insertedMessages[1].Role)
|
||||
require.Equal(t, StatusPending, repo.insertedMessages[1].Status)
|
||||
|
||||
// hub.Start was called.
|
||||
hub.mu.Lock()
|
||||
require.Len(t, hub.started, 1)
|
||||
require.Equal(t, aID, hub.started[0].AssistantMessageID)
|
||||
hub.mu.Unlock()
|
||||
}
|
||||
|
||||
func TestMessageService_Send_NotOwner_Forbidden(t *testing.T) {
|
||||
repo := newMsgFakeRepo()
|
||||
hub := &fakeHub{}
|
||||
owner := uuid.New()
|
||||
caller := uuid.New()
|
||||
conv := newConv(owner, uuid.New())
|
||||
repo.addConversation(conv)
|
||||
|
||||
svc := buildMessageService(repo, hub, defaultMsgConfig())
|
||||
_, _, err := svc.Send(context.Background(), caller, conv.ID, "hi", nil)
|
||||
|
||||
require.Error(t, err)
|
||||
var ae *errs.AppError
|
||||
require.True(t, errors.As(err, &ae))
|
||||
require.Equal(t, errs.CodeForbidden, ae.Code)
|
||||
}
|
||||
|
||||
func TestMessageService_Send_AttachmentOwnershipMismatch(t *testing.T) {
|
||||
repo := newMsgFakeRepo()
|
||||
hub := &fakeHub{}
|
||||
userID := uuid.New()
|
||||
conv := newConv(userID, uuid.New())
|
||||
repo.addConversation(conv)
|
||||
|
||||
// Register only 2 of 3 requested attachment IDs.
|
||||
a1 := &Attachment{ID: uuid.New(), UserID: userID, MimeType: "image/png", SizeBytes: 100}
|
||||
a2 := &Attachment{ID: uuid.New(), UserID: userID, MimeType: "image/png", SizeBytes: 100}
|
||||
repo.addAttachment(a1)
|
||||
repo.addAttachment(a2)
|
||||
|
||||
svc := buildMessageService(repo, hub, defaultMsgConfig())
|
||||
_, _, err := svc.Send(context.Background(), userID, conv.ID, "hi", []uuid.UUID{a1.ID, a2.ID, uuid.New()})
|
||||
|
||||
require.Error(t, err)
|
||||
var ae *errs.AppError
|
||||
require.True(t, errors.As(err, &ae))
|
||||
require.Equal(t, errs.CodeChatAttachmentTooLarge, ae.Code)
|
||||
}
|
||||
|
||||
func TestMessageService_Send_AttachmentTooLarge_SingleFile(t *testing.T) {
|
||||
repo := newMsgFakeRepo()
|
||||
hub := &fakeHub{}
|
||||
userID := uuid.New()
|
||||
conv := newConv(userID, uuid.New())
|
||||
repo.addConversation(conv)
|
||||
|
||||
a := &Attachment{ID: uuid.New(), UserID: userID, MimeType: "image/png", SizeBytes: 20 * 1024 * 1024} // 20 MB > 10 MB limit
|
||||
repo.addAttachment(a)
|
||||
|
||||
svc := buildMessageService(repo, hub, defaultMsgConfig())
|
||||
_, _, err := svc.Send(context.Background(), userID, conv.ID, "hi", []uuid.UUID{a.ID})
|
||||
|
||||
require.Error(t, err)
|
||||
var ae *errs.AppError
|
||||
require.True(t, errors.As(err, &ae))
|
||||
require.Equal(t, errs.CodeChatAttachmentTooLarge, ae.Code)
|
||||
}
|
||||
|
||||
func TestMessageService_Send_AttachmentTooLarge_Total(t *testing.T) {
|
||||
repo := newMsgFakeRepo()
|
||||
hub := &fakeHub{}
|
||||
userID := uuid.New()
|
||||
conv := newConv(userID, uuid.New())
|
||||
repo.addConversation(conv)
|
||||
|
||||
// Each file is 9 MB (under per-file limit), but total 27 MB > 50 MB? Let's use 30+30 = 60 MB.
|
||||
a1 := &Attachment{ID: uuid.New(), UserID: userID, MimeType: "image/png", SizeBytes: 30 * 1024 * 1024}
|
||||
a2 := &Attachment{ID: uuid.New(), UserID: userID, MimeType: "image/png", SizeBytes: 30 * 1024 * 1024}
|
||||
repo.addAttachment(a1)
|
||||
repo.addAttachment(a2)
|
||||
|
||||
// Use config with smaller maxMsg to make the test reliable.
|
||||
cfg := defaultMsgConfig()
|
||||
cfg.MaxFileBytes = 50 * 1024 * 1024 // 50 MB per file — both pass individually
|
||||
cfg.MaxMsgBytes = 40 * 1024 * 1024 // 40 MB total — 60 MB fails
|
||||
|
||||
svc := buildMessageService(repo, hub, cfg)
|
||||
_, _, err := svc.Send(context.Background(), userID, conv.ID, "hi", []uuid.UUID{a1.ID, a2.ID})
|
||||
|
||||
require.Error(t, err)
|
||||
var ae *errs.AppError
|
||||
require.True(t, errors.As(err, &ae))
|
||||
require.Equal(t, errs.CodeChatAttachmentTooLarge, ae.Code)
|
||||
}
|
||||
|
||||
func TestMessageService_Send_AttachmentMimeNotAllowed(t *testing.T) {
|
||||
repo := newMsgFakeRepo()
|
||||
hub := &fakeHub{}
|
||||
userID := uuid.New()
|
||||
conv := newConv(userID, uuid.New())
|
||||
repo.addConversation(conv)
|
||||
|
||||
a := &Attachment{ID: uuid.New(), UserID: userID, MimeType: "video/mp4", SizeBytes: 100}
|
||||
repo.addAttachment(a)
|
||||
|
||||
svc := buildMessageService(repo, hub, defaultMsgConfig())
|
||||
_, _, err := svc.Send(context.Background(), userID, conv.ID, "hi", []uuid.UUID{a.ID})
|
||||
|
||||
require.Error(t, err)
|
||||
var ae *errs.AppError
|
||||
require.True(t, errors.As(err, &ae))
|
||||
require.Equal(t, errs.CodeChatAttachmentUnsupportedMime, ae.Code)
|
||||
}
|
||||
|
||||
func TestMessageService_Send_CapabilityMismatch(t *testing.T) {
|
||||
repo := newMsgFakeRepo()
|
||||
hub := &fakeHub{}
|
||||
userID := uuid.New()
|
||||
|
||||
// Model without image capability.
|
||||
modelID := uuid.New()
|
||||
conv := newConv(userID, modelID)
|
||||
repo.addConversation(conv)
|
||||
|
||||
// Override GetLLMModel to return a model without image capability.
|
||||
specialRepo := &capMismatchRepo{msgFakeRepo: *newMsgFakeRepo()}
|
||||
specialRepo.conversations = repo.conversations
|
||||
specialRepo.nextMsgID = 1
|
||||
|
||||
a := &Attachment{ID: uuid.New(), UserID: userID, MimeType: "image/png", SizeBytes: 100}
|
||||
specialRepo.addAttachment(a)
|
||||
|
||||
svc := NewMessageService(specialRepo, newFakeStorage(), hub, noopAuditor{}, discardLogger(), defaultMsgConfig())
|
||||
_, _, err := svc.Send(context.Background(), userID, conv.ID, "hi", []uuid.UUID{a.ID})
|
||||
|
||||
require.Error(t, err)
|
||||
var ae *errs.AppError
|
||||
require.True(t, errors.As(err, &ae))
|
||||
require.Equal(t, errs.CodeChatModelCapabilityMismatch, ae.Code)
|
||||
}
|
||||
|
||||
// capMismatchRepo overrides GetLLMModel to return a model without image capability.
|
||||
type capMismatchRepo struct {
|
||||
msgFakeRepo
|
||||
}
|
||||
|
||||
func (r *capMismatchRepo) GetLLMModel(_ context.Context, id uuid.UUID) (*LLMModel, error) {
|
||||
return &LLMModel{
|
||||
ID: id,
|
||||
EndpointID: uuid.New(),
|
||||
ContextWindow: 4096,
|
||||
MaxOutputTokens: 1024,
|
||||
Capabilities: ModelCapabilities{Image: false, PDF: false},
|
||||
}, nil
|
||||
}
|
||||
|
||||
func TestMessageService_Send_HubStartFailure_MarksError(t *testing.T) {
|
||||
repo := newMsgFakeRepo()
|
||||
hub := &fakeHub{startErr: errors.New("hub overload")}
|
||||
userID := uuid.New()
|
||||
conv := newConv(userID, uuid.New())
|
||||
repo.addConversation(conv)
|
||||
|
||||
svc := buildMessageService(repo, hub, defaultMsgConfig())
|
||||
uID, aID, err := svc.Send(context.Background(), userID, conv.ID, "hello", nil)
|
||||
|
||||
require.Error(t, err)
|
||||
require.Greater(t, uID, int64(0))
|
||||
require.Greater(t, aID, int64(0))
|
||||
|
||||
// UpdateMessageError must have been called for the assistant message.
|
||||
repo.mu.Lock()
|
||||
require.Len(t, repo.errorUpdates, 1)
|
||||
require.Equal(t, aID, repo.errorUpdates[0].id)
|
||||
require.NotNil(t, repo.errorUpdates[0].msg)
|
||||
require.Equal(t, "hub overload", *repo.errorUpdates[0].msg)
|
||||
repo.mu.Unlock()
|
||||
}
|
||||
|
||||
func TestMessageService_Send_NoAttachments(t *testing.T) {
|
||||
repo := newMsgFakeRepo()
|
||||
hub := &fakeHub{}
|
||||
userID := uuid.New()
|
||||
conv := newConv(userID, uuid.New())
|
||||
repo.addConversation(conv)
|
||||
|
||||
svc := buildMessageService(repo, hub, defaultMsgConfig())
|
||||
_, _, err := svc.Send(context.Background(), userID, conv.ID, "hello", []uuid.UUID{})
|
||||
|
||||
require.NoError(t, err)
|
||||
|
||||
// AttachToMessage must NOT have been called.
|
||||
repo.mu.Lock()
|
||||
require.Empty(t, repo.attachToMessageCalls)
|
||||
repo.mu.Unlock()
|
||||
}
|
||||
|
||||
func TestMessageService_Send_WithAttachments_CallsAttachToMessage(t *testing.T) {
|
||||
repo := newMsgFakeRepo()
|
||||
hub := &fakeHub{}
|
||||
userID := uuid.New()
|
||||
conv := newConv(userID, uuid.New())
|
||||
repo.addConversation(conv)
|
||||
|
||||
a1 := &Attachment{ID: uuid.New(), UserID: userID, MimeType: "image/png", SizeBytes: 100}
|
||||
a2 := &Attachment{ID: uuid.New(), UserID: userID, MimeType: "image/jpeg", SizeBytes: 200}
|
||||
repo.addAttachment(a1)
|
||||
repo.addAttachment(a2)
|
||||
|
||||
svc := buildMessageService(repo, hub, defaultMsgConfig())
|
||||
uID, _, err := svc.Send(context.Background(), userID, conv.ID, "hello", []uuid.UUID{a1.ID, a2.ID})
|
||||
|
||||
require.NoError(t, err)
|
||||
|
||||
repo.mu.Lock()
|
||||
require.Len(t, repo.attachToMessageCalls, 1)
|
||||
require.Equal(t, uID, repo.attachToMessageCalls[0].messageID)
|
||||
require.Len(t, repo.attachToMessageCalls[0].ids, 2)
|
||||
repo.mu.Unlock()
|
||||
}
|
||||
|
||||
// ===== Cancel tests =====
|
||||
|
||||
func TestMessageService_Cancel_NotPending(t *testing.T) {
|
||||
repo := newMsgFakeRepo()
|
||||
hub := &fakeHub{}
|
||||
userID := uuid.New()
|
||||
conv := newConv(userID, uuid.New())
|
||||
repo.addConversation(conv)
|
||||
|
||||
m := &Message{ID: 99, ConversationID: conv.ID, Status: StatusOK, Role: RoleAssistant}
|
||||
repo.addMessage(m)
|
||||
|
||||
svc := buildMessageService(repo, hub, defaultMsgConfig())
|
||||
err := svc.Cancel(context.Background(), userID, 99)
|
||||
|
||||
require.Error(t, err)
|
||||
var ae *errs.AppError
|
||||
require.True(t, errors.As(err, &ae))
|
||||
require.Equal(t, errs.CodeChatMessageNotPending, ae.Code)
|
||||
}
|
||||
|
||||
func TestMessageService_Cancel_HappyPath(t *testing.T) {
|
||||
repo := newMsgFakeRepo()
|
||||
hub := &fakeHub{}
|
||||
userID := uuid.New()
|
||||
conv := newConv(userID, uuid.New())
|
||||
repo.addConversation(conv)
|
||||
|
||||
m := &Message{ID: 42, ConversationID: conv.ID, Status: StatusPending, Role: RoleAssistant}
|
||||
repo.addMessage(m)
|
||||
|
||||
svc := buildMessageService(repo, hub, defaultMsgConfig())
|
||||
err := svc.Cancel(context.Background(), userID, 42)
|
||||
|
||||
require.NoError(t, err)
|
||||
|
||||
hub.mu.Lock()
|
||||
require.Len(t, hub.cancelled, 1)
|
||||
require.Equal(t, int64(42), hub.cancelled[0])
|
||||
hub.mu.Unlock()
|
||||
}
|
||||
|
||||
func TestMessageService_Cancel_NotOwner_Forbidden(t *testing.T) {
|
||||
repo := newMsgFakeRepo()
|
||||
hub := &fakeHub{}
|
||||
owner := uuid.New()
|
||||
caller := uuid.New()
|
||||
conv := newConv(owner, uuid.New())
|
||||
repo.addConversation(conv)
|
||||
|
||||
m := &Message{ID: 77, ConversationID: conv.ID, Status: StatusPending, Role: RoleAssistant}
|
||||
repo.addMessage(m)
|
||||
|
||||
svc := buildMessageService(repo, hub, defaultMsgConfig())
|
||||
err := svc.Cancel(context.Background(), caller, 77)
|
||||
|
||||
require.Error(t, err)
|
||||
var ae *errs.AppError
|
||||
require.True(t, errors.As(err, &ae))
|
||||
require.Equal(t, errs.CodeForbidden, ae.Code)
|
||||
}
|
||||
|
||||
// ===== Retry tests =====
|
||||
|
||||
func TestMessageService_Retry_NotRetriable(t *testing.T) {
|
||||
repo := newMsgFakeRepo()
|
||||
hub := &fakeHub{}
|
||||
userID := uuid.New()
|
||||
conv := newConv(userID, uuid.New())
|
||||
repo.addConversation(conv)
|
||||
|
||||
m := &Message{ID: 55, ConversationID: conv.ID, Status: StatusOK, Role: RoleAssistant}
|
||||
repo.addMessage(m)
|
||||
|
||||
svc := buildMessageService(repo, hub, defaultMsgConfig())
|
||||
_, err := svc.Retry(context.Background(), userID, 55)
|
||||
|
||||
require.Error(t, err)
|
||||
var ae *errs.AppError
|
||||
require.True(t, errors.As(err, &ae))
|
||||
require.Equal(t, errs.CodeChatMessageNotPending, ae.Code)
|
||||
}
|
||||
|
||||
func TestMessageService_Retry_HappyPath(t *testing.T) {
|
||||
repo := newMsgFakeRepo()
|
||||
hub := &fakeHub{}
|
||||
userID := uuid.New()
|
||||
conv := newConv(userID, uuid.New())
|
||||
repo.addConversation(conv)
|
||||
|
||||
// Old failed message.
|
||||
old := &Message{ID: 10, ConversationID: conv.ID, Status: StatusError, Role: RoleAssistant}
|
||||
repo.addMessage(old)
|
||||
|
||||
svc := buildMessageService(repo, hub, defaultMsgConfig())
|
||||
newID, err := svc.Retry(context.Background(), userID, 10)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.Greater(t, newID, int64(0))
|
||||
|
||||
// Old message deleted.
|
||||
repo.mu.Lock()
|
||||
require.Contains(t, repo.deletedMessageIDs, int64(10))
|
||||
|
||||
// New assistant message inserted as Pending.
|
||||
var newMsg *Message
|
||||
for _, m := range repo.insertedMessages {
|
||||
if m.ID == newID {
|
||||
newMsg = m
|
||||
break
|
||||
}
|
||||
}
|
||||
repo.mu.Unlock()
|
||||
require.NotNil(t, newMsg)
|
||||
require.Equal(t, StatusPending, newMsg.Status)
|
||||
require.Equal(t, RoleAssistant, newMsg.Role)
|
||||
|
||||
// hub.Start called.
|
||||
hub.mu.Lock()
|
||||
require.Len(t, hub.started, 1)
|
||||
require.Equal(t, newID, hub.started[0].AssistantMessageID)
|
||||
hub.mu.Unlock()
|
||||
}
|
||||
|
||||
func TestMessageService_Retry_Cancelled_HappyPath(t *testing.T) {
|
||||
repo := newMsgFakeRepo()
|
||||
hub := &fakeHub{}
|
||||
userID := uuid.New()
|
||||
conv := newConv(userID, uuid.New())
|
||||
repo.addConversation(conv)
|
||||
|
||||
old := &Message{ID: 11, ConversationID: conv.ID, Status: StatusCancelled, Role: RoleAssistant}
|
||||
repo.addMessage(old)
|
||||
|
||||
svc := buildMessageService(repo, hub, defaultMsgConfig())
|
||||
newID, err := svc.Retry(context.Background(), userID, 11)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.Greater(t, newID, int64(0))
|
||||
}
|
||||
|
||||
// ===== Stream tests =====
|
||||
|
||||
// closedSubscription immediately returns closed=true with one event.
|
||||
type closedSubscription struct {
|
||||
events []SSEEvent
|
||||
}
|
||||
|
||||
func (s *closedSubscription) Next(_ context.Context, lastID int64) ([]SSEEvent, bool, error) {
|
||||
if lastID < int64(len(s.events)) {
|
||||
out := s.events[lastID:]
|
||||
return out, true, nil
|
||||
}
|
||||
return nil, true, nil
|
||||
}
|
||||
|
||||
func TestMessageService_Stream_HappyPath_OK(t *testing.T) {
|
||||
repo := newMsgFakeRepo()
|
||||
hub := &fakeHub{}
|
||||
userID := uuid.New()
|
||||
conv := newConv(userID, uuid.New())
|
||||
repo.addConversation(conv)
|
||||
|
||||
m := &Message{
|
||||
ID: 1,
|
||||
ConversationID: conv.ID,
|
||||
Status: StatusOK,
|
||||
Role: RoleAssistant,
|
||||
Content: "hello world",
|
||||
}
|
||||
repo.addMessage(m)
|
||||
|
||||
svc := buildMessageService(repo, hub, defaultMsgConfig())
|
||||
ch, err := svc.Stream(context.Background(), userID, conv.ID, 0)
|
||||
|
||||
require.NoError(t, err)
|
||||
|
||||
var events []SSEEvent
|
||||
for ev := range ch {
|
||||
events = append(events, ev)
|
||||
}
|
||||
|
||||
// Expect: message_meta + text_delta + done
|
||||
require.GreaterOrEqual(t, len(events), 3)
|
||||
require.Equal(t, "message_meta", events[0].Event)
|
||||
require.Equal(t, "text_delta", events[1].Event)
|
||||
require.Equal(t, "done", events[2].Event)
|
||||
}
|
||||
|
||||
func TestMessageService_Stream_PendingFallbackToError(t *testing.T) {
|
||||
repo := newMsgFakeRepo()
|
||||
hub := &fakeHub{
|
||||
subscribeFn: func(_ int64) (Subscription, bool) {
|
||||
return nil, false // simulate server restart
|
||||
},
|
||||
}
|
||||
userID := uuid.New()
|
||||
conv := newConv(userID, uuid.New())
|
||||
repo.addConversation(conv)
|
||||
|
||||
m := &Message{
|
||||
ID: 1,
|
||||
ConversationID: conv.ID,
|
||||
Status: StatusPending,
|
||||
Role: RoleAssistant,
|
||||
}
|
||||
repo.addMessage(m)
|
||||
|
||||
svc := buildMessageService(repo, hub, defaultMsgConfig())
|
||||
ch, err := svc.Stream(context.Background(), userID, conv.ID, 0)
|
||||
|
||||
require.NoError(t, err)
|
||||
|
||||
var events []SSEEvent
|
||||
for ev := range ch {
|
||||
events = append(events, ev)
|
||||
}
|
||||
|
||||
// Should receive message_meta + error event.
|
||||
require.GreaterOrEqual(t, len(events), 2)
|
||||
var gotError bool
|
||||
for _, ev := range events {
|
||||
if ev.Event == "error" {
|
||||
gotError = true
|
||||
}
|
||||
}
|
||||
require.True(t, gotError, "expected error event in stream")
|
||||
|
||||
// UpdateMessageError should have been called.
|
||||
repo.mu.Lock()
|
||||
require.NotEmpty(t, repo.errorUpdates)
|
||||
require.Equal(t, int64(1), repo.errorUpdates[0].id)
|
||||
require.NotNil(t, repo.errorUpdates[0].msg)
|
||||
require.Equal(t, "server restart", *repo.errorUpdates[0].msg)
|
||||
repo.mu.Unlock()
|
||||
}
|
||||
|
||||
func TestMessageService_Stream_NotOwner_Forbidden(t *testing.T) {
|
||||
repo := newMsgFakeRepo()
|
||||
hub := &fakeHub{}
|
||||
owner := uuid.New()
|
||||
caller := uuid.New()
|
||||
conv := newConv(owner, uuid.New())
|
||||
repo.addConversation(conv)
|
||||
|
||||
svc := buildMessageService(repo, hub, defaultMsgConfig())
|
||||
_, err := svc.Stream(context.Background(), caller, conv.ID, 0)
|
||||
|
||||
require.Error(t, err)
|
||||
var ae *errs.AppError
|
||||
require.True(t, errors.As(err, &ae))
|
||||
require.Equal(t, errs.CodeForbidden, ae.Code)
|
||||
}
|
||||
|
||||
// ===== ListMessages tests =====
|
||||
|
||||
func TestMessageService_ListMessages_DefaultLimit(t *testing.T) {
|
||||
repo := newMsgFakeRepo()
|
||||
hub := &fakeHub{}
|
||||
userID := uuid.New()
|
||||
conv := newConv(userID, uuid.New())
|
||||
repo.addConversation(conv)
|
||||
|
||||
svc := buildMessageService(repo, hub, defaultMsgConfig())
|
||||
|
||||
// limit=0 should return up to 50 (no error even if no messages exist).
|
||||
msgs, err := svc.ListMessages(context.Background(), userID, conv.ID, nil, 0)
|
||||
require.NoError(t, err)
|
||||
_ = msgs // empty slice is fine; we're testing the limit clamp, not the content
|
||||
}
|
||||
|
||||
func TestMessageService_ListMessages_NotOwner(t *testing.T) {
|
||||
repo := newMsgFakeRepo()
|
||||
hub := &fakeHub{}
|
||||
owner := uuid.New()
|
||||
caller := uuid.New()
|
||||
conv := newConv(owner, uuid.New())
|
||||
repo.addConversation(conv)
|
||||
|
||||
svc := buildMessageService(repo, hub, defaultMsgConfig())
|
||||
_, err := svc.ListMessages(context.Background(), caller, conv.ID, nil, 10)
|
||||
|
||||
require.Error(t, err)
|
||||
var ae *errs.AppError
|
||||
require.True(t, errors.As(err, &ae))
|
||||
require.Equal(t, errs.CodeForbidden, ae.Code)
|
||||
}
|
||||
|
||||
func TestMessageService_ListMessages_CapClamped(t *testing.T) {
|
||||
repo := newMsgFakeRepo()
|
||||
hub := &fakeHub{}
|
||||
userID := uuid.New()
|
||||
conv := newConv(userID, uuid.New())
|
||||
repo.addConversation(conv)
|
||||
|
||||
svc := buildMessageService(repo, hub, defaultMsgConfig())
|
||||
|
||||
// limit=999 should be clamped to 50 (no error even if no messages exist).
|
||||
msgs, err := svc.ListMessages(context.Background(), userID, conv.ID, nil, 999)
|
||||
require.NoError(t, err)
|
||||
_ = msgs // empty slice is fine; we're testing the cap clamp, not the content
|
||||
}
|
||||
|
||||
// ===== composeHistory / truncateByTokens tests =====
|
||||
|
||||
func TestComposeHistory_TruncatesByTokens(t *testing.T) {
|
||||
// Build a list of messages with substantial text content.
|
||||
// With a very small budget, only the most recent message should survive.
|
||||
msgs := make([]llm.Message, 5)
|
||||
for i := range msgs {
|
||||
text := "This is message number " + int64ToString(int64(i)) + " with some content that takes up token budget in the conversation history."
|
||||
msgs[i] = llm.Message{
|
||||
Role: llm.RoleUser,
|
||||
Blocks: []llm.ContentBlock{{Type: "text", Text: text}},
|
||||
}
|
||||
}
|
||||
|
||||
systemPrompt := ""
|
||||
// Budget of 10 tokens: only the last message should fit (roughly 30+ tokens each).
|
||||
budget := 10
|
||||
result := truncateByTokens(msgs, systemPrompt, budget)
|
||||
|
||||
// With a budget of 10, fewer messages than the original 5 should survive.
|
||||
require.Less(t, len(result), len(msgs))
|
||||
}
|
||||
|
||||
func TestTruncateByTokens_AllFit(t *testing.T) {
|
||||
msgs := []llm.Message{
|
||||
{Role: llm.RoleUser, Blocks: []llm.ContentBlock{{Type: "text", Text: "hi"}}},
|
||||
{Role: llm.RoleAssistant, Blocks: []llm.ContentBlock{{Type: "text", Text: "hello"}}},
|
||||
}
|
||||
|
||||
result := truncateByTokens(msgs, "", 10000)
|
||||
require.Equal(t, len(msgs), len(result))
|
||||
}
|
||||
|
||||
func TestTruncateByTokens_ZeroBudget(t *testing.T) {
|
||||
msgs := []llm.Message{
|
||||
{Role: llm.RoleUser, Blocks: []llm.ContentBlock{{Type: "text", Text: "hi"}}},
|
||||
}
|
||||
|
||||
result := truncateByTokens(msgs, "", 0)
|
||||
require.Empty(t, result)
|
||||
}
|
||||
|
||||
func TestTruncateByTokens_PreservesNewest(t *testing.T) {
|
||||
// Make messages with distinct content.
|
||||
oldest := llm.Message{Role: llm.RoleUser, Blocks: []llm.ContentBlock{{Type: "text", Text: "oldest message that is quite long and should be dropped from history"}}}
|
||||
newest := llm.Message{Role: llm.RoleAssistant, Blocks: []llm.ContentBlock{{Type: "text", Text: "hi"}}}
|
||||
|
||||
msgs := []llm.Message{oldest, newest}
|
||||
|
||||
// Budget tight enough for only one message.
|
||||
result := truncateByTokens(msgs, "", 5)
|
||||
|
||||
require.Len(t, result, 1)
|
||||
require.Equal(t, "hi", result[0].Blocks[0].Text)
|
||||
}
|
||||
|
||||
// ===== mimeAllowed / capabilityAllows helpers =====
|
||||
|
||||
func TestMimeAllowed(t *testing.T) {
|
||||
svc := &messageService{
|
||||
mimeWhite: map[string]bool{
|
||||
"image/png": true,
|
||||
"application/pdf": true,
|
||||
"text/*": true,
|
||||
},
|
||||
}
|
||||
|
||||
require.True(t, svc.mimeAllowed("image/png"))
|
||||
require.True(t, svc.mimeAllowed("application/pdf"))
|
||||
require.True(t, svc.mimeAllowed("text/plain"))
|
||||
require.True(t, svc.mimeAllowed("text/html"))
|
||||
require.False(t, svc.mimeAllowed("video/mp4"))
|
||||
require.False(t, svc.mimeAllowed("audio/mpeg"))
|
||||
}
|
||||
|
||||
func TestCapabilityAllows(t *testing.T) {
|
||||
caps := ModelCapabilities{Image: true, PDF: false}
|
||||
|
||||
require.True(t, capabilityAllows(caps, "image/png"))
|
||||
require.True(t, capabilityAllows(caps, "image/jpeg"))
|
||||
require.False(t, capabilityAllows(caps, "application/pdf"))
|
||||
require.True(t, capabilityAllows(caps, "text/plain"))
|
||||
require.False(t, capabilityAllows(caps, "video/mp4"))
|
||||
}
|
||||
Reference in New Issue
Block a user