weichenandCursor
3997a04bdb
Use an isolated dual-pool database so idempotent finalize, cancel CAS, commit-fail SSE, and receipt-gated cleanup are proven without a process mutex. Co-authored-by: Cursor <cursoragent@cursor.com>
364 lines
13 KiB
Go
364 lines
13 KiB
Go
package taskrun_test
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"net/url"
|
|
"os"
|
|
"path/filepath"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"chat-agent-system/internal/app/uow"
|
|
"chat-agent-system/internal/artifactvault"
|
|
"chat-agent-system/internal/bizdev"
|
|
"chat-agent-system/internal/ctxmgr"
|
|
"chat-agent-system/internal/db"
|
|
"chat-agent-system/internal/docgov"
|
|
"chat-agent-system/internal/stream"
|
|
"chat-agent-system/internal/taskrun"
|
|
"chat-agent-system/internal/workspace"
|
|
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
)
|
|
|
|
func init() {
|
|
if _, err := os.Stat("deploy/postgres/goose"); err == nil {
|
|
return
|
|
}
|
|
for _, cand := range []string{"../..", "../../.."} {
|
|
if fi, err := os.Stat(cand + "/deploy/postgres/goose"); err == nil && fi.IsDir() {
|
|
_ = os.Chdir(cand)
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
var (
|
|
isoCleanup func()
|
|
)
|
|
|
|
func TestMain(m *testing.M) {
|
|
code := m.Run()
|
|
if isoCleanup != nil {
|
|
isoCleanup()
|
|
}
|
|
os.Exit(code)
|
|
}
|
|
|
|
func rewriteDatabaseName(dsn, name string) string {
|
|
if u, err := url.Parse(dsn); err == nil && u.Scheme != "" {
|
|
u.Path = "/" + name
|
|
return u.String()
|
|
}
|
|
return dsn
|
|
}
|
|
|
|
var isoOnce sync.Once
|
|
var isoPoolA, isoPoolB *pgxpool.Pool
|
|
var isoErr error
|
|
|
|
func openIsolatedM3(t *testing.T) (*pgxpool.Pool, *pgxpool.Pool) {
|
|
t.Helper()
|
|
if testing.Short() {
|
|
t.Skip("U8 PG races are not part of -short")
|
|
}
|
|
dsn := db.SkipIfNoTestDatabaseURL(t)
|
|
db.RefuseSharedEnvDSN(t, dsn)
|
|
|
|
isoOnce.Do(func() {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
|
|
defer cancel()
|
|
adminDSN := rewriteDatabaseName(dsn, "postgres")
|
|
admin, err := pgxpool.New(ctx, adminDSN)
|
|
if err != nil {
|
|
isoErr = err
|
|
return
|
|
}
|
|
const isolated = "agent_rel_m3_u8"
|
|
_, _ = admin.Exec(ctx, `SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname=$1 AND pid <> pg_backend_pid()`, isolated)
|
|
_, _ = admin.Exec(ctx, `DROP DATABASE IF EXISTS `+isolated)
|
|
if _, err := admin.Exec(ctx, `CREATE DATABASE `+isolated); err != nil {
|
|
admin.Close()
|
|
isoErr = err
|
|
return
|
|
}
|
|
isoDSN := rewriteDatabaseName(dsn, isolated)
|
|
poolA, err := pgxpool.New(ctx, isoDSN)
|
|
if err != nil {
|
|
admin.Close()
|
|
isoErr = err
|
|
return
|
|
}
|
|
if err := db.RunMigrations(ctx, poolA, db.ResolveMigrationsDir()); err != nil {
|
|
poolA.Close()
|
|
admin.Close()
|
|
isoErr = err
|
|
return
|
|
}
|
|
poolB, err := pgxpool.New(ctx, isoDSN)
|
|
if err != nil {
|
|
poolA.Close()
|
|
admin.Close()
|
|
isoErr = err
|
|
return
|
|
}
|
|
isoPoolA, isoPoolB = poolA, poolB
|
|
isoCleanup = func() {
|
|
poolA.Close()
|
|
poolB.Close()
|
|
_, _ = admin.Exec(context.Background(), `SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname=$1 AND pid <> pg_backend_pid()`, isolated)
|
|
_, _ = admin.Exec(context.Background(), `DROP DATABASE IF EXISTS `+isolated)
|
|
admin.Close()
|
|
}
|
|
})
|
|
if isoErr != nil {
|
|
t.Fatalf("isolated db: %v", isoErr)
|
|
}
|
|
return isoPoolA, isoPoolB
|
|
}
|
|
|
|
func seedUserAndProject(t *testing.T, pool *pgxpool.Pool) (userID, projectID int64) {
|
|
t.Helper()
|
|
ctx := context.Background()
|
|
err := pool.QueryRow(ctx, `SELECT id FROM users ORDER BY id LIMIT 1`).Scan(&userID)
|
|
if err != nil {
|
|
t.Fatalf("seed user: %v", err)
|
|
}
|
|
err = pool.QueryRow(ctx, `SELECT id FROM projects ORDER BY id LIMIT 1`).Scan(&projectID)
|
|
if err != nil {
|
|
t.Fatalf("seed project: %v", err)
|
|
}
|
|
return userID, projectID
|
|
}
|
|
|
|
func insertRunning(t *testing.T, pool *pgxpool.Pool, rec *taskrun.RunRecord) {
|
|
t.Helper()
|
|
_, err := pool.Exec(context.Background(), `
|
|
INSERT INTO task_runs (
|
|
run_id, session_id, project_id, user_id, client_turn_id, task_key, status,
|
|
workspace_id, assistant_message_id, created_at, updated_at
|
|
) VALUES ($1,$2,$3,$4,$5,$6,'running',$7,$8,now(),now())
|
|
`, rec.RunID, rec.SessionID, rec.ProjectID, rec.UserID, rec.ClientTurnID, rec.TaskKey, rec.WorkspaceID, rec.AssistantMsgID)
|
|
if err != nil {
|
|
t.Fatalf("insert run: %v", err)
|
|
}
|
|
}
|
|
|
|
func TestPG_DuplicateFinalizeIdempotent(t *testing.T) {
|
|
poolA, poolB := openIsolatedM3(t)
|
|
userID, projectID := seedUserAndProject(t, poolA)
|
|
ctxMgr := ctxmgr.NewContextManager(poolA, nil, ctxmgr.TimeDecayConfig{})
|
|
sessionID := "sess_u8_dup"
|
|
asstID, err := ctxMgr.RecordUserMessageWithID(context.Background(), sessionID, userID, "assistant", "任务执行中…", "TRIGGER_TASK", "", map[string]any{"project_id": projectID})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
rec := &taskrun.RunRecord{
|
|
RunID: "run_u8_dup", SessionID: sessionID, ProjectID: projectID, UserID: userID,
|
|
TaskKey: "intake", AssistantMsgID: asstID, WorkspaceID: "ws_u8_dup",
|
|
}
|
|
insertRunning(t, poolA, rec)
|
|
body := []byte("# committed intake body\n")
|
|
res := &bizdev.TaskRunResult{
|
|
Success: true, TaskKey: "intake", ArtifactTitle: "用户初始需求收集表", MarkdownBody: string(body),
|
|
DocRecord: &docgov.DocumentRecord{DocID: "doc_u8_dup", ContentSHA256: artifactvault.SHA256Hex(body), Version: 1, ProjectID: projectID},
|
|
}
|
|
cmdsA := &taskrun.Commands{UoW: uow.NewPG(poolA), Ctx: ctxMgr, Fav: &artifactvault.Ingest{Env: "dev"}, Env: "dev", StorageIdentity: "env:dev"}
|
|
cmdsB := &taskrun.Commands{UoW: uow.NewPG(poolB), Ctx: ctxmgr.NewContextManager(poolB, nil, ctxmgr.TimeDecayConfig{}), Fav: &artifactvault.Ingest{Env: "dev"}, Env: "dev", StorageIdentity: "env:dev"}
|
|
st1, err := cmdsA.Finalize(context.Background(), rec, res)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
st2, err := cmdsB.Finalize(context.Background(), rec, res)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if st1 != "completed" || st2 != "completed" {
|
|
t.Fatalf("status %s %s", st1, st2)
|
|
}
|
|
var msgs, mans, envs, stgs int
|
|
_ = poolA.QueryRow(context.Background(), `SELECT COUNT(*) FROM chat_messages WHERE session_id=$1 AND role='assistant'`, sessionID).Scan(&msgs)
|
|
_ = poolA.QueryRow(context.Background(), `SELECT COUNT(*) FROM task_run_manifests WHERE run_id=$1`, rec.RunID).Scan(&mans)
|
|
_ = poolA.QueryRow(context.Background(), `SELECT COUNT(*) FROM event_envelopes WHERE aggregate_id=$1`, rec.RunID).Scan(&envs)
|
|
_ = poolA.QueryRow(context.Background(), `SELECT COUNT(*) FROM artifact_payload_staging WHERE source_event_id=$1`, "taskrun:"+rec.RunID+":completed").Scan(&stgs)
|
|
if msgs != 1 || mans != 1 || envs != 1 || stgs != 1 {
|
|
t.Fatalf("counts messages=%d manifests=%d envelopes=%d staging=%d", msgs, mans, envs, stgs)
|
|
}
|
|
var sha1, sha2 string
|
|
_ = poolA.QueryRow(context.Background(), `SELECT sha256 FROM task_run_manifests WHERE run_id=$1`, rec.RunID).Scan(&sha1)
|
|
_ = poolB.QueryRow(context.Background(), `SELECT sha256 FROM task_run_manifests WHERE run_id=$1`, rec.RunID).Scan(&sha2)
|
|
if sha1 == "" || sha1 != sha2 {
|
|
t.Fatalf("manifest sha %q %q", sha1, sha2)
|
|
}
|
|
}
|
|
|
|
func TestPG_FinalizeCancelRaceOneTerminal(t *testing.T) {
|
|
poolA, poolB := openIsolatedM3(t)
|
|
userID, projectID := seedUserAndProject(t, poolA)
|
|
ctxMgr := ctxmgr.NewContextManager(poolA, nil, ctxmgr.TimeDecayConfig{})
|
|
sessionID := "sess_u8_race"
|
|
asstID, err := ctxMgr.RecordUserMessageWithID(context.Background(), sessionID, userID, "assistant", "任务执行中…", "TRIGGER_TASK", "", map[string]any{"project_id": projectID})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
rec := &taskrun.RunRecord{
|
|
RunID: "run_u8_race", SessionID: sessionID, ProjectID: projectID, UserID: userID,
|
|
TaskKey: "intake", AssistantMsgID: asstID, WorkspaceID: "ws_u8_race",
|
|
}
|
|
insertRunning(t, poolA, rec)
|
|
body := []byte("# race body\n")
|
|
res := &bizdev.TaskRunResult{
|
|
Success: true, TaskKey: "intake", MarkdownBody: string(body),
|
|
DocRecord: &docgov.DocumentRecord{DocID: "doc_u8_race", ContentSHA256: artifactvault.SHA256Hex(body), Version: 1},
|
|
}
|
|
cmdsA := &taskrun.Commands{UoW: uow.NewPG(poolA), Ctx: ctxMgr, Fav: &artifactvault.Ingest{Env: "dev"}, Env: "dev", StorageIdentity: "env:dev"}
|
|
cmdsB := &taskrun.Commands{UoW: uow.NewPG(poolB), Ctx: ctxmgr.NewContextManager(poolB, nil, ctxmgr.TimeDecayConfig{}), Env: "dev", StorageIdentity: "env:dev"}
|
|
var wg sync.WaitGroup
|
|
wg.Add(2)
|
|
go func() {
|
|
defer wg.Done()
|
|
_, _ = cmdsA.Finalize(context.Background(), rec, res)
|
|
}()
|
|
go func() {
|
|
defer wg.Done()
|
|
_, _ = cmdsB.Cancel(context.Background(), rec)
|
|
}()
|
|
wg.Wait()
|
|
var status string
|
|
if err := poolA.QueryRow(context.Background(), `SELECT status FROM task_runs WHERE run_id=$1`, rec.RunID).Scan(&status); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if status != "completed" && status != "cancelled" {
|
|
t.Fatalf("illegal terminal %q", status)
|
|
}
|
|
var n int
|
|
_ = poolA.QueryRow(context.Background(), `SELECT COUNT(*) FROM chat_messages WHERE session_id=$1 AND role='assistant'`, sessionID).Scan(&n)
|
|
if n != 1 {
|
|
t.Fatalf("assistant turns=%d", n)
|
|
}
|
|
}
|
|
|
|
func TestPG_SSEDoneNotEmittedWhenCommitFails(t *testing.T) {
|
|
poolA, _ := openIsolatedM3(t)
|
|
userID, projectID := seedUserAndProject(t, poolA)
|
|
ctxMgr := ctxmgr.NewContextManager(poolA, nil, ctxmgr.TimeDecayConfig{})
|
|
sessionID := "sess_u8_sse"
|
|
asstID, err := ctxMgr.RecordUserMessageWithID(context.Background(), sessionID, userID, "assistant", "任务执行中…", "TRIGGER_TASK", "", map[string]any{"project_id": projectID})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
rec := &taskrun.RunRecord{
|
|
RunID: "run_u8_sse", SessionID: sessionID, ProjectID: projectID, UserID: userID,
|
|
TaskKey: "intake", AssistantMsgID: asstID, Status: "pending",
|
|
}
|
|
unit := uow.NewPG(poolA)
|
|
unit.BeforeCommit = func(uow.TxScope) error { return errors.New("inject commit fail") }
|
|
cmds := &taskrun.Commands{UoW: unit, Ctx: ctxMgr, Env: "dev", StorageIdentity: "env:dev"}
|
|
mgr := taskrun.NewManager(nil, ctxMgr)
|
|
mgr.SetStore(taskrun.NewPGStore(poolA))
|
|
mgr.SetCommands(cmds)
|
|
mgr.SetExecutor(func(context.Context, string, bizdev.TaskRunRequest, *stream.Streamer) (*bizdev.TaskRunResult, error) {
|
|
return &bizdev.TaskRunResult{Success: true, MarkdownBody: "# sse\n"}, nil
|
|
}, nil, nil)
|
|
if _, err := mgr.CreateAndStart(context.Background(), rec.RunID, rec, bizdev.TaskRunRequest{TaskKey: "intake"}, bizdev.StageRunRequest{}, ""); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
deadline := time.Now().Add(8 * time.Second)
|
|
for time.Now().Before(deadline) {
|
|
got, err := mgr.GetRun(rec.RunID)
|
|
if err == nil && got.Status == "finalizing_failed" {
|
|
evs, listErr := taskrun.NewPGStore(poolA).ListEvents(context.Background(), rec.RunID, 0)
|
|
if listErr != nil {
|
|
t.Fatal(listErr)
|
|
}
|
|
for _, ev := range evs {
|
|
if ev.Event.Type == stream.EventDone {
|
|
t.Fatal("done emitted on failed commit")
|
|
}
|
|
}
|
|
msg, _ := ctxMgr.GetMessageByID(context.Background(), asstID)
|
|
if msg != nil && msg.Content != "任务执行中…" {
|
|
t.Fatalf("placeholder mutated after rollback: %q", msg.Content)
|
|
}
|
|
return
|
|
}
|
|
time.Sleep(30 * time.Millisecond)
|
|
}
|
|
t.Fatal("expected finalizing_failed")
|
|
}
|
|
|
|
func TestPG_CleanupWithoutReceiptFailsAndStagingSurvives(t *testing.T) {
|
|
poolA, _ := openIsolatedM3(t)
|
|
userID, projectID := seedUserAndProject(t, poolA)
|
|
ctxMgr := ctxmgr.NewContextManager(poolA, nil, ctxmgr.TimeDecayConfig{})
|
|
sessionID := "sess_u8_clean"
|
|
asstID, err := ctxMgr.RecordUserMessageWithID(context.Background(), sessionID, userID, "assistant", "任务执行中…", "TRIGGER_TASK", "", map[string]any{"project_id": projectID})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
dataRoot := t.TempDir()
|
|
cap := workspace.NewCapability(dataRoot, nil)
|
|
ref, err := cap.Issue(workspace.WorkspaceRef{ID: "ws_u8_clean", ProjectScope: projectID, RunID: "run_u8_clean"})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
body := []byte("# vaulted body for cleanup\n")
|
|
if err := cap.WriteArtifact(context.Background(), ref, workspace.DeliverableKey, body); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
rec := &taskrun.RunRecord{
|
|
RunID: "run_u8_clean", SessionID: sessionID, ProjectID: projectID, UserID: userID,
|
|
TaskKey: "intake", AssistantMsgID: asstID, WorkspaceID: ref.ID,
|
|
}
|
|
insertRunning(t, poolA, rec)
|
|
res := &bizdev.TaskRunResult{
|
|
Success: true, TaskKey: "intake", MarkdownBody: string(body), WorkDir: ref.ID,
|
|
DocRecord: &docgov.DocumentRecord{DocID: "doc_u8_clean", ContentSHA256: artifactvault.SHA256Hex(body), Version: 2},
|
|
}
|
|
cmds := &taskrun.Commands{UoW: uow.NewPG(poolA), Ctx: ctxMgr, Fav: &artifactvault.Ingest{Env: "dev"}, Env: "dev", StorageIdentity: "env:dev"}
|
|
if _, err := cmds.Finalize(context.Background(), rec, res); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var manSHA, storageID string
|
|
var members []byte
|
|
if err := poolA.QueryRow(context.Background(), `SELECT sha256, storage_identity, members FROM task_run_manifests WHERE run_id=$1`, rec.RunID).
|
|
Scan(&manSHA, &storageID, &members); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var parsed []taskrun.ManifestMember
|
|
if err := json.Unmarshal(members, &parsed); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
frozen := workspace.FrozenManifest{SHA256: manSHA, StorageIdentity: storageID, WorkspaceID: ref.ID}
|
|
for _, m := range parsed {
|
|
frozen.Members = append(frozen.Members, workspace.FrozenMember{LogicalID: m.LogicalID, Digest: m.Digest, Kind: m.Kind})
|
|
}
|
|
if err := cap.Cleanup(context.Background(), ref, frozen, nil); err == nil {
|
|
t.Fatal("cleanup without receipts must fail")
|
|
}
|
|
if _, err := os.Stat(filepath.Join(dataRoot, "workspaces", ref.ID, workspace.DeliverableKey)); err != nil {
|
|
t.Fatalf("workspace removed without receipts: %v", err)
|
|
}
|
|
var receipts []workspace.RetentionReceipt
|
|
for _, m := range parsed {
|
|
receipts = append(receipts, workspace.RetentionReceipt{
|
|
LogicalID: m.LogicalID, Digest: m.Digest, StorageIdentity: storageID, Verified: true,
|
|
})
|
|
}
|
|
if err := cap.Cleanup(context.Background(), ref, frozen, receipts); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
eventID := "taskrun:" + rec.RunID + ":completed"
|
|
reg := artifactvault.NewPGRegistry(poolA)
|
|
st, err := reg.LoadStaging(context.Background(), eventID)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if artifactvault.SHA256Hex(st.Payload) != artifactvault.SHA256Hex(body) {
|
|
t.Fatalf("staging sha drifted")
|
|
}
|
|
}
|