feat(chat): MessageService Send/Stream/Cancel/Retry with attachment + capability + history truncation

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

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