xinmi/internal/workbench/service.go

package workbench

import (
	"context"
	"errors"
	"fmt"
	"github.com/Tencent/WeKnora/internal/accesspolicy"
	"sort"
	"strings"
	"time"

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

var ErrForbidden = errors.New("没有此资源的操作权限")
var ErrConflict = errors.New("内容已更新或有任务正在执行，请刷新后重试")
var ErrInvalid = errors.New("请求参数无效")

type Service struct {
	db                    *gorm.DB
	models                interfaces.ModelService
	agents                interfaces.CustomAgentService
	sessions              interfaces.SessionService
	instructions          InstructionResolver
	knowledgeBases        interfaces.KnowledgeBaseService
	knowledge             interfaces.KnowledgeService
	knowledgeShares       interfaces.KBShareService
	accessSnapshot        func(context.Context) (*accesspolicy.Snapshot, error)
	memories              interfaces.MemoryService
	rememberInTransaction func(context.Context, *gorm.DB, types.MemoryItem) (*types.MemoryItem, error)
}

func (s *Service) SetMemoryService(memories interfaces.MemoryService) { s.memories = memories }

func (s *Service) SetTransactionalMemoryWriter(writer func(context.Context, *gorm.DB, types.MemoryItem) (*types.MemoryItem, error)) {
	s.rememberInTransaction = writer
}

func (s *Service) SetAccessSnapshot(loader func(context.Context) (*accesspolicy.Snapshot, error)) {
	s.accessSnapshot = loader
}

func (s *Service) SetKnowledgeServices(bases interfaces.KnowledgeBaseService, knowledge interfaces.KnowledgeService, shares interfaces.KBShareService) {
	s.knowledgeBases, s.knowledge, s.knowledgeShares = bases, knowledge, shares
}

// InstructionResolver reuses the native, access-checked Markdown skill catalog.
// Scripts and sandbox tools are not executed by the content compiler.
type InstructionResolver interface {
	ResolveAgentInstructions(context.Context, *types.CustomAgent) (string, error)
}

func (s *Service) SetInstructionResolver(resolver InstructionResolver) {
	s.instructions = resolver
}

func NewService(db *gorm.DB) *Service {
	return &Service{db: db}
}

// SetDependencies is called once after the full WeKnora container has built its
// native model, agent and session services. Keeping construction DB-only avoids
// an initialization cycle with those services.
func (s *Service) SetDependencies(models interfaces.ModelService, agents interfaces.CustomAgentService, sessions interfaces.SessionService) {
	s.models, s.agents, s.sessions = models, agents, sessions
}
func identity(ctx context.Context) (uint64, string, bool, error) {
	tenant, ok := types.TenantIDFromContext(ctx)
	user, userOK := types.UserIDFromContext(ctx)
	if !ok || tenant == 0 || !userOK || types.IsSyntheticUserID(user) {
		return 0, "", false, ErrForbidden
	}
	return tenant, user, types.IsSystemAdminFromContext(ctx) || types.TenantRoleFromContext(ctx).HasPermission(types.TenantRoleAdmin), nil
}
func scope(ctx context.Context, db *gorm.DB) (*gorm.DB, error) {
	tenant, user, admin, e := identity(ctx)
	if e != nil {
		return nil, e
	}
	q := db.Where("tenant_id = ?", tenant)
	if !admin {
		q = q.Where("owner_id = ? OR (kind = 'flows' AND id IN ?)", user, accesspolicy.SharedIDs(ctx, "flows"))
		if hidden := accesspolicy.HiddenIDs(ctx, "flows"); len(hidden) > 0 {
			q = q.Where("kind <> 'flows' OR id NOT IN ?", hidden)
		}
	}
	return q, nil
}
func validKind(k string) bool { return k == "training-maps" || k == "content-tasks" || k == "flows" }
func (s *Service) get(ctx context.Context, kind, id string) (*Resource, error) {
	if !validKind(kind) {
		return nil, ErrInvalid
	}
	q, e := scope(ctx, s.db.WithContext(ctx))
	if e != nil {
		return nil, e
	}
	var r Resource
	e = q.Where("id = ? AND kind = ?", id, kind).First(&r).Error
	return &r, e
}
func (s *Service) save(ctx context.Context, r *Resource, expected int) error {
	tenant, user, _, e := identity(ctx)
	if e != nil {
		return e
	}
	if tenant != r.TenantID {
		return ErrForbidden
	}
	e = s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		lease, hasLease := ctx.Value(leaseContextKey{}).(jobLease)
		lockDependencies := ctx.Value(configurationSaveKey{}) == true || hasLease
		if hasLease && ctx.Value(configurationSaveKey{}) != true && r.Kind == "content-tasks" {
			// Use the stored action; reviewing prose does not depend on source configuration.
			var job Job
			q := tx.Select("action").Where("id = ? AND tenant_id = ? AND owner_id = ? AND resource_id = ? AND kind = ? AND lease_token = ? AND status = ? AND lease_until > ?", lease.ID, tenant, user, r.ID, r.Kind, lease.Token, "running", time.Now()).Limit(1).Find(&job)
			if q.Error != nil {
				return q.Error
			}
			if q.RowsAffected != 1 || lease.Token == "" {
				return ErrConflict
			}
			lockDependencies = !taskActionUsesOnlySavedContent(job.Action)
		}
		if lockDependencies {
			if err := lockReferences(ctx, tx, r); err != nil {
				return err
			}
		}
		current, err := lockResource(ctx, tx, r.Kind, r.ID)
		if err != nil {
			return err
		}
		if current.Version != expected {
			return ErrConflict
		}
		if err = validateLegacyMetadataUpdate(current.Data, r.Data); err != nil {
			return err
		}
		r.Data = resourceData(current, r.Data)
		discardLegacyFlowKnowledge(r)
		if err = writableResource(ctx, tx, current); err != nil {
			return err
		}
		if ctx.Value(configurationSaveKey{}) == true {
			if _, imported := r.Data["legacy_import"]; imported {
				_, actor, _, _ := identity(ctx)
				r.Data["legacy_activation"] = JSON{"status": "active", "configured": true, "configured_by": actor, "configured_at": time.Now().UTC().Format(time.RFC3339), "version": expected + 1}
			}
		}
		if err = archiveResource(ctx, tx, current); err != nil {
			return err
		}
		if err = persistResource(ctx, tx, r, expected); err != nil {
			return err
		}
		if err = syncContentFeedbackMemories(tx, current, r); err != nil {
			return err
		}
		if lease, ok := ctx.Value(leaseContextKey{}).(jobLease); ok {
			return finishConsumption(tx, &Job{ID: lease.ID, TenantID: tenant, OwnerID: r.OwnerID}, true)
		}
		return nil
	})
	if e == nil {
		r.pendingMemory = nil
	}
	return e
}

func (s *Service) create(ctx context.Context, r *Resource) error {
	tenant, user, _, err := identity(ctx)
	if err != nil {
		return err
	}
	if r.TenantID != tenant || r.OwnerID != user || !validKind(r.Kind) {
		return ErrForbidden
	}
	return s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		if err := lockReferences(ctx, tx, r); err != nil {
			return err
		}
		discardLegacyFlowKnowledge(r)
		return tx.Create(r).Error
	})
}

func discardLegacyFlowKnowledge(r *Resource) {
	if r.Kind != "flows" {
		return
	}
	// Canvas knowledge comes from node Agents, including for rename-only saves.
	for _, key := range []string{"knowledge_base_ids", "knowledge_ids", "knowledge_folders", "knowledge_tags"} {
		delete(r.Data, key)
	}
}

// Lock parents before a content task so deletion cannot race a new binding.
func lockReferences(ctx context.Context, tx *gorm.DB, r *Resource) error {
	if r.Kind != "content-tasks" {
		return nil
	}
	d, err := decode[TaskData](r.Data)
	if err != nil || d.FlowID == "" {
		return ErrInvalid
	}
	refs := []struct{ id, kind string }{{d.FlowID, "flows"}}
	if d.MapID != "" {
		refs = append(refs, struct{ id, kind string }{d.MapID, "training-maps"})
	}
	sort.Slice(refs, func(i, j int) bool { return refs[i].id < refs[j].id })
	for _, ref := range refs {
		q, err := scope(ctx, tx)
		if err != nil {
			return err
		}
		var parent Resource
		if err = q.Clauses(clause.Locking{Strength: "SHARE"}).Where("id = ? AND kind = ?", ref.id, ref.kind).First(&parent).Error; err != nil {
			return err
		}
		if err = requireExecutable(&parent); err != nil {
			return err
		}
		if parent.Kind == "training-maps" && parent.Status != "confirmed" {
			return fmt.Errorf("%w: 请先确认训练营导图", ErrInvalid)
		}
	}
	return nil
}

func archiveResource(ctx context.Context, tx *gorm.DB, r *Resource) error {
	_, actor, _, err := identity(ctx)
	if err != nil {
		return err
	}
	return tx.Clauses(clause.OnConflict{Columns: []clause.Column{{Name: "resource_id"}, {Name: "version"}}, DoNothing: true}).Create(&Revision{ID: uuid.NewString(), TenantID: r.TenantID, ResourceID: r.ID, Version: r.Version, ActorID: actor, Data: r.Data, CreatedAt: time.Now()}).Error
}

// Tenant administrators may edit resources without changing their ownership.
func lockResource(ctx context.Context, tx *gorm.DB, kind, id string) (*Resource, error) {
	if !validKind(kind) {
		return nil, ErrInvalid
	}
	tenant, user, admin, err := identity(ctx)
	if err != nil {
		return nil, err
	}
	q := tx.Where("tenant_id = ?", tenant)
	if !admin {
		q = q.Where("owner_id = ?", user)
	}
	var r Resource
	err = q.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id = ? AND kind = ?", id, kind).First(&r).Error
	return &r, err
}

func writableResource(ctx context.Context, tx *gorm.DB, r *Resource) error {
	lease, hasLease := ctx.Value(leaseContextKey{}).(jobLease)
	if !hasLease {
		return requireIdle(tx, r.TenantID, r.ID)
	}
	_, user, _, err := identity(ctx)
	if err != nil {
		return err
	}
	var j Job
	q := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id = ? AND tenant_id = ? AND owner_id = ? AND resource_id = ? AND kind = ? AND lease_token = ? AND status = ? AND lease_until > ?", lease.ID, r.TenantID, user, r.ID, r.Kind, lease.Token, "running", time.Now()).Limit(1).Find(&j)
	if q.Error != nil {
		return q.Error
	}
	if q.RowsAffected != 1 || lease.Token == "" || j.LeaseUntil == nil || !j.LeaseUntil.After(time.Now()) {
		return ErrConflict
	}
	return nil
}

// Caller holds the resource lock and has validated either idleness or a lease.
func persistResource(ctx context.Context, tx *gorm.DB, r *Resource, expected int) error {
	tenant, user, admin, err := identity(ctx)
	if err != nil {
		return err
	}
	if r.TenantID != tenant || (r.OwnerID != user && !admin) {
		return ErrForbidden
	}
	now := time.Now()
	next := expected + 1
	q := tx.Model(&Resource{}).Where("id = ? AND tenant_id = ? AND owner_id = ? AND kind = ? AND version = ?", r.ID, tenant, r.OwnerID, r.Kind, expected)
	if lease, ok := ctx.Value(leaseContextKey{}).(jobLease); ok {
		validLease := tx.Session(&gorm.Session{NewDB: true}).Model(&Job{}).Select("1").Where("id = ? AND tenant_id = ? AND owner_id = ? AND resource_id = ? AND kind = ? AND lease_token = ? AND status = ?", lease.ID, tenant, user, r.ID, r.Kind, lease.Token, "running")
		if tx.Dialector.Name() == "postgres" {
			validLease = validLease.Where("lease_until > clock_timestamp()")
		} else {
			validLease = validLease.Where("lease_until > ?", now)
		}
		q = q.Where("EXISTS (?)", validLease)
	}
	q = q.Updates(map[string]any{"name": r.Name, "data": r.Data, "status": r.Status, "version": next, "updated_at": now})
	if q.Error != nil {
		return q.Error
	}
	if q.RowsAffected != 1 {
		return ErrConflict
	}
	if err = tx.Create(&Revision{ID: uuid.NewString(), TenantID: tenant, ResourceID: r.ID, Version: next, ActorID: user, Data: r.Data, CreatedAt: now}).Error; err != nil {
		return err
	}
	r.Version, r.UpdatedAt = next, now
	return nil
}
func (s *Service) event(ctx context.Context, id, action, message, level string) error {
	tenant, user, _, e := identity(ctx)
	if e != nil {
		return e
	}
	return s.db.WithContext(ctx).Create(&Event{ID: uuid.NewString(), TenantID: tenant, OwnerID: user, ResourceID: id, Action: action, Message: message, Level: level, CreatedAt: time.Now()}).Error
}
func (s *Service) validateAgent(ctx context.Context, id string) error {
	return s.validateAgentDuty(ctx, id, "")
}
func (s *Service) validateAgentDuty(ctx context.Context, id, duty string) error {
	if id == "" {
		return fmt.Errorf("%w: 请先选择 Agent", ErrInvalid)
	}
	if s.agents == nil {
		return errors.New("Agent 服务未初始化")
	}
	a, e := s.agents.GetAgentByID(ctx, id)
	if e != nil {
		return fmt.Errorf("%w: Agent 不存在或无权使用", ErrInvalid)
	}
	tenant, _, _, _ := identity(ctx)
	if a == nil || a.TenantID != tenant {
		return ErrForbidden
	}
	if duty != "" && !a.Config.SupportsDuty(duty) {
		return fmt.Errorf("%w: Agent 职责与当前任务不匹配", ErrInvalid)
	}
	return nil
}
func (s *Service) validate(ctx context.Context, r *Resource) error {
	if strings.TrimSpace(r.Name) == "" || len([]rune(r.Name)) > 200 {
		return fmt.Errorf("%w: 名称必填且最多 200 字", ErrInvalid)
	}
	switch r.Kind {
	case "training-maps":
		d, e := decode[MapData](r.Data)
		if e != nil {
			return ErrInvalid
		}
		if d.ReadingMode != "" && d.ReadingMode != "quick" && d.ReadingMode != "full" {
			return fmt.Errorf("%w: 资料阅读模式无效", ErrInvalid)
		}
		if d.DaysCount < 0 || d.DaysCount > 60 || strings.TrimSpace(d.Topic) == "" {
			return fmt.Errorf("%w: 主题必填，天数可留空由 AI 推荐，或填写 1 到 60", ErrInvalid)
		}
		if e = s.validateKnowledgeScope(ctx, d.KBIDs, d.KnowledgeIDs, d.KnowledgeTags); e != nil {
			return e
		}
		return s.validateAgentDuty(ctx, d.AgentID, types.AgentRoleTrainingMap)
	case "flows":
		d, e := decode[FlowData](r.Data)
		if e != nil {
			return ErrInvalid
		}
		if e = validateFlow(d); e != nil {
			return e
		}
		for _, n := range d.Nodes {
			if e = s.validateAgentDuty(ctx, n.AgentID, types.AgentRoleProduction); e != nil {
				return e
			}
		}
		return nil
	case "content-tasks":
		d, e := decode[TaskData](r.Data)
		if e != nil {
			return ErrInvalid
		}
		if e = validateTaskDuration(d.DurationMinutes); e != nil {
			return e
		}
		if e = d.QualitySettings.validate(); e != nil {
			return e
		}
		ctx, e = s.taskExecutionContext(ctx, d)
		if e != nil {
			return e
		}
		if d.ReadingMode != "" && d.ReadingMode != "quick" && d.ReadingMode != "full" {
			return fmt.Errorf("%w: 资料阅读模式无效", ErrInvalid)
		}
		if strings.TrimSpace(d.SourceText) == "" && d.SourceImage == "" && d.MapID == "" {
			return fmt.Errorf("%w: 请选择已确认导图，或粘贴框架文字、上传导图图片", ErrInvalid)
		}
		if d.SourceImage != "" {
			if _, e = decodeSourceImage(d.SourceImage); e != nil {
				return e
			}
		}
		if d.MapID != "" {
			m, e := s.get(ctx, "training-maps", d.MapID)
			if e != nil {
				return fmt.Errorf("%w: 导图不存在或无权使用", ErrInvalid)
			}
			if e = requireExecutable(m); e != nil {
				return e
			}
			md, e := decode[MapData](m.Data)
			if e != nil {
				return fmt.Errorf("%w: 导图配置格式不完整", ErrInvalid)
			}
			if m.Status != "confirmed" || d.Day < 1 || d.Day > len(md.Days) {
				return fmt.Errorf("%w: 请选择已确认导图中的对应天数", ErrInvalid)
			}
		}
		flow, e := s.taskFlow(ctx, d)
		if e != nil {
			return e
		}
		if e = s.inheritFlow(ctx, &d, flow); e != nil {
			return e
		}
		if d.FlowSnapshot == nil {
			d.FlowSnapshot = &flow
		}
		for i := range d.Segments {
			if d.Segments[i].AgentID == "" {
				d.Segments[i].AgentID = d.AgentID
			}
			if e = s.validateAgentDuty(ctx, d.Segments[i].AgentID, types.AgentRoleProduction); e != nil {
				_, imported := r.Data["legacy_import"]
				if !imported || ctx.Value(configurationSaveKey{}) != true || r.Data["legacy_activation"] != nil {
					return e
				}
				d.Segments[i].AgentID = d.AgentID
				if d.Segments[i].Content != "" && d.Segments[i].Status != "skipped" {
					d.Segments[i].Status = "needs_review"
					r.Status = "needs_review"
				}
			}
		}
		r.Data = resourceData(r, d)
		return nil
	}
	return ErrInvalid
}

func taskActionUsesOnlySavedContent(action string) bool {
	return action == "accept" || action == "skip" || action == "review" || action == "review-settings" || action == "compose"
}

func (s *Service) validateTaskAction(ctx context.Context, r *Resource, action string) error {
	// Reviewing existing prose neither retrieves source material nor regenerates it.
	if taskActionUsesOnlySavedContent(action) {
		d, err := decode[TaskData](r.Data)
		if err != nil || (action != "review-settings" && len(d.Segments) == 0) {
			return fmt.Errorf("%w: 请先拆解内容段", ErrInvalid)
		}
		return nil
	}
	return s.validateTaskExecution(ctx, r)
}

func (s *Service) validateTaskExecution(ctx context.Context, r *Resource) error {
	d, err := decode[TaskData](r.Data)
	if err != nil {
		return ErrInvalid
	}
	if d.FlowID == "" || (d.MapID == "" && strings.TrimSpace(d.SourceText) == "" && d.SourceImage == "") || (d.MapID != "" && d.Day < 1) {
		return fmt.Errorf("%w: 任务缺少导图或框架来源、日期、画布配置，请先补全配置；历史正文、质检和成稿仍保留", ErrInvalid)
	}
	return s.validate(ctx, r)
}

func validateFlow(f FlowData) error {
	if len(f.Nodes) == 0 || len(f.Nodes) > 100 {
		return fmt.Errorf("%w: 画布需要 1 到 100 个节点", ErrInvalid)
	}
	nodes := map[string]bool{}
	degree := map[string]int{}
	adj := map[string][]string{}
	for _, n := range f.Nodes {
		if n.ID == "" || nodes[n.ID] {
			return fmt.Errorf("%w: 节点编号缺失或重复", ErrInvalid)
		}
		nodes[n.ID] = true
		degree[n.ID] = 0
	}
	for _, e := range f.Edges {
		if !nodes[e.Source] || !nodes[e.Target] {
			return fmt.Errorf("%w: 连线节点不存在", ErrInvalid)
		}
		adj[e.Source] = append(adj[e.Source], e.Target)
		degree[e.Target]++
	}
	queue := []string{}
	for id, d := range degree {
		if d == 0 {
			queue = append(queue, id)
		}
	}
	visited := 0
	for len(queue) > 0 {
		id := queue[0]
		queue = queue[1:]
		visited++
		for _, n := range adj[id] {
			degree[n]--
			if degree[n] == 0 {
				queue = append(queue, n)
			}
		}
	}
	if visited != len(nodes) {
		return fmt.Errorf("%w: 画布连线不能形成循环", ErrInvalid)
	}
	return nil
}