From c42887498ba6da591c429fc43cb7fcdf34fbcd3d Mon Sep 17 00:00:00 2001 From: Jerry Yan <792602257@qq.com> Date: Mon, 4 May 2026 17:39:28 +0800 Subject: [PATCH] feat(chat): MessageService Send/Stream/Cancel/Retry with attachment + capability + history truncation --- internal/chat/service_message.go | 553 +++++++++++++++ internal/chat/service_message_test.go | 935 ++++++++++++++++++++++++++ 2 files changed, 1488 insertions(+) create mode 100644 internal/chat/service_message.go create mode 100644 internal/chat/service_message_test.go diff --git a/internal/chat/service_message.go b/internal/chat/service_message.go new file mode 100644 index 0000000..03549d4 --- /dev/null +++ b/internal/chat/service_message.go @@ -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) + } +} diff --git a/internal/chat/service_message_test.go b/internal/chat/service_message_test.go new file mode 100644 index 0000000..154725c --- /dev/null +++ b/internal/chat/service_message_test.go @@ -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")) +}