Agent-System/internal/server/task_run_resume_pg_test.go

package server_test

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"testing"
	"time"

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

func pgResumeConfig(t *testing.T) *config.Config {
	t.Helper()
	cfg := configFromTestDatabaseURL(t)
	cfg.DataDir = t.TempDir()
	cfg.AgentBackend = "stub"
	cfg.AllowInMemory = false
	return cfg
}

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

	cfg := pgResumeConfig(t)
	svc := bizdev.MustNewService(t, cfg)
	srv := server.NewServer(cfg, svc)
	handler := srv.Routes()
	token := loginTestUser(t, handler, "zhanggong", "123456")
	userID := fetchAuthUserID(t, handler, token)

	workDir := filepath.Join(cfg.DataDir, "workspaces", "bizdev", "resume_test")
	if err := os.MkdirAll(workDir, 0o755); err != nil {
		t.Fatal(err)
	}
	_ = os.WriteFile(filepath.Join(workDir, "output.md"), []byte("# resumed\n"), 0o644)

	store := taskrun.NewPGStore(pool)
	projectID := seedProjectForRun(t, pool, "run-resume", userID)
	runID := "run_resume_" + time.Now().Format("150405.000000")
	_, err := pool.Exec(context.Background(), `
		INSERT INTO task_runs (
			run_id, session_id, project_id, user_id, task_key, status, workspace_id, error_message, created_at, updated_at
		) VALUES ($1, 'sess_resume', $2, $3, 'intake', 'interrupted', $4, 'worker lease expired', now(), now())
	`, runID, projectID, userID, workDir)
	if err != nil {
		t.Fatalf("seed interrupted run: %v", err)
	}
	t.Cleanup(func() {
		_, _ = pool.Exec(context.Background(), `DELETE FROM task_runs WHERE run_id = $1`, runID)
	})

	srv.TaskRuns().SetExecutor(
		func(_ context.Context, _ string, _ bizdev.TaskRunRequest, _ *stream.Streamer) (*bizdev.TaskRunResult, error) {
			return &bizdev.TaskRunResult{
				Success: true, WorkDir: workDir, TaskKey: "intake", ArtifactTitle: "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)
	})
	srv.TaskRuns().SetStore(store)

	body, _ := json.Marshal(map[string]bool{"confirm": true})
	req := httptest.NewRequest(http.MethodPost, "/api/v1/tasks/runs/"+runID+"/resume", bytes.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	withBearerAuth(req, token)
	rec := httptest.NewRecorder()
	handler.ServeHTTP(rec, req)
	if rec.Code != http.StatusAccepted {
		t.Fatalf("resume expected 202, got %d %s", rec.Code, rec.Body.String())
	}

	deadline := time.Now().Add(8 * time.Second)
	for time.Now().Before(deadline) {
		got, getErr := store.GetRun(context.Background(), runID)
		if getErr == nil && got.Status == "completed" {
			if got.Attempt < 2 {
				t.Fatalf("expected attempt>=2 after resume, got %d", got.Attempt)
			}
			return
		}
		time.Sleep(50 * time.Millisecond)
	}
	t.Fatal("resume did not complete")
}

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

	cfg := pgResumeConfig(t)
	srv := server.NewServer(cfg, bizdev.MustNewService(t, cfg))
	handler := srv.Routes()
	token := loginTestUser(t, handler, "zhanggong", "123456")
	userID := fetchAuthUserID(t, handler, token)
	workDir := filepath.Join(cfg.DataDir, "ws")
	_ = os.MkdirAll(workDir, 0o755)

	runID := "run_no_confirm_" + time.Now().Format("150405.000000")
	_, _ = pool.Exec(context.Background(), `
		INSERT INTO task_runs (run_id, session_id, user_id, task_key, status, workspace_id, created_at, updated_at)
		VALUES ($1, 's', $2, 'intake', 'interrupted', $3, now(), now())
	`, runID, userID, workDir)
	t.Cleanup(func() {
		_, _ = pool.Exec(context.Background(), `DELETE FROM task_runs WHERE run_id = $1`, runID)
	})
	_ = srv

	req := httptest.NewRequest(http.MethodPost, "/api/v1/tasks/runs/"+runID+"/resume", bytes.NewReader([]byte(`{"confirm":false}`)))
	req.Header.Set("Content-Type", "application/json")
	withBearerAuth(req, token)
	rec := httptest.NewRecorder()
	handler.ServeHTTP(rec, req)
	if rec.Code != http.StatusBadRequest {
		t.Fatalf("expected 400 without confirm, got %d", rec.Code)
	}
}

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

	cfg := pgResumeConfig(t)
	handler := server.NewServer(cfg, bizdev.MustNewService(t, cfg)).Routes()
	token := loginTestUser(t, handler, "zhanggong", "123456")
	userID := fetchAuthUserID(t, handler, token)
	workDir := filepath.Join(cfg.DataDir, "workspaces", "bizdev", "resumable_flag")
	_ = os.MkdirAll(workDir, 0o755)

	runID := "run_flag_" + time.Now().Format("150405.000000")
	_, _ = pool.Exec(context.Background(), `
		INSERT INTO task_runs (run_id, session_id, user_id, task_key, status, workspace_id, created_at, updated_at)
		VALUES ($1, 's', $2, 'intake', 'interrupted', $3, now(), now())
	`, runID, userID, workDir)
	t.Cleanup(func() {
		_, _ = pool.Exec(context.Background(), `DELETE FROM task_runs WHERE run_id = $1`, runID)
	})

	req := httptest.NewRequest(http.MethodGet, fmt.Sprintf("/api/v1/tasks/runs/%s", runID), nil)
	withBearerAuth(req, token)
	rec := httptest.NewRecorder()
	handler.ServeHTTP(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("GET run: %d", rec.Code)
	}
	var body map[string]any
	_ = json.Unmarshal(rec.Body.Bytes(), &body)
	if body["resumable"] != true {
		t.Fatalf("expected resumable=true, got %v", body["resumable"])
	}
}