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

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
	briefReader ContextReader // 可为 nil；非 nil 时 composeHistory 注入 FK 上下文
}

// NewMessageService constructs a MessageService. briefReader 可为 nil：nil 时
// composeHistory 完全保留旧行为（仅用 conv.SystemPrompt），不做 FK 上下文注入。
func NewMessageService(
	repo Repository,
	st storage.Storage,
	hub StreamerHub,
	auditor audit.Recorder,
	log *slog.Logger,
	cfg MessageServiceConfig,
	briefReader ContextReader,
) 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,
		briefReader: briefReader,
	}
}

// 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
				}
				nextEventID := int64(0)
				for {
					events, closed, err := sub.Next(ctx, nextEventID)
					if err != nil {
						return
					}
					for _, ev := range events {
						out <- ev
						nextEventID = ev.ID + 1
					}
					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, uuid.UUID, error) {
	old, err := s.repo.GetMessage(ctx, messageID)
	if err != nil {
		return 0, uuid.Nil, err
	}
	if old.Status != StatusError && old.Status != StatusCancelled {
		return 0, uuid.Nil, errs.New(errs.CodeChatMessageNotPending, "message is not retriable")
	}
	conv, err := s.repo.GetConversation(ctx, old.ConversationID)
	if err != nil {
		return 0, uuid.Nil, err
	}
	if conv.UserID != userID {
		return 0, uuid.Nil, errs.New(errs.CodeForbidden, "not conversation owner")
	}

	repo, tx, err := s.repo.WithTx(ctx)
	if err != nil {
		return 0, uuid.Nil, err
	}
	committed := false
	defer func() {
		if !committed {
			_ = tx.Rollback(ctx)
		}
	}()

	if err := repo.DeleteMessage(ctx, messageID); err != nil {
		return 0, uuid.Nil, err
	}
	asstMsg, err := repo.InsertMessage(ctx, conv.ID, RoleAssistant, "", StatusPending)
	if err != nil {
		return 0, uuid.Nil, err
	}

	if err := tx.Commit(ctx); err != nil {
		return 0, uuid.Nil, 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, conv.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, conv.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, conv.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, conv.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, conv.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
	}
	msgs, err := s.repo.ListMessagesByConversation(ctx, conversationID, beforeID, limit)
	if err != nil {
		return nil, err
	}
	// Populate each message's attachments for the user-facing read path.
	for i := range msgs {
		atts, err := s.repo.ListAttachmentsForMessage(ctx, msgs[i].ID)
		if err != nil {
			return nil, err
		}
		msgs[i].Attachments = atts
	}
	return msgs, nil
}

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

	// Resolve the effective system prompt by merging conv.SystemPrompt with the
	// FK-derived context brief (requirement/issue/project mounts). The brief is
	// pre-bounded by the composer, so it cannot starve message history to zero
	// (truncateByTokens still reserves the remaining budget for messages).
	effectiveSystemPrompt := s.effectiveSystemPrompt(ctx, conv)

	// 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, effectiveSystemPrompt, budget)

	return llm.Request{
		SystemPrompt:    effectiveSystemPrompt,
		Messages:        msgs,
		MaxOutputTokens: model.MaxOutputTokens,
		Reasoning:       model.Capabilities.Reasoning,
	}, nil
}

// effectiveSystemPrompt 把会话静态 system prompt 与 FK 派生上下文简报合并。
// briefReader 为 nil、会话无任何 FK 挂载、或简报为空时，直接返回 conv.SystemPrompt
// （与旧行为完全一致）。合并是追加式的（FK 简报置于静态 prompt 之后），便于
// 让旧会话（曾把角色 prompt 烘焙进 SystemPrompt）与新会话都能正常工作。
func (s *messageService) effectiveSystemPrompt(ctx context.Context, conv *Conversation) string {
	if s.briefReader == nil || !conv.HasFKMount() {
		return conv.SystemPrompt
	}
	fkBrief, err := s.briefReader.BriefForConversation(ctx, conv)
	if err != nil {
		if s.log != nil {
			s.log.Warn("chat.compose_history.brief_failed",
				"conversation_id", conv.ID, "err", err.Error())
		}
		return conv.SystemPrompt
	}
	fkBrief = strings.TrimSpace(fkBrief)
	if fkBrief == "" {
		return conv.SystemPrompt
	}
	if strings.TrimSpace(conv.SystemPrompt) == "" {
		return fkBrief
	}
	return conv.SystemPrompt + "\n\n" + fkBrief
}

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