Agent-System/internal/projectcontext/service.go

package projectcontext

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"sort"
	"strings"
	"time"

	"github.com/jackc/pgx/v5/pgxpool"

	"chat-agent-system/internal/squad"
)

const (
	AuthorityConfirmedFact = "confirmed_fact"
	AuthorityApprovedDoc   = "approved_document"
	AuthorityUpload        = "upload_material"
	AuthorityMilestone     = "milestone"
	AuthorityDialogue      = "dialogue"
	AuthorityKnowledge     = "knowledge"
)

// Source 来源清单条目（全量枚举，按任务检索后分级注入）。
type Source struct {
	Type            string `json:"type"`
	ID              string `json:"id"`
	Version         string `json:"version,omitempty"`
	ContentHash     string `json:"content_hash"`
	AuthorityLevel  string `json:"authority_level"`
	PermissionScope string `json:"permission_scope"`
	Referenced      bool   `json:"referenced"`
	Snippet         string `json:"snippet,omitempty"`
}

// Snapshot 不可变项目上下文快照。
type Snapshot struct {
	ID                       int64     `json:"id"`
	ProjectID                int64     `json:"project_id"`
	TaskKey                  string    `json:"task_key"`
	CutoffAt                 time.Time `json:"cutoff_at"`
	AggregateHash            string    `json:"aggregate_hash"`
	SourceCount              int       `json:"source_count"`
	ExcludedPrivateCount     int       `json:"excluded_private_count"`
	ExcludedNoAccessCount    int       `json:"excluded_no_access_count"`
	Sources                  []Source  `json:"sources"`
}

// BuildInput 构建快照所需输入。
type BuildInput struct {
	ProjectID   int64
	TaskKey     string
	UserID      int64
	ProjectName string
	Facts       map[string]string
	Materials   []squad.MaterialItem
	PriorDocs   []squad.PriorDocItem
	Knowledge   []squad.KnowledgeHit
	Dialogue    []string
	Decisions   []string
}

// Service 项目上下文快照服务。
type Service struct {
	pool *pgxpool.Pool
}

func NewService(pool *pgxpool.Pool) *Service {
	return &Service{pool: pool}
}

func (s *Service) Available() bool {
	return s != nil && s.pool != nil
}

// BuildSnapshot 建立来源清单并按任务检索相关片段（不 dump 全量原文）。
func (s *Service) BuildSnapshot(ctx context.Context, in BuildInput) (*Snapshot, error) {
	cutoff := time.Now()
	sources, excludedPrivate, excludedNoAccess := enumerateSources(in)
	sources = selectForTask(in.TaskKey, sources)
	aggHash := aggregateHash(sources)

	snap := &Snapshot{
		ProjectID:             in.ProjectID,
		TaskKey:               in.TaskKey,
		CutoffAt:              cutoff,
		AggregateHash:         aggHash,
		SourceCount:           len(sources),
		ExcludedPrivateCount:  excludedPrivate,
		ExcludedNoAccessCount: excludedNoAccess,
		Sources:               sources,
	}
	if !s.Available() || in.ProjectID <= 0 {
		return snap, nil
	}

	var id int64
	err := s.pool.QueryRow(ctx, `
		INSERT INTO project_context_snapshots (project_id, task_key, cutoff_at, aggregate_hash, source_count, created_by, frozen)
		VALUES ($1, $2, $3, $4, $5, $6, true) RETURNING id
	`, in.ProjectID, in.TaskKey, cutoff, aggHash, len(sources), in.UserID).Scan(&id)
	if err != nil {
		return snap, err
	}
	snap.ID = id
	for _, src := range sources {
		_, _ = s.pool.Exec(ctx, `
			INSERT INTO project_context_snapshot_sources
			(snapshot_id, source_type, source_id, source_version, content_hash, authority_level, permission_scope, referenced)
			VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
			ON CONFLICT (snapshot_id, source_type, source_id, source_version) DO NOTHING
		`, id, src.Type, src.ID, src.Version, src.ContentHash, src.AuthorityLevel, src.PermissionScope, src.Referenced)
	}
	return snap, nil
}

// MaterialsForWorkspace 按权威等级与任务目标注入工作区材料（核心固定 + 相关分批）。
func MaterialsForWorkspace(snap *Snapshot, maxItems int) []squad.MaterialItem {
	if snap == nil || len(snap.Sources) == 0 {
		return nil
	}
	if maxItems <= 0 {
		maxItems = 12
	}
	out := make([]squad.MaterialItem, 0, maxItems)
	for _, src := range snap.Sources {
		if !src.Referenced || strings.TrimSpace(src.Snippet) == "" {
			continue
		}
		out = append(out, squad.MaterialItem{
			Filename: fmt.Sprintf("%s_%s.md", src.Type, src.ID),
			Content:  []byte(src.Snippet),
		})
		if len(out) >= maxItems {
			break
		}
	}
	return out
}

func enumerateSources(in BuildInput) ([]Source, int, int) {
	out := make([]Source, 0)
	excludedPrivate := 0
	excludedNoAccess := 0
	for k, v := range in.Facts {
		if strings.TrimSpace(v) == "" {
			continue
		}
		out = append(out, Source{
			Type: AuthorityConfirmedFact, ID: k, ContentHash: hashText(v),
			AuthorityLevel: AuthorityConfirmedFact, PermissionScope: "project",
		})
	}
	for i, m := range in.Materials {
		body := string(m.Content)
		if len(body) == 0 {
			continue
		}
		id := m.Filename
		if id == "" {
			id = fmt.Sprintf("material_%d", i)
		}
		out = append(out, Source{
			Type: "upload_material", ID: id, ContentHash: hashText(body),
			AuthorityLevel: AuthorityUpload, PermissionScope: "project", Snippet: truncate(body, 4000),
		})
	}
	for _, d := range in.PriorDocs {
		body := d.Content
		if body == "" {
			continue
		}
		id := d.DocKey
		if id == "" {
			id = d.Title
		}
		out = append(out, Source{
			Type: "document", ID: id, Version: d.Stage, ContentHash: hashText(body),
			AuthorityLevel: AuthorityApprovedDoc, PermissionScope: "project", Snippet: truncate(body, 4000),
		})
	}
	for i, line := range in.Dialogue {
		trimmed := strings.TrimSpace(line)
		if trimmed == "" {
			excludedPrivate++
			continue
		}
		if strings.HasPrefix(trimmed, "[private]") || strings.HasPrefix(trimmed, "[draft]") {
			excludedPrivate++
			continue
		}
		out = append(out, Source{
			Type: "dialogue", ID: fmt.Sprintf("turn_%d", i), ContentHash: hashText(trimmed),
			AuthorityLevel: AuthorityDialogue, PermissionScope: "project_shared", Snippet: truncate(trimmed, 2000),
		})
	}
	for i, d := range in.Decisions {
		out = append(out, Source{
			Type: "decision", ID: fmt.Sprintf("decision_%d", i), ContentHash: hashText(d),
			AuthorityLevel: AuthorityConfirmedFact, PermissionScope: "project", Snippet: truncate(d, 2000),
		})
	}
	for _, h := range in.Knowledge {
		body := h.Content
		if body == "" {
			body = h.Title
		}
		out = append(out, Source{
			Type: "knowledge", ID: h.CiteKey, ContentHash: hashText(body),
			AuthorityLevel: AuthorityKnowledge, PermissionScope: "global_kb", Snippet: truncate(body, 2000),
		})
	}
	sort.Slice(out, func(i, j int) bool {
		return authorityRank(out[i].AuthorityLevel) < authorityRank(out[j].AuthorityLevel)
	})
	return out, excludedPrivate, excludedNoAccess
}

func selectForTask(taskKey string, sources []Source) []Source {
	if len(sources) == 0 {
		return sources
	}
	keywords := taskKeywords(taskKey)
	for i := range sources {
		src := &sources[i]
		if src.AuthorityLevel == AuthorityConfirmedFact || src.AuthorityLevel == AuthorityApprovedDoc {
			src.Referenced = true
			continue
		}
		if len(keywords) == 0 {
			src.Referenced = src.AuthorityLevel != AuthorityDialogue
			continue
		}
		hay := strings.ToLower(src.Snippet + " " + src.ID + " " + src.Type)
		for _, kw := range keywords {
			if strings.Contains(hay, kw) {
				src.Referenced = true
				break
			}
		}
	}
	return sources
}

func taskKeywords(taskKey string) []string {
	switch taskKey {
	case "system_reqs":
		return []string{"需求", "requirement", "功能", "system_reqs"}
	case "test_manager":
		return []string{"测试", "验证", "风险", "test", "verification", "requirement"}
	default:
		return nil
	}
}

func authorityRank(level string) int {
	switch level {
	case AuthorityConfirmedFact:
		return 0
	case AuthorityApprovedDoc:
		return 1
	case AuthorityUpload:
		return 2
	case AuthorityMilestone:
		return 3
	case AuthorityDialogue:
		return 4
	default:
		return 5
	}
}

func aggregateHash(sources []Source) string {
	raw, _ := json.Marshal(sources)
	sum := sha256.Sum256(raw)
	return hex.EncodeToString(sum[:])
}

func hashText(s string) string {
	sum := sha256.Sum256([]byte(s))
	return hex.EncodeToString(sum[:])
}

func truncate(s string, max int) string {
	if len(s) <= max {
		return s
	}
	return s[:max] + "\n…(truncated)"
}