ARTEX-chixuweihu/server/worker_pools_test.go

package server

import (
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"net/http/httptest"
	"strconv"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"github.com/Autumn-27/artex/agent"
	"github.com/Autumn-27/artex/db"
	"github.com/Autumn-27/norma/llm"
)

func TestExpandWorkerSlotsLegacy(t *testing.T) {
	for _, tc := range []struct {
		name  string
		count int
		want  int
	}{
		{"configured", 2, 2},
		{"zero defaults", 0, 3},
		{"negative defaults", -1, 3},
	} {
		t.Run(tc.name, func(t *testing.T) {
			slots, err := ExpandWorkerSlots(nil, tc.count)
			if err != nil || len(slots) != tc.want {
				t.Fatalf("slots=%v err=%v, want %d", slots, err, tc.want)
			}
			for i, slot := range slots {
				if slot.Name != "work#"+strconv.Itoa(i+1) || slot.ProfileID != 0 {
					t.Fatalf("slot[%d]=%+v", i, slot)
				}
			}
		})
	}
}

func TestExpandWorkerSlotsPoolOrder(t *testing.T) {
	slots, err := ExpandWorkerSlots([]WorkerPool{{ProfileID: 7, Workers: 2}, {ProfileID: 12, Workers: 3}}, 99)
	if err != nil || len(slots) != 5 {
		t.Fatalf("slots=%v err=%v", slots, err)
	}
	for i, slot := range slots {
		want := int64(7)
		if i >= 2 {
			want = 12
		}
		if slot.Name != "work#"+strconv.Itoa(i+1) || slot.ProfileID != want {
			t.Fatalf("slot[%d]=%+v want profile %d", i, slot, want)
		}
	}
}

func TestValidateWorkerPools(t *testing.T) {
	for _, tc := range []struct {
		name  string
		pools []WorkerPool
		valid bool
	}{
		{"empty", nil, true},
		{"bad profile", []WorkerPool{{ProfileID: 0, Workers: 1}}, false},
		{"bad workers", []WorkerPool{{ProfileID: 1, Workers: 0}}, false},
		{"duplicate", []WorkerPool{{ProfileID: 1, Workers: 1}, {ProfileID: 1, Workers: 1}}, false},
		{"too many", []WorkerPool{{ProfileID: 1, Workers: 51}}, false},
		{"valid", []WorkerPool{{ProfileID: 1, Workers: 1}}, true},
	} {
		t.Run(tc.name, func(t *testing.T) {
			if got := ValidateWorkerPools(tc.pools) == nil; got != tc.valid {
				t.Fatalf("valid=%v want %v", got, tc.valid)
			}
		})
	}
}

func TestWorkerPoolSettingsRoundTrip(t *testing.T) {
	m, err := NewManager(t.TempDir(), "")
	if err != nil {
		t.Skipf("postgres unavailable: %v", err)
	}
	defer m.Close()
	pID, err := m.pg.SaveProfile(&db.LLMProfile{Name: "pool-test", Format: "openai", Model: "pool-test", APIKey: "pool-test-key"})
	if err != nil {
		t.Fatal(err)
	}
	defer m.pg.DeleteProfile(pID)
	oldPools, _, _ := m.pg.GetSetting(settingWorkerPools)
	oldWorkers, workersSet, _ := m.pg.GetSetting(settingWorkers)
	defer func() {
		if workersSet {
			_ = m.pg.SetSetting(settingWorkers, oldWorkers)
		} else {
			_, _ = m.pg.Exec("DELETE FROM settings WHERE key=$1", settingWorkers)
		}
		if oldPools != "" {
			_ = m.pg.SetSetting(settingWorkerPools, oldPools)
		} else {
			_, _ = m.pg.Exec("DELETE FROM settings WHERE key=$1", settingWorkerPools)
		}
	}()
	if err := m.SetWorkerPools([]WorkerPool{{ProfileID: 9223372036854775807, Workers: 1}}); err == nil {
		t.Fatal("nonexistent profile accepted")
	}
	if err := m.SetWorkerPools([]WorkerPool{{ProfileID: pID, Workers: 2}}); err != nil {
		t.Fatal(err)
	}
	got := m.WorkerPools()
	if len(got) != 1 || got[0].ProfileID != pID || got[0].Workers != 2 || m.Workers() != 2 {
		t.Fatalf("round trip pools=%+v workers=%d", got, m.Workers())
	}
	if err := m.SetWorkerPools(nil); err != nil {
		t.Fatal(err)
	}
	if got := m.WorkerPools(); got == nil || len(got) != 0 {
		t.Fatalf("cleared pools=%+v", got)
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	s := New(ctx, m, t.TempDir(), t.TempDir(), t.TempDir())
	if _, err := m.pg.Exec("DELETE FROM settings WHERE key=$1", settingWorkerPools); err != nil {
		t.Fatal(err)
	}
	res := httptest.NewRecorder()
	s.getSettings(res, httptest.NewRequest(http.MethodGet, "/api/settings", nil))
	var unset map[string]any
	if err := json.Unmarshal(res.Body.Bytes(), &unset); err != nil {
		t.Fatal(err)
	}
	if pools, ok := unset["worker_pools"].([]any); !ok || pools == nil || len(pools) != 0 {
		t.Fatalf("GET unset worker_pools=%#v, want empty array", unset["worker_pools"])
	}
	put := httptest.NewRequest(http.MethodPut, "/api/settings", strings.NewReader(`{"workers":2,"worker_pools":[{"profile_id":"`+strconv.FormatInt(pID, 10)+`","workers":2}]}`))
	putResult := httptest.NewRecorder()
	s.putSettings(putResult, put)
	if putResult.Code != http.StatusOK {
		t.Fatalf("PUT worker pools status=%d body=%s", putResult.Code, putResult.Body.String())
	}
	if m.Workers() != 2 || len(m.WorkerPools()) != 1 {
		t.Fatalf("PUT worker pools did not sync settings: workers=%d pools=%+v", m.Workers(), m.WorkerPools())
	}
	workersOnly := httptest.NewRequest(http.MethodPut, "/api/settings", strings.NewReader(`{"workers":9}`))
	workersOnlyResult := httptest.NewRecorder()
	s.putSettings(workersOnlyResult, workersOnly)
	if workersOnlyResult.Code != http.StatusOK || m.Workers() != 2 {
		t.Fatalf("workers-only PUT changed pool total: status=%d workers=%d", workersOnlyResult.Code, m.Workers())
	}
	req := httptest.NewRequest(http.MethodGet, "/api/settings", nil)
	res = httptest.NewRecorder()
	s.getSettings(res, req)
	if res.Code != http.StatusOK {
		t.Fatalf("GET settings status=%d", res.Code)
	}
	var payload map[string]any
	if err := json.Unmarshal(res.Body.Bytes(), &payload); err != nil {
		t.Fatal(err)
	}
	pools, ok := payload["worker_pools"].([]any)
	if !ok || len(pools) != 1 || pools[0].(map[string]any)["profile_id"] != strconv.FormatInt(pID, 10) {
		t.Fatalf("GET worker_pools=%#v, want string profile id", payload["worker_pools"])
	}
	clear := httptest.NewRequest(http.MethodPut, "/api/settings", strings.NewReader(`{"workers":4,"worker_pools":[]}`))
	clearResult := httptest.NewRecorder()
	s.putSettings(clearResult, clear)
	if clearResult.Code != http.StatusOK || m.Workers() != 4 || len(m.WorkerPools()) != 0 {
		t.Fatalf("PUT legacy clear status=%d body=%s workers=%d pools=%+v", clearResult.Code, clearResult.Body.String(), m.Workers(), m.WorkerPools())
	}
	res = httptest.NewRecorder()
	s.getSettings(res, httptest.NewRequest(http.MethodGet, "/api/settings", nil))
	var cleared map[string]any
	if err := json.Unmarshal(res.Body.Bytes(), &cleared); err != nil {
		t.Fatal(err)
	}
	if pools, ok := cleared["worker_pools"].([]any); !ok || pools == nil || len(pools) != 0 {
		t.Fatalf("GET cleared worker_pools=%#v, want empty array", cleared["worker_pools"])
	}
}

func TestPinnedRuntimeMissingProfileDoesNotFallback(t *testing.T) {
	m, err := NewManager(t.TempDir(), "")
	if err != nil {
		t.Skipf("postgres unavailable: %v", err)
	}
	defer m.Close()
	task, err := m.CreateTask("pinned runtime test", "test", nil, 0, 0)
	if err != nil {
		t.Fatal(err)
	}
	defer m.DeleteTask(task.ID, DeleteTaskOptions{})
	fallback, err := m.pg.SaveProfile(&db.LLMProfile{Name: "pinned-fallback", Format: "openai", Model: "fallback", APIKey: "fallback-key"})
	if err != nil {
		t.Fatal(err)
	}
	defer m.pg.DeleteProfile(fallback)
	if _, err := m.ReplaceTaskLLMProfiles(task.ID, []int64{fallback}, fallback); err != nil {
		t.Fatal(err)
	}
	missing := int64(9223372036854775807)
	s := &Server{m: m, llmOn: true, llmProv: &scriptedLLMProvider{}}
	runtime := &taskLLMRuntime{s: s, taskID: task.ID, agentKey: "worker", pinProfileID: &missing}
	if _, err := runtime.current(); err == nil {
		t.Fatal("missing pinned profile silently fell back")
	}
	if _, ok := runtime.activeCfg(); ok || runtime.nonStreaming() || runtime.maxTokens() != 0 {
		t.Fatal("missing pinned profile exposed fallback runtime settings")
	}
}

func TestPinnedRuntimeQuotaDoesNotAdvanceTaskFailover(t *testing.T) {
	m, err := NewManager(t.TempDir(), "")
	if err != nil {
		t.Skipf("postgres unavailable: %v", err)
	}
	defer m.Close()
	task, err := m.CreateTask("pinned quota", "test", nil, 0, 0)
	if err != nil {
		t.Fatal(err)
	}
	defer m.DeleteTask(task.ID, DeleteTaskOptions{})
	profileID, err := m.pg.SaveProfile(&db.LLMProfile{Name: "pinned-quota", Format: "openai", Model: "pinned", APIKey: "key"})
	if err != nil {
		t.Fatal(err)
	}
	defer m.pg.DeleteProfile(profileID)
	if _, err := m.ReplaceTaskLLMProfiles(task.ID, []int64{profileID}, profileID); err != nil {
		t.Fatal(err)
	}
	before, err := m.pg.GetTask(taskIDInt(t, task.ID))
	if err != nil || before == nil || before.ActiveLLMProfileID == nil {
		t.Fatalf("load task failover state: task=%+v err=%v", before, err)
	}
	provider := &scriptedLLMProvider{err: errors.New("insufficient_quota")}
	s := &Server{m: m, provByProfile: map[int64]*provEntry{profileID: {prov: provider, cfg: agent.Config{}}}}
	id := profileID
	runtime := &taskLLMRuntime{s: s, taskID: task.ID, agentKey: "worker", pinProfileID: &id}
	if _, _, _, err := runtime.Complete(context.Background(), llm.CompletionRequest{}); err == nil {
		t.Fatal("pinned quota error unexpectedly succeeded")
	}
	after, err := m.pg.GetTask(taskIDInt(t, task.ID))
	if err != nil || after == nil || after.LLMChainRevision != before.LLMChainRevision || after.ActiveLLMProfileID == nil || *after.ActiveLLMProfileID != *before.ActiveLLMProfileID {
		t.Fatalf("pinned quota changed task failover state: before=%+v after=%+v err=%v", before, after, err)
	}
}

func taskIDInt(t *testing.T, id string) int64 {
	t.Helper()
	n, err := strconv.ParseInt(id, 10, 64)
	if err != nil {
		t.Fatal(err)
	}
	return n
}

func TestPinnedWorkerLoopDoesNotConsultTaskRouterAndReportsFailureOnce(t *testing.T) {
	m, err := NewManager(t.TempDir(), "")
	if err != nil {
		t.Skipf("postgres unavailable: %v", err)
	}
	defer m.Close()
	task, err := m.CreateTask("pinned worker loop", "test", nil, 0, 0)
	if err != nil {
		t.Fatal(err)
	}
	defer m.DeleteTask(task.ID, DeleteTaskOptions{})

	e := NewEngine(m)
	var routerCalls atomic.Int32
	e.SetAuthoritativeAgentResolver(func(*Task) (*agent.Planner, *agent.Worker) {
		routerCalls.Add(1)
		return nil, nil
	})
	e.SetWorkerResolver(func(*Task, WorkerSlot) (*agent.Worker, error) {
		return nil, errors.New("profile unavailable")
	})
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan struct{})
	go func() {
		e.workerLoop(ctx, task, WorkerSlot{Name: "work#4", ProfileID: 12})
		close(done)
	}()
	time.Sleep(100 * time.Millisecond)
	cancel()
	select {
	case <-done:
	case <-time.After(time.Second):
		t.Fatal("worker loop did not stop")
	}
	if routerCalls.Load() != 0 {
		t.Fatalf("pinned worker consulted task router %d times", routerCalls.Load())
	}
	activities, _, err := task.Store.ActivityList(nil, 0, 20)
	if err != nil {
		t.Fatal(err)
	}
	var failures int
	for _, activity := range activities {
		if activity.Worker == "system" && activity.IsError {
			failures++
		}
	}
	if failures != 1 {
		t.Fatalf("pinned failure activities=%d, want one throttled entry: %+v", failures, activities)
	}
}

func TestPinnedWorkerAvailabilityRecoveryUsesSystemSession(t *testing.T) {
	m, err := NewManager(t.TempDir(), "")
	if err != nil {
		t.Skipf("postgres unavailable: %v", err)
	}
	defer m.Close()
	task, err := m.CreateTask("pinned availability", "test", nil, 0, 0)
	if err != nil {
		t.Fatal(err)
	}
	defer m.DeleteTask(task.ID, DeleteTaskOptions{})
	e := NewEngine(m)
	slot := WorkerSlot{Name: "work#4", ProfileID: 12}
	e.reportSlotAvailability(task, slot, errors.New("profile unavailable"))
	e.reportSlotAvailability(task, slot, nil)
	activities, _, err := task.Store.ActivityList(nil, 0, 20)
	if err != nil {
		t.Fatal(err)
	}
	var failure, recovery int
	for _, activity := range activities {
		if activity.Worker != "system" {
			continue
		}
		if activity.IsError && strings.Contains(activity.Summary, "work#4") && strings.Contains(activity.Summary, "#12") {
			failure++
		}
		if !activity.IsError && strings.Contains(activity.Summary, "work#4") && strings.Contains(activity.Summary, "#12") {
			recovery++
		}
	}
	if failure != 1 || recovery != 1 {
		t.Fatalf("system availability activities failure=%d recovery=%d: %+v", failure, recovery, activities)
	}
}

func TestDetachedIntentPinnedOwnerUsesPinnedResolver(t *testing.T) {
	m, err := NewManager(t.TempDir(), "")
	if err != nil {
		t.Skipf("postgres unavailable: %v", err)
	}
	defer m.Close()
	task, err := m.CreateTask("detached pinned owner", "test", nil, 0, 0)
	if err != nil {
		t.Fatal(err)
	}
	defer m.DeleteTask(task.ID, DeleteTaskOptions{})
	intentID, err := task.Store.AddIntent(map[string]any{"text": "continue"}, 0, nil, "test")
	if err != nil {
		t.Fatal(err)
	}
	if ok, err := task.Store.ClaimIntent(intentID, "work#4"); err != nil || !ok {
		t.Fatalf("claim intent: ok=%v err=%v", ok, err)
	}
	if changed, err := task.Store.CompareAndSetIntentState(intentID, "running", "paused"); err != nil || !changed {
		t.Fatalf("pause intent: changed=%v err=%v", changed, err)
	}
	e := NewEngine(m)
	e.slotPlans.Store(task.ID, map[string]int64{"work#4": 12})
	var pinnedCalls, snapshotCalls atomic.Int32
	e.SetWorkerResolver(func(_ *Task, slot WorkerSlot) (*agent.Worker, error) {
		pinnedCalls.Add(1)
		if slot.Name != "work#4" || slot.ProfileID != 12 {
			t.Fatalf("unexpected pinned slot: %+v", slot)
		}
		return nil, errors.New("profile unavailable")
	})
	e.SetAuthoritativeAgentResolver(func(*Task) (*agent.Planner, *agent.Worker) {
		snapshotCalls.Add(1)
		return nil, nil
	})
	err = e.runDetachedIntent(context.Background(), task, intentID, "request", "human", "agent")
	if err == nil || !strings.Contains(err.Error(), "LLM 配置 #12 不可用") {
		t.Fatalf("runDetachedIntent err=%v", err)
	}
	if pinnedCalls.Load() != 1 || snapshotCalls.Load() != 0 {
		t.Fatalf("pinned calls=%d snapshot calls=%d", pinnedCalls.Load(), snapshotCalls.Load())
	}
	node, err := task.Store.GetNode(intentID)
	if err != nil || node == nil || node.State != "paused" {
		t.Fatalf("intent changed after unavailable resolver: node=%+v err=%v", node, err)
	}
}