Agent-System/internal/app/chatflow/chatlog.go

package chatflow

import (
	"context"
	"sort"
	"strings"

	"chat-agent-system/internal/ctxmgr"
	"chat-agent-system/internal/projectinstance"
)

type sessionAuth struct {
	SessionID      string
	SessionOwnerID int64
	ProjectID      int64
	ViewerUserID   int64
	StageAccess    *projectinstance.StageAccess
}

// Adapter implements ChatLog on top of ctxmgr without moving table ownership.
type Adapter struct {
	Mgr        *ctxmgr.ContextManager
	Resolve    StageAccessResolver
	FailClosed bool
}

func NewAdapter(mgr *ctxmgr.ContextManager, resolve StageAccessResolver, failClosed bool) *Adapter {
	return &Adapter{Mgr: mgr, Resolve: resolve, FailClosed: failClosed}
}

func (a *Adapter) authorize(ctx context.Context, p Principal, sessionID string) (*sessionAuth, error) {
	if a == nil || a.Mgr == nil {
		return nil, ErrClosed
	}
	if p.UserID <= 0 {
		return nil, ErrUnauthenticated
	}
	if strings.TrimSpace(sessionID) == "" {
		return nil, ErrNotFound
	}
	ownerID, err := a.Mgr.GetSessionOwnerID(ctx, sessionID)
	if err != nil || ownerID <= 0 {
		return nil, ErrNotFound
	}
	if ownerID != p.UserID {
		return nil, ErrForbidden
	}
	auth := &sessionAuth{
		SessionID:      sessionID,
		SessionOwnerID: ownerID,
		ViewerUserID:   p.UserID,
	}
	if projectID, err := a.Mgr.GetSessionProjectID(ctx, sessionID); err == nil {
		auth.ProjectID = projectID
	}
	if auth.ProjectID > 0 && a.Resolve != nil {
		access, err := a.Resolve(ctx, auth.ProjectID, p.UserID)
		if err != nil {
			if a.FailClosed {
				return nil, ErrForbidden
			}
		} else {
			auth.StageAccess = access
		}
	}
	return auth, nil
}

func (a *Adapter) visible(auth *sessionAuth, msg ctxmgr.ChatMessageRecord) bool {
	if auth == nil {
		return false
	}
	return projectinstance.MessageRecordVisible(
		msg.ScopeKind,
		msg.StageKey,
		msg.Metadata,
		auth.StageAccess,
		auth.SessionOwnerID,
		auth.ViewerUserID,
	)
}

func (a *Adapter) filterMessages(auth *sessionAuth, msgs []ctxmgr.ChatMessageRecord) []ctxmgr.ChatMessageRecord {
	if len(msgs) == 0 {
		return msgs
	}
	out := make([]ctxmgr.ChatMessageRecord, 0, len(msgs))
	for _, msg := range msgs {
		if a.visible(auth, msg) {
			out = append(out, msg)
		}
	}
	return out
}

func (a *Adapter) GetWindow(ctx context.Context, p Principal, q ctxmgr.ChatWindowQuery) (*ctxmgr.ChatWindowResult, error) {
	auth, err := a.authorize(ctx, p, q.SessionID)
	if err != nil {
		return nil, err
	}
	res, err := a.Mgr.GetChatWindow(ctx, q)
	if err != nil {
		return nil, err
	}
	applyWindowFilter(auth, a, res)
	return res, nil
}

func applyWindowFilter(auth *sessionAuth, a *Adapter, res *ctxmgr.ChatWindowResult) {
	if res == nil || auth == nil {
		return
	}
	filtered := make([]ctxmgr.ChatDayBucket, 0, len(res.Days))
	for i := range res.Days {
		msgs := a.filterMessages(auth, res.Days[i].Messages)
		if len(msgs) == 0 {
			continue
		}
		res.Days[i].Messages = msgs
		res.Days[i].LoadedCount = len(msgs)
		res.Days[i].Count = len(msgs)
		res.Days[i].FirstMessageID = msgs[0].ID
		res.Days[i].LastMessageID = msgs[len(msgs)-1].ID
		filtered = append(filtered, res.Days[i])
	}
	res.Days = filtered
	if len(filtered) > 0 {
		res.OldestLoadedDay = filtered[0].DayKey
	} else {
		res.OldestLoadedDay = ""
	}
}

func (a *Adapter) GetDay(ctx context.Context, p Principal, q DayQuery) (*ctxmgr.ChatDayMessagesResult, error) {
	auth, err := a.authorize(ctx, p, q.SessionID)
	if err != nil {
		return nil, err
	}
	res, err := a.Mgr.GetChatDayMessages(ctx, q.SessionID, q.Date, q.AnchorID, q.BeforeID, q.AfterID, q.Limit)
	if err != nil {
		return nil, err
	}
	res.Messages = a.filterMessages(auth, res.Messages)
	res.LoadedCount = len(res.Messages)
	res.TotalCount = len(res.Messages)
	return res, nil
}

func (a *Adapter) History(ctx context.Context, p Principal, q HistoryQuery) ([]ctxmgr.ChatMessageRecord, error) {
	auth, err := a.authorize(ctx, p, q.SessionID)
	if err != nil {
		return nil, err
	}
	if q.Limit <= 0 {
		q.Limit = 50
	}
	if q.Offset < 0 {
		q.Offset = 0
	}
	msgs, err := a.Mgr.GetChatHistory(ctx, q.SessionID, q.Limit, q.Offset)
	if err != nil {
		return nil, err
	}
	return a.filterMessages(auth, msgs), nil
}

func (a *Adapter) ListVisible(ctx context.Context, p Principal, sessionID string) ([]ctxmgr.ChatMessageRecord, error) {
	auth, err := a.authorize(ctx, p, sessionID)
	if err != nil {
		return nil, err
	}
	msgs, err := a.Mgr.ListSessionMessagesForAuth(ctx, sessionID)
	if err != nil {
		return nil, err
	}
	return a.filterMessages(auth, msgs), nil
}

func (a *Adapter) CategoryStats(ctx context.Context, p Principal, sessionID string) ([]ctxmgr.ChatCategoryStat, error) {
	visible, err := a.ListVisible(ctx, p, sessionID)
	if err != nil {
		return nil, err
	}
	return aggregateCategoryStats(visible), nil
}

func (a *Adapter) DateStats(ctx context.Context, p Principal, sessionID string) ([]ctxmgr.ChatDateStat, error) {
	visible, err := a.ListVisible(ctx, p, sessionID)
	if err != nil {
		return nil, err
	}
	return aggregateDateStats(visible), nil
}

func (a *Adapter) Search(ctx context.Context, p Principal, q ctxmgr.ChatSearchQuery) ([]ctxmgr.ChatSearchHit, error) {
	auth, err := a.authorize(ctx, p, q.SessionID)
	if err != nil {
		return nil, err
	}
	q.UserID = p.UserID
	hits, err := a.Mgr.SearchChatMessages(ctx, q)
	if err != nil {
		return nil, err
	}
	out := make([]ctxmgr.ChatSearchHit, 0, len(hits))
	for _, hit := range hits {
		meta := map[string]any{"stage_key": hit.StageKey}
		if projectinstance.MessageRecordVisible("", hit.StageKey, meta, auth.StageAccess, auth.SessionOwnerID, auth.ViewerUserID) {
			out = append(out, hit)
		}
	}
	return out, nil
}

func (a *Adapter) Unread(ctx context.Context, p Principal, q UnreadQuery) (*ctxmgr.ChatUnreadSummary, error) {
	auth, err := a.authorize(ctx, p, q.SessionID)
	if err != nil {
		return nil, err
	}
	summary, err := a.Mgr.GetChatUnreadSummary(ctx, q.SessionID, p.UserID, q.SinceMessageID)
	if err != nil {
		return &ctxmgr.ChatUnreadSummary{UnreadMessageIDs: []int64{}, Items: []ctxmgr.UnreadDayItem{}}, nil
	}
	if summary == nil {
		return &ctxmgr.ChatUnreadSummary{UnreadMessageIDs: []int64{}, Items: []ctxmgr.UnreadDayItem{}}, nil
	}
	visible, err := a.Mgr.ListSessionMessagesForAuth(ctx, q.SessionID)
	if err != nil {
		return summary, nil
	}
	allowed := make(map[int64]bool, len(visible))
	for _, msg := range a.filterMessages(auth, visible) {
		allowed[msg.ID] = true
	}
	filteredIDs := make([]int64, 0, len(summary.UnreadMessageIDs))
	for _, id := range summary.UnreadMessageIDs {
		if allowed[id] {
			filteredIDs = append(filteredIDs, id)
		}
	}
	summary.UnreadMessageIDs = filteredIDs
	summary.Total = len(filteredIDs)
	return summary, nil
}

func (a *Adapter) MarkRead(ctx context.Context, p Principal, req ctxmgr.MarkChatReadRequest) (*ctxmgr.ChatReadPosition, error) {
	if _, err := a.authorize(ctx, p, req.SessionID); err != nil {
		return nil, err
	}
	req.UserID = p.UserID
	return a.Mgr.MarkChatRead(ctx, req)
}

func (a *Adapter) Export(ctx context.Context, p Principal, q ExportQuery) ([]ctxmgr.ChatMessageRecord, error) {
	auth, err := a.authorize(ctx, p, q.SessionID)
	if err != nil {
		return nil, err
	}
	msgs, err := a.Mgr.GetChatHistory(ctx, q.SessionID, 5000, 0)
	if err != nil {
		return nil, err
	}
	visible := a.filterMessages(auth, msgs)
	out := make([]ctxmgr.ChatMessageRecord, 0, len(visible))
	for _, m := range visible {
		dk := m.DayKey
		if dk == "" && !m.CreatedAt.IsZero() {
			dk = m.CreatedAt.Format("2006-01-02")
		}
		if q.DateFrom != "" && dk < q.DateFrom {
			continue
		}
		if q.DateTo != "" && dk > q.DateTo {
			continue
		}
		if q.Category != "" && q.Category != "all" && m.Category != q.Category {
			continue
		}
		out = append(out, m)
	}
	return out, nil
}

func (a *Adapter) ListDigests(ctx context.Context, p Principal, sessionID, from, to string) ([]ctxmgr.DailyDigest, error) {
	visible, err := a.ListVisible(ctx, p, sessionID)
	if err != nil {
		return nil, err
	}
	days := make(map[string]struct{}, len(visible))
	for _, msg := range visible {
		if msg.DayKey != "" {
			days[msg.DayKey] = struct{}{}
		}
	}
	digests, err := a.Mgr.ListDailyDigests(ctx, sessionID, from, to)
	if err != nil {
		return nil, err
	}
	out := make([]ctxmgr.DailyDigest, 0, len(digests))
	for _, d := range digests {
		if _, ok := days[d.DayKey]; ok {
			out = append(out, d)
		}
	}
	return out, nil
}

func (a *Adapter) GenerateDigest(ctx context.Context, p Principal, sessionID, dayKey string) (*ctxmgr.DailyDigest, error) {
	visible, err := a.ListVisible(ctx, p, sessionID)
	if err != nil {
		return nil, err
	}
	dayMsgs := make([]ctxmgr.ChatMessageRecord, 0, len(visible))
	for _, msg := range visible {
		if msg.DayKey == dayKey {
			dayMsgs = append(dayMsgs, msg)
		}
	}
	return a.Mgr.GenerateDailyDigestFromMessages(ctx, sessionID, dayKey, dayMsgs)
}

// GenerateRollingSummary builds a rolling summary from authorized-visible messages only.
// Prompt-injection content cannot expand the message set; auth is decided before summarization.
func (a *Adapter) GenerateRollingSummary(ctx context.Context, p Principal, sessionID string, earlyMsgs []ctxmgr.ChatMessageRecord) (*ctxmgr.ChatSummary, error) {
	auth, err := a.authorize(ctx, p, sessionID)
	if err != nil {
		return nil, err
	}
	visible := a.filterMessages(auth, earlyMsgs)
	oldSum, _ := a.Mgr.GetChatSummary(ctx, sessionID)
	return a.Mgr.GenerateRollingSummary(ctx, sessionID, oldSum, visible)
}

func (a *Adapter) Restore(ctx context.Context, p Principal, sessionID string) (*ctxmgr.RestoredSessionContext, error) {
	auth, err := a.authorize(ctx, p, sessionID)
	if err != nil {
		return nil, err
	}
	restored, err := a.Mgr.RestoreSessionContext(ctx, sessionID)
	if err != nil {
		return nil, err
	}
	if restored == nil {
		return nil, nil
	}
	restored.Messages = a.filterMessages(auth, restored.Messages)
	return restored, nil
}

func (a *Adapter) SearchMemory(ctx context.Context, p Principal, q MemoryQuery) ([]ctxmgr.RecalledMemoryItem, error) {
	if a == nil || a.Mgr == nil {
		return nil, ErrClosed
	}
	if p.UserID <= 0 {
		return nil, ErrUnauthenticated
	}
	items, err := a.Mgr.SearchChatMemory(ctx, ctxmgr.ChatMemorySearchOptions{
		Query:            q.Query,
		UserID:           p.UserID,
		ProjectID:        q.ProjectID,
		SessionID:        q.SessionID,
		ExcludeSessionID: q.ExcludeSessionID,
		TopK:             q.TopK,
		MinScore:         q.MinScore,
	})
	if err != nil {
		return nil, err
	}
	out := make([]ctxmgr.RecalledMemoryItem, 0, len(items))
	for _, item := range items {
		ownerID := item.UserID
		if item.SessionID != "" {
			if sidOwner, err := a.Mgr.GetSessionOwnerID(ctx, item.SessionID); err == nil && sidOwner > 0 {
				ownerID = sidOwner
			}
		}
		var access *projectinstance.StageAccess
		if item.ProjectID > 0 && a.Resolve != nil {
			if resolved, err := a.Resolve(ctx, item.ProjectID, p.UserID); err == nil {
				access = resolved
			} else if a.FailClosed {
				continue
			}
		}
		if projectinstance.MessageRecordVisible("", "", item.Metadata, access, ownerID, p.UserID) {
			out = append(out, item)
		}
	}
	return out, nil
}

func aggregateDateStats(msgs []ctxmgr.ChatMessageRecord) []ctxmgr.ChatDateStat {
	statMap := make(map[string]*ctxmgr.ChatDateStat)
	var orderedKeys []string
	for _, msg := range msgs {
		dKey := msg.DayKey
		if dKey == "" {
			continue
		}
		st, ok := statMap[dKey]
		if !ok {
			st = &ctxmgr.ChatDateStat{
				DayKey:     dKey,
				Categories: make(map[string]int),
			}
			statMap[dKey] = st
			orderedKeys = append(orderedKeys, dKey)
		}
		st.Count++
		cat := msg.Category
		if cat == "" {
			cat = ctxmgr.CategoryUncategorized
		}
		st.Categories[cat]++
	}
	sort.Strings(orderedKeys)
	res := make([]ctxmgr.ChatDateStat, 0, len(orderedKeys))
	for _, k := range orderedKeys {
		res = append(res, *statMap[k])
	}
	return res
}

func aggregateCategoryStats(msgs []ctxmgr.ChatMessageRecord) []ctxmgr.ChatCategoryStat {
	statMap := make(map[string]*ctxmgr.ChatCategoryStat)
	for _, msg := range msgs {
		cat := msg.Category
		if cat == "" {
			cat = ctxmgr.CategoryUncategorized
		}
		st, ok := statMap[cat]
		if !ok {
			st = &ctxmgr.ChatCategoryStat{
				Category: cat,
				Label:    ctxmgr.GetCategoryLabel(cat),
			}
			statMap[cat] = st
		}
		st.Count++
		if st.FirstMessageID == 0 || msg.ID < st.FirstMessageID {
			st.FirstMessageID = msg.ID
		}
		if msg.ID > st.LastMessageID {
			st.LastMessageID = msg.ID
		}
	}
	res := make([]ctxmgr.ChatCategoryStat, 0, len(statMap))
	for _, st := range statMap {
		res = append(res, *st)
	}
	sort.Slice(res, func(i, j int) bool {
		return res[i].Count > res[j].Count
	})
	return res
}