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