weichen
4fdbfbc9a2
Details (original message): Add the vector RAG pipeline: paragraph chunker, OpenAI-compatible embeddings and rerank clients, dual-tier vector store with cosine search, and RRF hybrid retrieval
399 lines
12 KiB
Go
399 lines
12 KiB
Go
package store
|
|
|
|
import (
|
|
"database/sql"
|
|
"os"
|
|
"path/filepath"
|
|
"strconv"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
|
|
core "agent-system-core"
|
|
)
|
|
|
|
// openTestStore 双档测试入口：设 AGENT_PG_TEST_DSN 时对 PostgreSQL 档逐用例建独立
|
|
// schema 运行（双档 CI 验证，D16）；默认嵌入式 SQLite 档。
|
|
func openTestStore(t *testing.T) *Store {
|
|
t.Helper()
|
|
dsn := os.Getenv("AGENT_PG_TEST_DSN")
|
|
if dsn == "" {
|
|
st, err := OpenSQLite(filepath.Join(t.TempDir(), "test.db"))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
t.Cleanup(func() { st.Close() })
|
|
return st
|
|
}
|
|
// 每用例独立 schema（随机名），测试后级联清理，避免跨用例/跨运行干扰。
|
|
schema := "t_" + strings.ToLower(newID("s")) + "_" + strconv.Itoa(os.Getpid())
|
|
admin, err := sql.Open("pgx", dsn)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := admin.Exec("CREATE SCHEMA " + schema); err != nil {
|
|
admin.Close()
|
|
t.Fatalf("create schema: %v", err)
|
|
}
|
|
t.Cleanup(func() {
|
|
admin.Exec("DROP SCHEMA IF EXISTS " + schema + " CASCADE")
|
|
admin.Close()
|
|
})
|
|
st, err := OpenPostgres(dsn+"&search_path="+schema, filepath.Join(t.TempDir(), "artifacts"))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
t.Cleanup(func() { st.Close() })
|
|
return st
|
|
}
|
|
|
|
func isSQLite(st *Store) bool { return st.Dialect() == "sqlite" }
|
|
|
|
// 迁移验收：全量应用 → 状态 1..4 → 回滚到 1（老库形态）→ 再次前进（既有部署升级路径）。
|
|
func TestMigrationUpgradesExistingDatabase(t *testing.T) {
|
|
st := openTestStore(t)
|
|
got, err := st.MigrationStatus()
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(got) != 5 {
|
|
t.Fatalf("expected versions [1..5], got %v", got)
|
|
}
|
|
// 老库形态（v2 前）：events 全局自增、settings 单列主键。
|
|
if err := st.MigrateDown(1); err != nil {
|
|
t.Fatalf("migrate down: %v", err)
|
|
}
|
|
var seqDefault any
|
|
if err := st.db.QueryRow(`SELECT COALESCE((SELECT MAX(seq) FROM events), 0)`).Scan(&seqDefault); err != nil {
|
|
t.Fatalf("old events table: %v", err)
|
|
}
|
|
// 写一行老形态数据再前进：升级必须保留数据。
|
|
now := nowISO()
|
|
if _, err := st.db.Exec(st.q(`INSERT INTO projects(id,name,settings,created_at,updated_at) VALUES('prj_old','老项目','{}',?,?)`), now, now); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := st.db.Exec(st.q(`INSERT INTO settings(scope,"values") VALUES('user','{"a":1}')`)); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := st.db.Exec(st.q(`INSERT INTO events(conversation_id,run_id,ts,v,type,payload) VALUES('c1','r1',?,1,'run.phase','{}')`), now); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
// 重新前进（走真实升级路径）。
|
|
if err := st.migrate(); err != nil {
|
|
t.Fatalf("re-migrate: %v", err)
|
|
}
|
|
if p, err := st.GetProject("", "prj_old"); err != nil || p.Name != "老项目" {
|
|
t.Fatalf("project survived upgrade: %v %v", p, err)
|
|
}
|
|
if m, err := st.GetSettingsScope("", "user"); err != nil || len(m) != 1 {
|
|
t.Fatalf("settings survived upgrade: %v %v", m, err)
|
|
}
|
|
// 事件重编号为会话内连续 seq。
|
|
evs, err := st.EventsAfter("c1", 0, 10)
|
|
if err != nil || len(evs) != 1 || evs[0].Seq != 1 {
|
|
t.Fatalf("event renumbered to per-conversation seq: %v %v", evs, err)
|
|
}
|
|
}
|
|
|
|
// 租户隔离（负例）：跨租户读写一律 NotFound，共享知识（owner=”）双方可见。
|
|
func TestTenantIsolation(t *testing.T) {
|
|
st := openTestStore(t)
|
|
a, b := "alice", "bob"
|
|
|
|
pa, err := st.CreateProject(a, "A的项目")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := st.GetProject(b, pa.ID); err != ErrNotFound {
|
|
t.Fatalf("cross-tenant project read must be NotFound, got %v", err)
|
|
}
|
|
if err := st.RenameProject(b, pa.ID, "劫持"); err != ErrNotFound {
|
|
t.Fatalf("cross-tenant project write must be NotFound, got %v", err)
|
|
}
|
|
ca, err := st.CreateConversation(a, pa.ID, "A的会话", "")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := st.GetConversation(b, ca.ID); err != ErrNotFound {
|
|
t.Fatalf("cross-tenant conversation read must be NotFound, got %v", err)
|
|
}
|
|
if _, err := st.CreateConversation(b, pa.ID, "越权会话", ""); err != ErrNotFound {
|
|
t.Fatalf("conversation in foreign project must be NotFound, got %v", err)
|
|
}
|
|
art := core.Artifact{Owner: a, ProjectID: pa.ID, Name: "a.txt", Kind: "data"}
|
|
if err := st.CreateArtifact(&art, []byte("x")); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := st.GetArtifact(b, art.ID); err != ErrNotFound {
|
|
t.Fatalf("cross-tenant artifact read must be NotFound, got %v", err)
|
|
}
|
|
if _, err := st.ArtifactContent(b, art.ID); err != ErrNotFound {
|
|
t.Fatalf("cross-tenant artifact content must be NotFound, got %v", err)
|
|
}
|
|
// 密钥与设置隔离。
|
|
keyA, err := st.CreateAPIKey(a, "ka")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if owner, ok := st.AuthOwner(keyA); !ok || owner != a {
|
|
t.Fatalf("auth owner: %q %v", owner, ok)
|
|
}
|
|
if _, ok := st.AuthOwner("ask_nonexistent"); ok {
|
|
t.Fatal("unknown key must not authenticate")
|
|
}
|
|
if err := st.MergeSettingsScope(a, "user", map[string]any{"x": 1}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if m, err := st.GetSettingsScope(b, "user"); err != nil || len(m) != 0 {
|
|
t.Fatalf("settings must be tenant-scoped, got %v %v", m, err)
|
|
}
|
|
// 共享知识：owner='' 对双方可见。
|
|
if err := st.AddKnowledge(core.KnowledgeDoc{Owner: "", Title: "共享", Chunk: "共享知识内容"}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := st.AddKnowledge(core.KnowledgeDoc{Owner: a, Title: "A私有", Chunk: "A 私有知识内容"}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
hitsB, err := st.SearchKnowledge(b, core.RetrievalQuery{Text: "共享"})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(hitsB) == 0 {
|
|
t.Fatal("shared knowledge must be visible to bob")
|
|
}
|
|
hitsB2, err := st.SearchKnowledge(b, core.RetrievalQuery{Text: "私有"})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(hitsB2) != 0 {
|
|
t.Fatal("alice's private knowledge must not be visible to bob")
|
|
}
|
|
}
|
|
|
|
// 并发验收：同会话并发追加消息，seq 唯一且连续（先查后写竞态的约束兜底）。
|
|
func TestConcurrentAppendMessageUniqueSeq(t *testing.T) {
|
|
st := openTestStore(t)
|
|
p, _ := st.CreateProject("o", "p")
|
|
c, _ := st.CreateConversation("o", p.ID, "c", "")
|
|
const n = 24
|
|
var wg sync.WaitGroup
|
|
errs := make(chan error, n)
|
|
for i := 0; i < n; i++ {
|
|
wg.Add(1)
|
|
go func(i int) {
|
|
defer wg.Done()
|
|
m := core.Message{ConversationID: c.ID, Role: core.RoleUser, Parts: []core.Part{{Type: core.PartText, Text: "m"}}}
|
|
if err := st.AppendMessage(&m); err != nil {
|
|
errs <- err
|
|
}
|
|
}(i)
|
|
}
|
|
wg.Wait()
|
|
close(errs)
|
|
for err := range errs {
|
|
t.Fatalf("concurrent append: %v", err)
|
|
}
|
|
msgs, err := st.ListMessages(c.ID)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(msgs) != n {
|
|
t.Fatalf("expected %d messages, got %d", n, len(msgs))
|
|
}
|
|
}
|
|
|
|
// 并发验收：同 scope 并发合并设置，键值不丢失（读改写事务 + 锁）。
|
|
func TestConcurrentSettingsMergeNoLostUpdate(t *testing.T) {
|
|
st := openTestStore(t)
|
|
const writers = 8
|
|
var wg sync.WaitGroup
|
|
for i := 0; i < writers; i++ {
|
|
wg.Add(1)
|
|
go func(i int) {
|
|
defer wg.Done()
|
|
_ = st.MergeSettingsScope("o", "user", map[string]any{string(rune('a' + i)): i})
|
|
}(i)
|
|
}
|
|
wg.Wait()
|
|
m, err := st.GetSettingsScope("o", "user")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(m) != writers {
|
|
t.Fatalf("lost update: expected %d keys, got %d (%v)", writers, len(m), m)
|
|
}
|
|
}
|
|
|
|
// 批量事件：连续 seq 分配 + 游标重放。
|
|
func TestAppendEventsBatchSeq(t *testing.T) {
|
|
st := openTestStore(t)
|
|
mk := func(t2 string) core.StreamEvent {
|
|
return core.StreamEvent{ConversationID: "c1", RunID: "r1", TS: nowISO(), V: 1, Type: core.StreamEventType(t2), Payload: []byte(`{}`)}
|
|
}
|
|
batch := []core.StreamEvent{mk("message.delta"), mk("message.delta"), mk("run.completed")}
|
|
if _, err := st.AppendEvents("c1", batch); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
for i, want := range []int64{1, 2, 3} {
|
|
if batch[i].Seq != want {
|
|
t.Fatalf("seq[%d] = %d, want %d", i, batch[i].Seq, want)
|
|
}
|
|
}
|
|
// 第二批接续。
|
|
b2 := []core.StreamEvent{mk("message.delta")}
|
|
if _, err := st.AppendEvents("c1", b2); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if b2[0].Seq != 4 {
|
|
t.Fatalf("second batch seq = %d, want 4", b2[0].Seq)
|
|
}
|
|
evs, err := st.EventsAfter("c1", 2, 10)
|
|
if err != nil || len(evs) != 2 {
|
|
t.Fatalf("replay after seq=2: %v %v", evs, err)
|
|
}
|
|
if last, _ := st.LastEventSeq("c1"); last != 4 {
|
|
t.Fatalf("last seq = %d", last)
|
|
}
|
|
// 独立会话互不影响（会话内 seq 各自从 1 开始）。
|
|
b3 := []core.StreamEvent{mk("run.phase")}
|
|
if _, err := st.AppendEvents("c2", b3); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if b3[0].Seq != 1 {
|
|
t.Fatalf("per-conversation seq expected, got %d", b3[0].Seq)
|
|
}
|
|
}
|
|
|
|
// HITL 原子 resolve：并发/重复 resolve 只成功一次。
|
|
func TestHITLResolveAtomic(t *testing.T) {
|
|
st := openTestStore(t)
|
|
card := core.HITLCard{ID: "hitl_1", RunID: "r1", ConversationID: "c1", Question: "q", Kind: core.HITLKindChoice}
|
|
if err := st.SaveHITLCard("o", &card); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var wg sync.WaitGroup
|
|
ok := make(chan bool, 2)
|
|
for i := 0; i < 2; i++ {
|
|
wg.Add(1)
|
|
go func() {
|
|
defer wg.Done()
|
|
err := st.ResolveHITLCard("o", "hitl_1", core.HITLResolvedPayload{Kind: "choice", Answer: "x"})
|
|
ok <- err == nil
|
|
}()
|
|
}
|
|
wg.Wait()
|
|
succeeded := 0
|
|
for i := 0; i < 2; i++ {
|
|
if <-ok {
|
|
succeeded++
|
|
}
|
|
}
|
|
if succeeded != 1 {
|
|
t.Fatalf("double resolve must succeed exactly once, got %d", succeeded)
|
|
}
|
|
}
|
|
|
|
// 备份验收：VACUUM INTO 产出可打开的一致性副本。
|
|
func TestBackup(t *testing.T) {
|
|
st := openTestStore(t)
|
|
p, _ := st.CreateProject("o", "p")
|
|
dest := filepath.Join(t.TempDir(), "backup.db")
|
|
err := st.Backup(dest)
|
|
if !isSQLite(st) {
|
|
// PG 档：备份走 pg_dump（运维路径），接口返回指引错误。
|
|
if err == nil || !strings.Contains(err.Error(), "pg_dump") {
|
|
t.Fatalf("PG tier backup must return guidance error, got %v", err)
|
|
}
|
|
return
|
|
}
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
// 打开副本验证数据完整。
|
|
cp, err := OpenSQLite(dest)
|
|
if err != nil {
|
|
t.Fatalf("open backup: %v", err)
|
|
}
|
|
defer cp.Close()
|
|
if _, err := cp.GetProject("o", p.ID); err != nil {
|
|
t.Fatalf("backup content: %v", err)
|
|
}
|
|
}
|
|
|
|
// 知识批量导入为单事务：全部成功（半批失败回滚由事务保证）。
|
|
func TestAddKnowledgeBatch(t *testing.T) {
|
|
st := openTestStore(t)
|
|
docs := []core.KnowledgeDoc{
|
|
{Owner: "o", Title: "t1", Chunk: "alpha beta"},
|
|
{Owner: "o", Title: "t2", Chunk: "gamma delta"},
|
|
}
|
|
if err := st.AddKnowledge(docs...); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
hits, err := st.SearchKnowledge("o", core.RetrievalQuery{Text: "alpha"})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(hits) == 0 {
|
|
t.Fatal("FTS must find alpha")
|
|
}
|
|
// 中文子串走 LIKE 兜底。
|
|
if err := st.AddKnowledge(core.KnowledgeDoc{Owner: "o", Title: "中文", Chunk: "知识库检索中文兜底路径"}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
hits, err = st.SearchKnowledge("o", core.RetrievalQuery{Text: "中文兜底"})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(hits) == 0 {
|
|
t.Fatal("CJK substring fallback must hit")
|
|
}
|
|
}
|
|
|
|
// 工件原子性：创建失败不留悬空行；内容更新走临时文件。
|
|
func TestArtifactAtomicity(t *testing.T) {
|
|
st := openTestStore(t)
|
|
art := core.Artifact{Owner: "o", ProjectID: "p", Name: "a.txt", Kind: "data"}
|
|
if err := st.CreateArtifact(&art, []byte("hello")); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if c, err := st.ArtifactContent("o", art.ID); err != nil || string(c) != "hello" {
|
|
t.Fatalf("content: %q %v", c, err)
|
|
}
|
|
if err := st.UpdateArtifactContent("o", art.ID, []byte("world")); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if c, _ := st.ArtifactContent("o", art.ID); string(c) != "world" {
|
|
t.Fatal("update lost")
|
|
}
|
|
// 删除联动。
|
|
if err := st.DeleteArtifact("o", art.ID); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := st.GetArtifact("o", art.ID); err != ErrNotFound {
|
|
t.Fatalf("deleted artifact: %v", err)
|
|
}
|
|
}
|
|
|
|
// 画像租户边界：builtin 共享，用户画像隔离。
|
|
func TestProfileTenantScope(t *testing.T) {
|
|
st := openTestStore(t)
|
|
if err := st.SaveProfile(core.AgentProfile{ID: "default", Builtin: true, SystemPrompt: "s", Name: "默认"}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := st.SaveProfile(core.AgentProfile{ID: "prof_alice", Owner: "alice", SystemPrompt: "s", Name: "A"}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
la, _ := st.ListProfiles("alice")
|
|
if len(la) != 2 { // builtin + alice 自己
|
|
t.Fatalf("alice profiles = %d", len(la))
|
|
}
|
|
lb, _ := st.ListProfiles("bob")
|
|
if len(lb) != 1 { // 仅 builtin
|
|
t.Fatalf("bob profiles = %d", len(lb))
|
|
}
|
|
if _, err := st.GetProfile("bob", "prof_alice"); err != ErrNotFound {
|
|
t.Fatalf("cross-tenant profile: %v", err)
|
|
}
|
|
}
|