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,553 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"log/slog"
|
||||
"strings"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"github.com/yan1h/agent-coding-workflow/internal/audit"
|
||||
"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"
|
||||
)
|
||||
|
||||
// compile-time assertion
|
||||
var _ MessageService = (*messageService)(nil)
|
||||
|
||||
// MessageServiceConfig carries tunable limits for messageService.
|
||||
type MessageServiceConfig struct {
|
||||
MimeWhitelist []string
|
||||
MaxFileBytes int64
|
||||
MaxMsgBytes int64
|
||||
SafetyPct float64
|
||||
}
|
||||
|
||||
type messageService struct {
|
||||
repo Repository
|
||||
storage storage.Storage
|
||||
hub StreamerHub
|
||||
auditor audit.Recorder
|
||||
log *slog.Logger
|
||||
mimeWhite map[string]bool
|
||||
maxFile int64
|
||||
maxMsg int64
|
||||
safetyPct float64
|
||||
}
|
||||
|
||||
// NewMessageService constructs a MessageService.
|
||||
func NewMessageService(
|
||||
repo Repository,
|
||||
st storage.Storage,
|
||||
hub StreamerHub,
|
||||
auditor audit.Recorder,
|
||||
log *slog.Logger,
|
||||
cfg MessageServiceConfig,
|
||||
) MessageService {
|
||||
mimeWhite := make(map[string]bool, len(cfg.MimeWhitelist))
|
||||
for _, m := range cfg.MimeWhitelist {
|
||||
mimeWhite[m] = true
|
||||
}
|
||||
return &messageService{
|
||||
repo: repo,
|
||||
storage: st,
|
||||
hub: hub,
|
||||
auditor: auditor,
|
||||
log: log,
|
||||
mimeWhite: mimeWhite,
|
||||
maxFile: cfg.MaxFileBytes,
|
||||
maxMsg: cfg.MaxMsgBytes,
|
||||
safetyPct: cfg.SafetyPct,
|
||||
}
|
||||
}
|
||||
|
||||
// Send persists a user message and a pending assistant message in one tx,
|
||||
// then kicks off the streamer hub. Returns both message IDs on success.
|
||||
func (s *messageService) Send(
|
||||
ctx context.Context,
|
||||
userID, conversationID uuid.UUID,
|
||||
content string,
|
||||
attachmentIDs []uuid.UUID,
|
||||
) (userMsgID, assistantMsgID int64, err error) {
|
||||
conv, err := s.repo.GetConversation(ctx, conversationID)
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
if conv.UserID != userID {
|
||||
return 0, 0, errs.New(errs.CodeForbidden, "not conversation owner")
|
||||
}
|
||||
|
||||
model, err := s.repo.GetLLMModel(ctx, conv.ModelID)
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
endpoint, err := s.repo.GetLLMEndpoint(ctx, model.EndpointID)
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
|
||||
// Resolve attachments and validate ownership/size/mime/capability.
|
||||
var atts []Attachment
|
||||
if len(attachmentIDs) > 0 {
|
||||
atts, err = s.repo.ListOrphanAttachmentsForUser(ctx, userID, attachmentIDs)
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
if len(atts) != len(attachmentIDs) {
|
||||
return 0, 0, errs.New(errs.CodeChatAttachmentTooLarge, "attachment ownership mismatch")
|
||||
}
|
||||
|
||||
var totalSize int64
|
||||
for _, a := range atts {
|
||||
if s.maxFile > 0 && a.SizeBytes > s.maxFile {
|
||||
return 0, 0, errs.New(errs.CodeChatAttachmentTooLarge, "file exceeds per-file size limit")
|
||||
}
|
||||
totalSize += a.SizeBytes
|
||||
if !s.mimeAllowed(a.MimeType) {
|
||||
return 0, 0, errs.New(errs.CodeChatAttachmentUnsupportedMime, "unsupported mime type: "+a.MimeType)
|
||||
}
|
||||
if !capabilityAllows(model.Capabilities, a.MimeType) {
|
||||
return 0, 0, errs.New(errs.CodeChatModelCapabilityMismatch, "model does not support mime type: "+a.MimeType)
|
||||
}
|
||||
}
|
||||
if s.maxMsg > 0 && totalSize > s.maxMsg {
|
||||
return 0, 0, errs.New(errs.CodeChatAttachmentTooLarge, "total attachment size exceeds limit")
|
||||
}
|
||||
}
|
||||
|
||||
// Open transaction: persist user message + pending assistant message atomically.
|
||||
repo, tx, err := s.repo.WithTx(ctx)
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
committed := false
|
||||
defer func() {
|
||||
if !committed {
|
||||
_ = tx.Rollback(ctx)
|
||||
}
|
||||
}()
|
||||
|
||||
userMsg, err := repo.InsertMessage(ctx, conversationID, RoleUser, content, StatusOK)
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
|
||||
if len(atts) > 0 {
|
||||
ids := make([]uuid.UUID, len(atts))
|
||||
for i, a := range atts {
|
||||
ids[i] = a.ID
|
||||
}
|
||||
if err := repo.AttachToMessage(ctx, ids, userMsg.ID); err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
}
|
||||
|
||||
asstMsg, err := repo.InsertMessage(ctx, conversationID, RoleAssistant, "", StatusPending)
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
committed = true
|
||||
|
||||
// Audit after commit so it never rolls back with the tx.
|
||||
s.tryAudit(ctx, audit.Entry{
|
||||
UserID: &userID,
|
||||
Action: "chat.message.sent",
|
||||
TargetType: "message",
|
||||
TargetID: int64ToString(userMsg.ID),
|
||||
Metadata: map[string]any{
|
||||
"conversation_id": conversationID.String(),
|
||||
"attachment_count": len(atts),
|
||||
},
|
||||
})
|
||||
|
||||
req, err := s.composeHistory(ctx, conv, model)
|
||||
if err != nil {
|
||||
e := err.Error()
|
||||
_ = s.repo.UpdateMessageError(ctx, asstMsg.ID, &e)
|
||||
return userMsg.ID, asstMsg.ID, err
|
||||
}
|
||||
|
||||
params := StreamParams{
|
||||
UserID: userID,
|
||||
ConversationID: conversationID,
|
||||
EndpointID: endpoint.ID,
|
||||
ModelID: model.ID,
|
||||
AssistantMessageID: asstMsg.ID,
|
||||
Request: req,
|
||||
ModelMeta: *model,
|
||||
}
|
||||
if err := s.hub.Start(ctx, params); err != nil {
|
||||
e := err.Error()
|
||||
_ = s.repo.UpdateMessageError(ctx, asstMsg.ID, &e)
|
||||
return userMsg.ID, asstMsg.ID, err
|
||||
}
|
||||
|
||||
return userMsg.ID, asstMsg.ID, nil
|
||||
}
|
||||
|
||||
// Stream replays message history for the conversation and tails any in-flight
|
||||
// pending message via hub.Subscribe.
|
||||
func (s *messageService) Stream(
|
||||
ctx context.Context,
|
||||
userID, conversationID uuid.UUID,
|
||||
sinceMessageID int64,
|
||||
) (<-chan SSEEvent, error) {
|
||||
conv, err := s.repo.GetConversation(ctx, conversationID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if conv.UserID != userID {
|
||||
return nil, errs.New(errs.CodeForbidden, "not conversation owner")
|
||||
}
|
||||
|
||||
out := make(chan SSEEvent, 16)
|
||||
go func() {
|
||||
defer close(out)
|
||||
|
||||
all, err := s.repo.ListMessagesAscending(ctx, conversationID)
|
||||
if err != nil {
|
||||
out <- SSEEvent{Event: "error", Data: map[string]string{"code": string(errs.CodeInternal), "message": err.Error()}}
|
||||
return
|
||||
}
|
||||
|
||||
for _, m := range all {
|
||||
if m.ID <= sinceMessageID {
|
||||
continue
|
||||
}
|
||||
|
||||
// Emit message_meta for every qualifying message.
|
||||
out <- SSEEvent{
|
||||
Event: "message_meta",
|
||||
Data: map[string]any{
|
||||
"id": m.ID,
|
||||
"role": string(m.Role),
|
||||
"status": string(m.Status),
|
||||
},
|
||||
}
|
||||
|
||||
switch m.Status {
|
||||
case StatusOK:
|
||||
if m.Content != "" {
|
||||
out <- SSEEvent{Event: "text_delta", Data: map[string]string{"text": m.Content}}
|
||||
}
|
||||
if m.Thinking != nil && *m.Thinking != "" {
|
||||
out <- SSEEvent{Event: "thinking_delta", Data: map[string]string{"text": *m.Thinking}}
|
||||
}
|
||||
out <- SSEEvent{Event: "done", Data: map[string]string{"finish_reason": "stop"}}
|
||||
|
||||
case StatusError:
|
||||
out <- SSEEvent{Event: "error", Data: map[string]string{
|
||||
"code": string(errs.CodeLLMUpstream),
|
||||
"message": derefString(m.ErrorMessage),
|
||||
}}
|
||||
|
||||
case StatusCancelled:
|
||||
out <- SSEEvent{Event: "error", Data: map[string]string{"code": "cancelled"}}
|
||||
|
||||
case StatusPending:
|
||||
sub, ok := s.hub.Subscribe(m.ID)
|
||||
if !ok {
|
||||
e := "server restart"
|
||||
_ = s.repo.UpdateMessageError(ctx, m.ID, &e)
|
||||
out <- SSEEvent{Event: "error", Data: map[string]string{"code": "server_restart", "message": e}}
|
||||
return
|
||||
}
|
||||
lastID := int64(-1)
|
||||
for {
|
||||
events, closed, err := sub.Next(ctx, lastID)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
for _, ev := range events {
|
||||
out <- ev
|
||||
lastID = ev.ID
|
||||
}
|
||||
if closed {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// Cancel signals the streamer hub to stop generating for the given pending message.
|
||||
func (s *messageService) Cancel(ctx context.Context, userID uuid.UUID, messageID int64) error {
|
||||
m, err := s.repo.GetMessage(ctx, messageID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if m.Status != StatusPending {
|
||||
return errs.New(errs.CodeChatMessageNotPending, "message is not pending")
|
||||
}
|
||||
conv, err := s.repo.GetConversation(ctx, m.ConversationID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if conv.UserID != userID {
|
||||
return errs.New(errs.CodeForbidden, "not conversation owner")
|
||||
}
|
||||
s.hub.Cancel(messageID)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Retry deletes the failed/cancelled assistant message, inserts a fresh pending
|
||||
// one, and restarts the streamer.
|
||||
func (s *messageService) Retry(ctx context.Context, userID uuid.UUID, messageID int64) (int64, error) {
|
||||
old, err := s.repo.GetMessage(ctx, messageID)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if old.Status != StatusError && old.Status != StatusCancelled {
|
||||
return 0, errs.New(errs.CodeChatMessageNotPending, "message is not retriable")
|
||||
}
|
||||
conv, err := s.repo.GetConversation(ctx, old.ConversationID)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if conv.UserID != userID {
|
||||
return 0, errs.New(errs.CodeForbidden, "not conversation owner")
|
||||
}
|
||||
|
||||
repo, tx, err := s.repo.WithTx(ctx)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
committed := false
|
||||
defer func() {
|
||||
if !committed {
|
||||
_ = tx.Rollback(ctx)
|
||||
}
|
||||
}()
|
||||
|
||||
if err := repo.DeleteMessage(ctx, messageID); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
asstMsg, err := repo.InsertMessage(ctx, conv.ID, RoleAssistant, "", StatusPending)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
committed = true
|
||||
|
||||
model, err := s.repo.GetLLMModel(ctx, conv.ModelID)
|
||||
if err != nil {
|
||||
e := err.Error()
|
||||
_ = s.repo.UpdateMessageError(ctx, asstMsg.ID, &e)
|
||||
return asstMsg.ID, err
|
||||
}
|
||||
endpoint, err := s.repo.GetLLMEndpoint(ctx, model.EndpointID)
|
||||
if err != nil {
|
||||
e := err.Error()
|
||||
_ = s.repo.UpdateMessageError(ctx, asstMsg.ID, &e)
|
||||
return asstMsg.ID, err
|
||||
}
|
||||
|
||||
req, err := s.composeHistory(ctx, conv, model)
|
||||
if err != nil {
|
||||
e := err.Error()
|
||||
_ = s.repo.UpdateMessageError(ctx, asstMsg.ID, &e)
|
||||
return asstMsg.ID, err
|
||||
}
|
||||
|
||||
params := StreamParams{
|
||||
UserID: userID,
|
||||
ConversationID: conv.ID,
|
||||
EndpointID: endpoint.ID,
|
||||
ModelID: model.ID,
|
||||
AssistantMessageID: asstMsg.ID,
|
||||
Request: req,
|
||||
ModelMeta: *model,
|
||||
}
|
||||
if err := s.hub.Start(ctx, params); err != nil {
|
||||
e := err.Error()
|
||||
_ = s.repo.UpdateMessageError(ctx, asstMsg.ID, &e)
|
||||
return asstMsg.ID, err
|
||||
}
|
||||
|
||||
// Audit after successful hub start.
|
||||
s.tryAudit(ctx, audit.Entry{
|
||||
UserID: &userID,
|
||||
Action: "chat.message.retried",
|
||||
TargetType: "message",
|
||||
TargetID: int64ToString(asstMsg.ID),
|
||||
Metadata: map[string]any{
|
||||
"conversation_id": conv.ID.String(),
|
||||
"old_message_id": messageID,
|
||||
},
|
||||
})
|
||||
|
||||
return asstMsg.ID, nil
|
||||
}
|
||||
|
||||
// ListMessages returns paginated messages for a conversation, enforcing ownership.
|
||||
func (s *messageService) ListMessages(
|
||||
ctx context.Context,
|
||||
userID, conversationID uuid.UUID,
|
||||
beforeID *int64,
|
||||
limit int,
|
||||
) ([]Message, error) {
|
||||
conv, err := s.repo.GetConversation(ctx, conversationID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if conv.UserID != userID {
|
||||
return nil, errs.New(errs.CodeForbidden, "not conversation owner")
|
||||
}
|
||||
if limit == 0 || limit > 200 {
|
||||
limit = 50
|
||||
}
|
||||
return s.repo.ListMessagesByConversation(ctx, conversationID, beforeID, limit)
|
||||
}
|
||||
|
||||
// composeHistory loads all OK/Pending messages for the conversation, fetches
|
||||
// attachment blobs from storage, truncates by token budget, and returns an
|
||||
// llm.Request ready to send.
|
||||
func (s *messageService) composeHistory(ctx context.Context, conv *Conversation, model *LLMModel) (llm.Request, error) {
|
||||
all, err := s.repo.ListMessagesAscending(ctx, conv.ID)
|
||||
if err != nil {
|
||||
return llm.Request{}, err
|
||||
}
|
||||
|
||||
var msgs []llm.Message
|
||||
for _, m := range all {
|
||||
if m.Status != StatusOK && m.Status != StatusPending {
|
||||
continue
|
||||
}
|
||||
|
||||
msg := llm.Message{Role: llm.Role(m.Role)}
|
||||
|
||||
if m.Content != "" {
|
||||
msg.Blocks = append(msg.Blocks, llm.ContentBlock{Type: "text", Text: m.Content})
|
||||
}
|
||||
|
||||
// Load attachments from storage.
|
||||
atts, err := s.repo.ListAttachmentsForMessage(ctx, m.ID)
|
||||
if err != nil {
|
||||
return llm.Request{}, err
|
||||
}
|
||||
for _, a := range atts {
|
||||
var blockType string
|
||||
switch {
|
||||
case strings.HasPrefix(a.MimeType, "image/"):
|
||||
blockType = "image"
|
||||
case a.MimeType == "application/pdf":
|
||||
blockType = "document"
|
||||
default:
|
||||
continue // skip unsupported types (capability check gates at Send time)
|
||||
}
|
||||
|
||||
rc, err := s.storage.Get(ctx, a.StoragePath)
|
||||
if err != nil {
|
||||
return llm.Request{}, err
|
||||
}
|
||||
data, readErr := io.ReadAll(rc)
|
||||
rc.Close()
|
||||
if readErr != nil {
|
||||
return llm.Request{}, readErr
|
||||
}
|
||||
|
||||
msg.Blocks = append(msg.Blocks, llm.ContentBlock{
|
||||
Type: blockType,
|
||||
MimeType: a.MimeType,
|
||||
Data: data,
|
||||
})
|
||||
}
|
||||
|
||||
// Only append if there's content or attachments to send.
|
||||
if len(msg.Blocks) > 0 {
|
||||
msgs = append(msgs, msg)
|
||||
}
|
||||
}
|
||||
|
||||
// Compute token budget: context window minus max output tokens, then apply safety margin.
|
||||
budget := model.ContextWindow - model.MaxOutputTokens
|
||||
budget = int(float64(budget) * (1 - s.safetyPct))
|
||||
|
||||
msgs = truncateByTokens(msgs, conv.SystemPrompt, budget)
|
||||
|
||||
return llm.Request{
|
||||
SystemPrompt: conv.SystemPrompt,
|
||||
Messages: msgs,
|
||||
MaxOutputTokens: model.MaxOutputTokens,
|
||||
Reasoning: model.Capabilities.Reasoning,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// truncateByTokens walks messages backwards and drops oldest messages until the
|
||||
// estimated token count (system prompt + messages) fits within budget.
|
||||
func truncateByTokens(msgs []llm.Message, systemPrompt string, budget int) []llm.Message {
|
||||
systemTokens := llm.EstimateTokens(systemPrompt)
|
||||
remaining := budget - systemTokens
|
||||
if remaining <= 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Walk backwards to find how many messages fit.
|
||||
tokens := 0
|
||||
cutAt := len(msgs)
|
||||
for i := len(msgs) - 1; i >= 0; i-- {
|
||||
msgTokens := 0
|
||||
for _, b := range msgs[i].Blocks {
|
||||
if b.Type == "text" {
|
||||
msgTokens += llm.EstimateTokens(b.Text)
|
||||
}
|
||||
}
|
||||
if tokens+msgTokens > remaining {
|
||||
cutAt = i + 1
|
||||
break
|
||||
}
|
||||
tokens += msgTokens
|
||||
cutAt = i
|
||||
}
|
||||
|
||||
return msgs[cutAt:]
|
||||
}
|
||||
|
||||
// ===== Helpers =====
|
||||
|
||||
func (s *messageService) mimeAllowed(mt string) bool {
|
||||
if s.mimeWhite[mt] {
|
||||
return true
|
||||
}
|
||||
if s.mimeWhite["text/*"] && strings.HasPrefix(mt, "text/") {
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func capabilityAllows(c ModelCapabilities, mt string) bool {
|
||||
switch {
|
||||
case strings.HasPrefix(mt, "image/"):
|
||||
return c.Image
|
||||
case mt == "application/pdf":
|
||||
return c.PDF
|
||||
case strings.HasPrefix(mt, "text/"):
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func derefString(p *string) string {
|
||||
if p == nil {
|
||||
return ""
|
||||
}
|
||||
return *p
|
||||
}
|
||||
|
||||
func (s *messageService) tryAudit(ctx context.Context, e audit.Entry) {
|
||||
if err := s.auditor.Record(ctx, e); err != nil {
|
||||
s.log.Warn("audit record failed", "action", e.Action, "err", err)
|
||||
}
|
||||
}
|
||||
@@ -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