agentic-coding-workflow/internal/chat/service_message_test.go

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"))
}