xinmi/internal/application/repository/message_memory_policy.go

package repository

import (
	"context"
	"errors"

	"github.com/Tencent/WeKnora/internal/types"
	"github.com/Tencent/WeKnora/internal/types/interfaces"
	"gorm.io/gorm"
)

// MemoryAllowedMessageIDs uses per-turn Agent identity; a session can switch
// Agents and its last-request settings cannot authorize older messages.
func (r *messageRepository) MemoryAllowedMessageIDs(ctx context.Context, scope interfaces.MemoryScope, sessionID string, messages []*types.Message) (map[string]bool, error) {
	allowed := map[string]bool{}
	owned, err := r.MemoryOwnedSessionIDs(ctx, scope, []string{sessionID})
	if err != nil || !owned[sessionID] {
		return allowed, err
	}
	for _, input := range messages {
		if input == nil || input.Role != "user" || input.SessionID != sessionID {
			continue
		}
		var answer types.Message
		query := r.db.WithContext(ctx).Where("session_id = ? AND role = ?", sessionID, "assistant")
		if input.RequestID != "" {
			query = query.Where("request_id = ?", input.RequestID)
		} else {
			// Without a request identity, only an adjacent assistant can safely
			// attribute a historical input. Never borrow another turn's Agent.
			query = r.db.WithContext(ctx).Where("session_id = ? AND created_at > ?", sessionID, input.CreatedAt)
		}
		if err := query.Order("created_at ASC, id ASC").First(&answer).Error; err != nil {
			if errors.Is(err, gorm.ErrRecordNotFound) {
				continue
			}
			return nil, err
		}
		if answer.Role != "assistant" {
			continue
		}
		if !answer.IsCompleted {
			// Keep the checkpoint behind a turn still streaming in this session.
			return nil, errors.New("memory extraction is waiting for a completed turn")
		}
		if answer.AgentID == "" {
			// Historical ordinary-assistant turns have no custom Agent binding.
			allowed[input.ID] = true
			continue
		}
		agentTenant := answer.AgentTenantID
		if agentTenant == 0 {
			agentTenant = scope.TenantID
		}
		var agent types.CustomAgent
		err := r.db.WithContext(ctx).Where("tenant_id = ? AND id = ?", agentTenant, answer.AgentID).First(&agent).Error
		if errors.Is(err, gorm.ErrRecordNotFound) {
			// Built-in defaults may have no tenant override. A soft-deleted
			// override must not be recreated through the built-in registry.
			var deleted int64
			if err := r.db.WithContext(ctx).Unscoped().Model(&types.CustomAgent{}).
				Where("tenant_id = ? AND id = ?", agentTenant, answer.AgentID).Count(&deleted).Error; err != nil {
				return nil, err
			}
			builtin := types.GetBuiltinAgent(answer.AgentID, agentTenant)
			if deleted != 0 || builtin == nil || agentTenant != scope.TenantID {
				continue
			}
			agent = *builtin
		} else if err != nil {
			return nil, err
		}
		if !agent.Enabled || (agent.Config.MemoryEnabled != nil && !*agent.Config.MemoryEnabled) {
			continue
		}
		if agentTenant != scope.TenantID {
			var shares int64
			err := r.db.WithContext(ctx).Table("agent_shares").
				Joins("JOIN organization_tenant_members otm ON otm.organization_id = agent_shares.organization_id").
				Joins("JOIN organizations ON organizations.id = agent_shares.organization_id AND organizations.deleted_at IS NULL").
				Where("agent_shares.agent_id = ? AND agent_shares.source_tenant_id = ? AND agent_shares.deleted_at IS NULL AND otm.tenant_id = ?", answer.AgentID, agentTenant, scope.TenantID).
				Count(&shares).Error
			if err != nil {
				return nil, err
			}
			if shares == 0 {
				continue
			}
		}
		allowed[input.ID] = true
	}
	return allowed, nil
}