Agent-System-Dev/tests/load/collaboration/seed/main.go

// 协作存储 U7/B5：大数据量 seed 工具（方案 §3 数据量档）。
//
// 直连隔离 PG（不经 HTTP），用 COPY 灌 messages/events/knowledge。
// 用户须已由 run_load.sh 经 CLI 创建（口令哈希与现网一致）；本工具只写
// 项目/会话拓扑与体量数据。不写 knowledge_vectors（向量扫描面向数万块）。
//
// 用法：
//   seed --dsn 'postgresql://…/u7load_app' --profile tiny|contract [--manifest path]
package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"os"
	"time"

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

type profile struct {
	users      int
	projects   int // per user
	convs      int // per project
	messages   int
	chunks     int
	sharedPct  int // % of knowledge with owner=''
}

var profiles = map[string]profile{
	// 接线冒烟：秒级。
	"tiny": {users: 3, projects: 2, convs: 2, messages: 1000, chunks: 1000, sharedPct: 10},
	// 负载合同数据量档（约需数分钟与数 GB 磁盘，视机器而定）。
	"contract": {users: 10, projects: 5, convs: 10, messages: 1_000_000, chunks: 100_000, sharedPct: 10},
}

type manifest struct {
	Profile       string    `json:"profile"`
	Owners        []string  `json:"owners"`
	Projects      int       `json:"projects"`
	Conversations int       `json:"conversations"`
	Messages      int       `json:"messages"`
	Events        int       `json:"events"`
	Knowledge     int       `json:"knowledge"`
	SharedChunks  int       `json:"sharedChunks"`
	ElapsedMS     int64     `json:"elapsedMs"`
	StartedAt     time.Time `json:"startedAt"`
	FinishedAt    time.Time `json:"finishedAt"`
}

func main() {
	dsn := flag.String("dsn", "", "PostgreSQL DSN（必填；目标须已迁移）")
	profName := flag.String("profile", "tiny", "tiny|contract")
	manifestPath := flag.String("manifest", "", "可选：写出 seed-manifest.json")
	usersOverride := flag.Int("users", 0, "覆盖用户数（0=用 profile）")
	projectsOverride := flag.Int("projects", 0, "覆盖每用户项目数")
	messagesOverride := flag.Int("messages", 0, "覆盖消息总数")
	chunksOverride := flag.Int("chunks", 0, "覆盖知识分块总数")
	flag.Parse()

	if *dsn == "" {
		fmt.Fprintln(os.Stderr, "seed: --dsn required")
		os.Exit(2)
	}
	p, ok := profiles[*profName]
	if !ok {
		fmt.Fprintf(os.Stderr, "seed: unknown profile %q (tiny|contract)\n", *profName)
		os.Exit(2)
	}
	if *usersOverride > 0 {
		p.users = *usersOverride
	}
	if *projectsOverride > 0 {
		p.projects = *projectsOverride
	}
	if *messagesOverride > 0 {
		p.messages = *messagesOverride
	}
	if *chunksOverride > 0 {
		p.chunks = *chunksOverride
	}

	ctx := context.Background()
	pool, err := pgxpool.New(ctx, *dsn)
	if err != nil {
		fmt.Fprintln(os.Stderr, "seed: connect:", err)
		os.Exit(1)
	}
	defer pool.Close()

	if err := requireSchema(ctx, pool); err != nil {
		fmt.Fprintln(os.Stderr, "seed:", err)
		os.Exit(1)
	}

	owners := make([]string, p.users)
	for i := 0; i < p.users; i++ {
		owners[i] = fmt.Sprintf("seed_u%d", i)
	}
	if err := requireOwners(ctx, pool, owners); err != nil {
		fmt.Fprintln(os.Stderr, "seed:", err)
		os.Exit(1)
	}

	started := time.Now().UTC()
	fmt.Printf("[seed] profile=%s users=%d projects/user=%d convs/project=%d messages=%d chunks=%d\n",
		*profName, p.users, p.projects, p.convs, p.messages, p.chunks)

	projIDs, convIDs, err := insertTopology(ctx, pool, owners, p.projects, p.convs)
	if err != nil {
		fmt.Fprintln(os.Stderr, "seed: topology:", err)
		os.Exit(1)
	}
	fmt.Printf("[seed] topology: projects=%d conversations=%d\n", len(projIDs), len(convIDs))

	nMsg, err := copyMessagesEvents(ctx, pool, convIDs, p.messages)
	if err != nil {
		fmt.Fprintln(os.Stderr, "seed: messages/events:", err)
		os.Exit(1)
	}
	fmt.Printf("[seed] messages=%d events=%d\n", nMsg, nMsg)

	nKnw, nShared, err := copyKnowledge(ctx, pool, projIDs, p.chunks, p.sharedPct)
	if err != nil {
		fmt.Fprintln(os.Stderr, "seed: knowledge:", err)
		os.Exit(1)
	}
	fmt.Printf("[seed] knowledge=%d (shared owner=''=%d)\n", nKnw, nShared)

	finished := time.Now().UTC()
	m := manifest{
		Profile:       *profName,
		Owners:        owners,
		Projects:      len(projIDs),
		Conversations: len(convIDs),
		Messages:      nMsg,
		Events:        nMsg,
		Knowledge:     nKnw,
		SharedChunks:  nShared,
		ElapsedMS:     finished.Sub(started).Milliseconds(),
		StartedAt:     started,
		FinishedAt:    finished,
	}
	if *manifestPath != "" {
		b, _ := json.MarshalIndent(m, "", "  ")
		if err := os.WriteFile(*manifestPath, b, 0o644); err != nil {
			fmt.Fprintln(os.Stderr, "seed: write manifest:", err)
			os.Exit(1)
		}
		fmt.Printf("[seed] manifest=%s\n", *manifestPath)
	}
	fmt.Printf("[seed] done in %dms\n", m.ElapsedMS)
}

func requireSchema(ctx context.Context, pool *pgxpool.Pool) error {
	for _, t := range []string{"projects", "conversations", "messages", "events", "knowledge", "users"} {
		var one int
		if err := pool.QueryRow(ctx, `SELECT 1 FROM information_schema.tables WHERE table_schema='public' AND table_name=$1`, t).Scan(&one); err != nil {
			return fmt.Errorf("schema missing table %s (run migrations via agent-system user create first): %w", t, err)
		}
	}
	return nil
}

func requireOwners(ctx context.Context, pool *pgxpool.Pool, owners []string) error {
	for _, o := range owners {
		var one int
		err := pool.QueryRow(ctx, `SELECT 1 FROM users WHERE owner=$1`, o).Scan(&one)
		if err != nil {
			return fmt.Errorf("user %q not found — run_load.sh must CLI-create seed_u* before seed", o)
		}
	}
	return nil
}

type projectRow struct {
	id, owner string
}

func insertTopology(ctx context.Context, pool *pgxpool.Pool, owners []string, projectsPerUser, convsPerProject int) (projects []projectRow, convIDs []string, err error) {
	now := time.Now().UTC().Format(time.RFC3339Nano)
	tx, err := pool.Begin(ctx)
	if err != nil {
		return nil, nil, err
	}
	defer tx.Rollback(ctx)

	var projBuf [][]any
	var convBuf [][]any
	for ui, owner := range owners {
		for pi := 0; pi < projectsPerUser; pi++ {
			pid := fmt.Sprintf("prj_seed_%d_%d", ui, pi)
			projBuf = append(projBuf, []any{pid, owner, fmt.Sprintf("seed-proj-%d-%d", ui, pi), "{}", now, now})
			projects = append(projects, projectRow{id: pid, owner: owner})
			for ci := 0; ci < convsPerProject; ci++ {
				cid := fmt.Sprintf("cnv_seed_%d_%d_%d", ui, pi, ci)
				convBuf = append(convBuf, []any{cid, pid, owner, fmt.Sprintf("seed-conv-%d-%d-%d", ui, pi, ci), "", now, now})
				convIDs = append(convIDs, cid)
			}
		}
	}
	if _, err := tx.CopyFrom(ctx,
		pgx.Identifier{"projects"},
		[]string{"id", "owner", "name", "settings", "created_at", "updated_at"},
		pgx.CopyFromRows(projBuf)); err != nil {
		return nil, nil, fmt.Errorf("projects: %w", err)
	}
	if _, err := tx.CopyFrom(ctx,
		pgx.Identifier{"conversations"},
		[]string{"id", "project_id", "owner", "title", "profile_id", "created_at", "updated_at"},
		pgx.CopyFromRows(convBuf)); err != nil {
		return nil, nil, fmt.Errorf("conversations: %w", err)
	}
	if err := tx.Commit(ctx); err != nil {
		return nil, nil, err
	}
	return projects, convIDs, nil
}

// copyMessagesEvents 均匀分配到各会话，每会话连续 seq 从 1 起；每条消息配对一条 event。
func copyMessagesEvents(ctx context.Context, pool *pgxpool.Pool, convIDs []string, total int) (int, error) {
	if total <= 0 || len(convIDs) == 0 {
		return 0, nil
	}
	const batch = 10000
	now := time.Now().UTC().Format(time.RFC3339Nano)
	perConv := make([]int, len(convIDs))
	for i := 0; i < total; i++ {
		perConv[i%len(convIDs)]++
	}

	written := 0
	msgBuf := make([][]any, 0, batch)
	evBuf := make([][]any, 0, batch)
	flush := func() error {
		if len(msgBuf) == 0 {
			return nil
		}
		tx, err := pool.Begin(ctx)
		if err != nil {
			return err
		}
		defer tx.Rollback(ctx)
		if _, err := tx.CopyFrom(ctx,
			pgx.Identifier{"messages"},
			[]string{"id", "conversation_id", "run_id", "role", "parts", "seq", "created_at", "request_id"},
			pgx.CopyFromRows(msgBuf)); err != nil {
			return fmt.Errorf("messages copy: %w", err)
		}
		if _, err := tx.CopyFrom(ctx,
			pgx.Identifier{"events"},
			[]string{"conversation_id", "seq", "run_id", "ts", "v", "type", "payload"},
			pgx.CopyFromRows(evBuf)); err != nil {
			return fmt.Errorf("events copy: %w", err)
		}
		if err := tx.Commit(ctx); err != nil {
			return err
		}
		msgBuf = msgBuf[:0]
		evBuf = evBuf[:0]
		return nil
	}

	global := 0
	for ci, cid := range convIDs {
		for seq := 1; seq <= perConv[ci]; seq++ {
			global++
			role := "user"
			if seq%2 == 0 {
				role = "assistant"
			}
			parts := fmt.Sprintf(`[{"type":"text","text":"seed msg %d"}]`, global)
			mid := fmt.Sprintf("msg_seed_%d", global)
			msgBuf = append(msgBuf, []any{mid, cid, "", role, parts, seq, now, ""})
			payload := fmt.Sprintf(`{"messageId":%q,"seq":%d}`, mid, seq)
			evBuf = append(evBuf, []any{cid, seq, "", now, 1, "message.upserted", payload})
			written++
			if len(msgBuf) >= batch {
				if err := flush(); err != nil {
					return written, err
				}
				if written%100000 == 0 {
					fmt.Printf("[seed] … messages %d/%d\n", written, total)
				}
			}
		}
	}
	if err := flush(); err != nil {
		return written, err
	}
	return written, nil
}

func copyKnowledge(ctx context.Context, pool *pgxpool.Pool, projects []projectRow, total, sharedPct int) (n, shared int, err error) {
	if total <= 0 || len(projects) == 0 {
		return 0, 0, nil
	}
	const batch = 5000
	now := time.Now().UTC().Format(time.RFC3339Nano)
	buf := make([][]any, 0, batch)
	flush := func() error {
		if len(buf) == 0 {
			return nil
		}
		tx, err := pool.Begin(ctx)
		if err != nil {
			return err
		}
		defer tx.Rollback(ctx)
		if _, err := tx.CopyFrom(ctx,
			pgx.Identifier{"knowledge"},
			[]string{"id", "owner", "project_id", "source", "title", "ref", "chunk", "created_at"},
			pgx.CopyFromRows(buf)); err != nil {
			return err
		}
		if err := tx.Commit(ctx); err != nil {
			return err
		}
		buf = buf[:0]
		return nil
	}

	for i := 0; i < total; i++ {
		pr := projects[i%len(projects)]
		owner := pr.owner
		if sharedPct > 0 && i%100 < sharedPct {
			owner = ""
			shared++
		}
		id := fmt.Sprintf("knw_seed_%d", i+1)
		chunk := fmt.Sprintf("seed knowledge chunk %d about project %s topic %d", i+1, pr.id, i%97)
		buf = append(buf, []any{id, owner, pr.id, "seed", fmt.Sprintf("seed-doc-%d", i+1), id, chunk, now})
		n++
		if len(buf) >= batch {
			if err := flush(); err != nil {
				return n, shared, err
			}
		}
	}
	if err := flush(); err != nil {
		return n, shared, err
	}
	return n, shared, nil
}