FranklinWH-Toolbox/service/repositories/mongo/agentSessionRepo.go

package mongo

import (
	"context"
	"errors"
	"time"

	"fwh-toolbox-server/database"

	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

// 新 AI 助手（AgentServer）的会话归属与子资源索引，见 docs/ai-assistant-agentserver/technical-design.md §3.3。
const (
	CollectionName_AgentSessions  = "AgentSessions"
	CollectionName_AgentResources = "AgentResources"
)

// AgentResourceTTL 子资源索引保留期（与会话保留期一致）
const AgentResourceTTL = 180 * 24 * time.Hour

// AgentSession 会话归属记录
type AgentSession struct {
	SessionId  string     `bson:"sessionId"`
	UserId     string     `bson:"userId"`
	ProfileId  string     `bson:"profileId"`
	Title      string     `bson:"title"`
	CreatedAt  time.Time  `bson:"createdAt"`
	LastActive time.Time  `bson:"lastActive"`
	DeletedAt  *time.Time `bson:"deletedAt,omitempty"`
}

// AgentResource 子资源（run / 交互 / 文件）所属会话
type AgentResource struct {
	ResourceId string    `bson:"resourceId"`
	SessionId  string    `bson:"sessionId"`
	CreatedAt  time.Time `bson:"createdAt"`
}

// InitAgentIndexes 创建索引：会话 ID 唯一、按用户列表、子资源 ID 唯一、子资源 TTL
func InitAgentIndexes(ctx context.Context) error {
	if _, err := database.GetCollection(CollectionName_AgentSessions).Indexes().CreateMany(ctx, []mongo.IndexModel{
		{Keys: bson.D{{Key: "sessionId", Value: 1}}, Options: options.Index().SetUnique(true)},
		{Keys: bson.D{{Key: "userId", Value: 1}, {Key: "deletedAt", Value: 1}, {Key: "lastActive", Value: -1}}},
	}); err != nil {
		return err
	}
	_, err := database.GetCollection(CollectionName_AgentResources).Indexes().CreateMany(ctx, []mongo.IndexModel{
		{Keys: bson.D{{Key: "resourceId", Value: 1}}, Options: options.Index().SetUnique(true)},
		{Keys: bson.D{{Key: "createdAt", Value: 1}}, Options: options.Index().SetExpireAfterSeconds(int32(AgentResourceTTL / time.Second))},
	})
	return err
}

// InsertAgentSession 写入新会话
func InsertAgentSession(ctx context.Context, s AgentSession) error {
	_, err := database.GetCollection(CollectionName_AgentSessions).InsertOne(ctx, s)
	return err
}

// GetAgentSession 查询会话（含已删除）；不存在返回 nil
func GetAgentSession(ctx context.Context, sessionId string) (*AgentSession, error) {
	var s AgentSession
	err := database.GetCollection(CollectionName_AgentSessions).FindOne(ctx, bson.M{"sessionId": sessionId}).Decode(&s)
	if errors.Is(err, mongo.ErrNoDocuments) {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}
	return &s, nil
}

// ListAgentSessionsByUser 用户未删除的会话，按 lastActive 倒序
func ListAgentSessionsByUser(ctx context.Context, userId string, limit int64) ([]AgentSession, error) {
	cur, err := database.GetCollection(CollectionName_AgentSessions).Find(ctx,
		bson.M{"userId": userId, "deletedAt": bson.M{"$exists": false}},
		options.Find().SetSort(bson.D{{Key: "lastActive", Value: -1}}).SetLimit(limit))
	if err != nil {
		return nil, err
	}
	defer cur.Close(ctx)
	list := make([]AgentSession, 0)
	if err := cur.All(ctx, &list); err != nil {
		return nil, err
	}
	return list, nil
}

// UpdateAgentSessionTitle 修改标题（只改本人的未删除会话）
func UpdateAgentSessionTitle(ctx context.Context, userId, sessionId, title string) error {
	_, err := database.GetCollection(CollectionName_AgentSessions).UpdateOne(ctx,
		bson.M{"userId": userId, "sessionId": sessionId, "deletedAt": bson.M{"$exists": false}},
		bson.M{"$set": bson.M{"title": title}})
	return err
}

// TouchAgentSession 刷新 lastActive
func TouchAgentSession(ctx context.Context, userId, sessionId string) error {
	_, err := database.GetCollection(CollectionName_AgentSessions).UpdateOne(ctx,
		bson.M{"userId": userId, "sessionId": sessionId, "deletedAt": bson.M{"$exists": false}},
		bson.M{"$set": bson.M{"lastActive": time.Now()}})
	return err
}

// SoftDeleteAgentSession 软删除
func SoftDeleteAgentSession(ctx context.Context, userId, sessionId string) error {
	_, err := database.GetCollection(CollectionName_AgentSessions).UpdateOne(ctx,
		bson.M{"userId": userId, "sessionId": sessionId, "deletedAt": bson.M{"$exists": false}},
		bson.M{"$set": bson.M{"deletedAt": time.Now()}})
	return err
}

// GetAgentResourceSession 子资源所属会话；不存在返回空串
func GetAgentResourceSession(ctx context.Context, resourceId string) (string, error) {
	var r AgentResource
	err := database.GetCollection(CollectionName_AgentResources).FindOne(ctx, bson.M{"resourceId": resourceId}).Decode(&r)
	if errors.Is(err, mongo.ErrNoDocuments) {
		return "", nil
	}
	if err != nil {
		return "", err
	}
	return r.SessionId, nil
}

// UpsertAgentResources 批量写入子资源索引；已存在的不覆盖所属会话
func UpsertAgentResources(ctx context.Context, rs []AgentResource) error {
	if len(rs) == 0 {
		return nil
	}
	models := make([]mongo.WriteModel, 0, len(rs))
	for _, r := range rs {
		models = append(models, mongo.NewUpdateOneModel().
			SetFilter(bson.M{"resourceId": r.ResourceId}).
			SetUpdate(bson.M{"$setOnInsert": r}).
			SetUpsert(true))
	}
	_, err := database.GetCollection(CollectionName_AgentResources).BulkWrite(ctx, models, options.BulkWrite().SetOrdered(false))
	return err
}