FranklinWH-Toolbox/service/services/agentServices.go

package services

import (
	"context"
	"errors"
	"strings"
	"time"

	"fwh-toolbox-server/config"
	mongoRepo "fwh-toolbox-server/repositories/mongo"
	"fwh-toolbox-server/utils/agentserver"

	"go.uber.org/zap"
)

// 新 AI 助手（AgentServer）：会话归属、按角色选择 Profile、代理组装。
// 设计见 docs/ai-assistant-agentserver/technical-design.md §3，与旧助手（aiChatServices.go）相互独立。

// ErrAgentSessionNotFound 会话不存在、已删除或不属于当前用户
var ErrAgentSessionNotFound = errors.New("会话不存在或无权访问")

// 每个用户列出的会话上限
const agentSessionListLimit = 50

// 标题最大长度（字符）
const agentTitleMaxLen = 100

// AgentSessionItem 会话列表项（返回给前端）
type AgentSessionItem struct {
	SessionId  string `json:"sessionId"`
	Title      string `json:"title"`
	ProfileId  string `json:"profileId"`
	CreatedAt  int64  `json:"createdAt"`
	LastActive int64  `json:"lastActive"`
}

// AgentService 新 AI 助手服务
type AgentService struct {
	cfg      config.AgentConfig
	client   *agentserver.Client
	proxy    *agentserver.Proxy
	recorder *agentResourceRecorder
}

// NewAgentService 创建服务并启动子资源索引的异步写入
func NewAgentService(cfg config.AgentConfig) *AgentService {
	client := agentserver.NewClient(cfg.ServerURL, cfg.ApiKey)
	recorder := newAgentResourceRecorder()
	s := &AgentService{cfg: cfg, client: client, recorder: recorder}
	s.proxy = &agentserver.Proxy{
		Client: client,
		Guard:  &agentserver.Guard{Store: agentMongoStore{}, Upstream: client, Sink: recorder},
		Sink:   recorder,
		OnMessage: func(userId, sessionId string) {
			go func() {
				ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
				defer cancel()
				if err := mongoRepo.TouchAgentSession(ctx, userId, sessionId); err != nil {
					zap.L().Warn("[Agent] 刷新会话活跃时间失败", zap.String("sessionId", sessionId), zap.Error(err))
				}
			}()
		},
		OnTitle: func(userId, sessionId, title string) {
			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
			defer cancel()
			if err := mongoRepo.UpdateAgentSessionTitle(ctx, userId, sessionId, clipTitle(title)); err != nil {
				zap.L().Warn("[Agent] 同步会话标题失败", zap.String("sessionId", sessionId), zap.Error(err))
			}
		},
	}
	if cfg.Enabled && !client.Configured() {
		zap.L().Warn("[Agent] 已开启新 AI 助手，但 AGENT_SERVER_URL / AGENT_SERVER_API_KEY 未配置")
	}
	return s
}

// Enabled 新助手对该用户是否可用：总开关打开、已配置 AgentServer、角色允许
func (s *AgentService) Enabled(roles []string) bool {
	return s.cfg.Enabled && s.client.Configured() && agentserver.RoleAllowed(roles, s.cfg.AllowedRoles)
}

// Proxy /agent/v1/* 代理
func (s *AgentService) Proxy() *agentserver.Proxy {
	return s.proxy
}

// ListSessions 本人未删除的会话，按最近活跃倒序
func (s *AgentService) ListSessions(ctx context.Context, userId string) ([]AgentSessionItem, error) {
	list, err := mongoRepo.ListAgentSessionsByUser(ctx, userId, agentSessionListLimit)
	if err != nil {
		return nil, err
	}
	out := make([]AgentSessionItem, 0, len(list))
	for _, r := range list {
		out = append(out, toAgentSessionItem(r))
	}
	return out, nil
}

// CreateSession 按角色选择 Profile，在 AgentServer 创建会话并记录归属
func (s *AgentService) CreateSession(ctx context.Context, userId string, roles []string) (*AgentSessionItem, error) {
	profileId := agentserver.SelectProfile(roles, s.cfg.ProfileByRole, s.cfg.ProfileDefault)
	dto, err := s.client.CreateSession(ctx, profileId)
	if err != nil {
		return nil, err
	}
	if dto.ProfileId != nil && *dto.ProfileId != "" {
		profileId = *dto.ProfileId
	}
	now := time.Now()
	rec := mongoRepo.AgentSession{
		SessionId:  dto.Id,
		UserId:     userId,
		ProfileId:  profileId,
		Title:      dto.Title,
		CreatedAt:  now,
		LastActive: now,
	}
	if err := mongoRepo.InsertAgentSession(ctx, rec); err != nil {
		// 归属写入失败时删除上游会话，避免出现无主会话
		if delErr := s.client.DeleteSession(context.Background(), dto.Id); delErr != nil {
			zap.L().Warn("[Agent] 回滚上游会话失败", zap.String("sessionId", dto.Id), zap.Error(delErr))
		}
		return nil, err
	}
	zap.L().Info("[Agent] 创建会话", zap.String("userId", userId), zap.String("sessionId", dto.Id), zap.String("profileId", profileId))
	item := toAgentSessionItem(rec)
	return &item, nil
}

// RenameSession 修改标题
func (s *AgentService) RenameSession(ctx context.Context, userId, sessionId, title string) error {
	if err := s.ensureOwned(ctx, userId, sessionId); err != nil {
		return err
	}
	title = clipTitle(title)
	if err := s.client.PatchSession(ctx, sessionId, map[string]any{"title": title}); err != nil {
		return err
	}
	return mongoRepo.UpdateAgentSessionTitle(ctx, userId, sessionId, title)
}

// DeleteSession 删除上游会话并软删除归属记录
func (s *AgentService) DeleteSession(ctx context.Context, userId, sessionId string) error {
	if err := s.ensureOwned(ctx, userId, sessionId); err != nil {
		return err
	}
	if err := s.client.DeleteSession(ctx, sessionId); err != nil {
		return err
	}
	return mongoRepo.SoftDeleteAgentSession(ctx, userId, sessionId)
}

func (s *AgentService) ensureOwned(ctx context.Context, userId, sessionId string) error {
	rec, err := mongoRepo.GetAgentSession(ctx, sessionId)
	if err != nil {
		return err
	}
	if rec == nil || rec.DeletedAt != nil || rec.UserId != userId {
		return ErrAgentSessionNotFound
	}
	return nil
}

func toAgentSessionItem(r mongoRepo.AgentSession) AgentSessionItem {
	return AgentSessionItem{
		SessionId:  r.SessionId,
		Title:      r.Title,
		ProfileId:  r.ProfileId,
		CreatedAt:  r.CreatedAt.UnixMilli(),
		LastActive: r.LastActive.UnixMilli(),
	}
}

func clipTitle(t string) string {
	t = strings.TrimSpace(t)
	if r := []rune(t); len(r) > agentTitleMaxLen {
		return string(r[:agentTitleMaxLen])
	}
	return t
}

// ============ 归属存储（Mongo） ============

type agentMongoStore struct{}

func (agentMongoStore) SessionOwner(ctx context.Context, sessionId string) (string, bool, error) {
	rec, err := mongoRepo.GetAgentSession(ctx, sessionId)
	if err != nil || rec == nil || rec.DeletedAt != nil {
		return "", false, err
	}
	return rec.UserId, true, nil
}

func (agentMongoStore) ResourceSession(ctx context.Context, resourceId string) (string, bool, error) {
	sid, err := mongoRepo.GetAgentResourceSession(ctx, resourceId)
	if err != nil {
		return "", false, err
	}
	return sid, sid != "", nil
}

func (agentMongoStore) RecentSessionIds(ctx context.Context, userId string, limit int) ([]string, error) {
	list, err := mongoRepo.ListAgentSessionsByUser(ctx, userId, int64(limit))
	if err != nil {
		return nil, err
	}
	ids := make([]string, 0, len(list))
	for _, r := range list {
		ids = append(ids, r.SessionId)
	}
	return ids, nil
}

// ============ 子资源索引的异步批量写入 ============

const (
	agentRecorderQueue    = 2048
	agentRecorderBatch    = 200
	agentRecorderInterval = 500 * time.Millisecond
)

// agentResourceRecorder 代理转发时解析到的子资源 ID 先进队列，后台按批写入，不阻塞转发
type agentResourceRecorder struct {
	ch chan agentserver.Resource
}

func newAgentResourceRecorder() *agentResourceRecorder {
	r := &agentResourceRecorder{ch: make(chan agentserver.Resource, agentRecorderQueue)}
	go r.loop()
	return r
}

// Record 实现 agentserver.ResourceSink；队列满时丢弃（之后由归属反查补写）
func (r *agentResourceRecorder) Record(rs []agentserver.Resource) {
	for _, res := range rs {
		select {
		case r.ch <- res:
		default:
			zap.L().Warn("[Agent] 子资源索引队列已满，丢弃", zap.String("resourceId", res.Id))
		}
	}
}

func (r *agentResourceRecorder) loop() {
	ticker := time.NewTicker(agentRecorderInterval)
	defer ticker.Stop()
	pending := map[string]string{}
	flush := func() {
		if len(pending) == 0 {
			return
		}
		now := time.Now()
		batch := make([]mongoRepo.AgentResource, 0, len(pending))
		for id, sid := range pending {
			batch = append(batch, mongoRepo.AgentResource{ResourceId: id, SessionId: sid, CreatedAt: now})
		}
		pending = map[string]string{}
		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel()
		if err := mongoRepo.UpsertAgentResources(ctx, batch); err != nil {
			zap.L().Warn("[Agent] 写入子资源索引失败", zap.Int("count", len(batch)), zap.Error(err))
		}
	}
	for {
		select {
		case res := <-r.ch:
			pending[res.Id] = res.SessionId
			if len(pending) >= agentRecorderBatch {
				flush()
			}
		case <-ticker.C:
			flush()
		}
	}
}