Agent-System/internal/server/task_run_persistence_pg_test.go

package server_test

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"chat-agent-system/internal/bizdev"
	"chat-agent-system/internal/stream"
	"chat-agent-system/internal/taskrun"
)

func TestTaskRun_PersistenceAndReplayAfterRestart_PG(t *testing.T) {
	pool := openPGPoolWithMigrations(t)
	if pool == nil {
		return
	}
	defer pool.Close()
	if !taskrun.TablesReady(context.Background(), pool) {
		t.Skip("task_runs table not ready")
	}

	handler, srv := setupPGGovernanceServer(t, pool)
	token := loginTestUser(t, handler, "zhanggong", "123456")

	runID := "run_pg_persist_" + time.Now().Format("150405")
	store := taskrun.NewPGStore(pool)
	srv.TaskRuns().SetStore(store)

	srv.TaskRuns().SetExecutor(
		func(_ context.Context, _ string, _ bizdev.TaskRunRequest, st *stream.Streamer) (*bizdev.TaskRunResult, error) {
			if st != nil {
				st.EmitPhase("test", "running", "step1", 1, 1, 50)
			}
			return &bizdev.TaskRunResult{Success: true, WorkDir: t.TempDir(), TaskKey: "intake"}, nil
		},
		nil,
		nil,
	)
	srv.TaskRuns().SetFinalizer(func(ctx context.Context, record *taskrun.RunRecord, result *bizdev.TaskRunResult) (string, error) {
		return srv.FinalizeTaskRunProjection(ctx, record, result)
	})

	record := &taskrun.RunRecord{
		RunID: runID, SessionID: "sess_pg_persist", ProjectID: 0,
		UserID: fetchAuthUserID(t, handler, token), TaskKey: "intake", Status: "running",
	}
	_, err := srv.TaskRuns().CreateAndStart(context.Background(), runID, record, bizdev.TaskRunRequest{TaskKey: "intake"}, bizdev.StageRunRequest{}, "")
	if err != nil {
		t.Fatalf("CreateAndStart: %v", err)
	}

	deadline := time.Now().Add(10 * time.Second)
	for time.Now().Before(deadline) {
		rec, getErr := store.GetRun(context.Background(), runID)
		if getErr == nil && rec.Status == "completed" {
			break
		}
		time.Sleep(50 * time.Millisecond)
	}

	persisted, err := store.GetRun(context.Background(), runID)
	if err != nil {
		t.Fatalf("GetRun from store: %v", err)
	}
	if persisted.Status != "completed" {
		t.Fatalf("expected completed, got %s", persisted.Status)
	}

	events, err := store.ListEvents(context.Background(), runID, 0)
	if err != nil {
		t.Fatalf("ListEvents: %v", err)
	}
	if len(events) == 0 {
		t.Fatal("expected persisted events")
	}
	hasDone := false
	for _, ev := range events {
		if ev.Event.Type == stream.EventDone {
			hasDone = true
		}
	}
	if !hasDone {
		t.Fatal("expected done event in persisted events")
	}

	// 模拟进程重启：新 Manager 仅挂载 Store
	restarted := taskrun.NewManager(srv.ProjectWorkflowService(), srv.BizdevService().ContextManager())
	restarted.SetStore(store)

	reloaded, err := restarted.GetRun(runID)
	if err != nil {
		t.Fatalf("GetRun after restart: %v", err)
	}
	if reloaded.Status != "completed" {
		t.Fatalf("reloaded status=%s", reloaded.Status)
	}

	ch, replay, err := restarted.Subscribe(runID, 0)
	if err != nil {
		t.Fatalf("Subscribe after restart: %v", err)
	}
	if len(replay) == 0 {
		t.Fatal("expected replay events from DB")
	}
	_, ok := <-ch
	if ok {
		t.Fatal("replay-only channel should be closed")
	}

	// HTTP GET 鉴权 + 查询
	req := httptest.NewRequest(http.MethodGet, "/api/v1/tasks/runs/"+runID, nil)
	withBearerAuth(req, token)
	rec := httptest.NewRecorder()
	handler.ServeHTTP(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("GET task run: %d %s", rec.Code, rec.Body.String())
	}
	var body taskrun.RunRecord
	if err := json.Unmarshal(rec.Body.Bytes(), &body); err != nil {
		t.Fatalf("decode: %v", err)
	}
	if body.RunID != runID {
		t.Fatalf("run_id=%s", body.RunID)
	}
}

func TestReadiness_StrictWhenPGConfigured(t *testing.T) {
	pool := openPGPoolWithMigrations(t)
	if pool == nil {
		return
	}
	defer pool.Close()

	handler, _ := setupPGGovernanceServer(t, pool)
	req := httptest.NewRequest(http.MethodGet, "/api/v1/ready", nil)
	rec := httptest.NewRecorder()
	handler.ServeHTTP(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("ready expected 200, got %d %s", rec.Code, rec.Body.String())
	}
	var resp map[string]any
	_ = json.Unmarshal(rec.Body.Bytes(), &resp)
	if resp["task_store_ready"] != true {
		t.Fatalf("task_store_ready=%v", resp["task_store_ready"])
	}
}