Agent-System/internal/taskrun/claim_idempotency_test.go

package taskrun

import (
	"context"
	"errors"
	"sync"
	"sync/atomic"
	"testing"
	"time"

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

// memIdemStore 最小内存 store：支持 idempotency_key 唯一占位。
type memIdemStore struct {
	mu     sync.Mutex
	byID   map[string]*RunRecord
	byKey  map[string]string
	events map[string][]StoredEvent
}

func newMemIdemStore() *memIdemStore {
	return &memIdemStore{
		byID:   make(map[string]*RunRecord),
		byKey:  make(map[string]string),
		events: make(map[string][]StoredEvent),
	}
}

func (m *memIdemStore) Available() bool { return true }
func (m *memIdemStore) InsertRun(_ context.Context, rec *RunRecord) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	if rec.IdempotencyKey != "" {
		if _, ok := m.byKey[rec.IdempotencyKey]; ok {
			return ErrIdempotencyConflict
		}
		m.byKey[rec.IdempotencyKey] = rec.RunID
	}
	cpy := *rec
	m.byID[rec.RunID] = &cpy
	return nil
}
func (m *memIdemStore) UpdateRun(_ context.Context, rec *RunRecord) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	cpy := *rec
	m.byID[rec.RunID] = &cpy
	return nil
}
func (m *memIdemStore) GetByIdempotencyKey(_ context.Context, key string) (*RunRecord, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	id, ok := m.byKey[key]
	if !ok {
		return nil, nil
	}
	rec := m.byID[id]
	if rec == nil {
		return nil, nil
	}
	cpy := *rec
	return &cpy, nil
}
func (m *memIdemStore) GetRun(_ context.Context, runID string) (*RunRecord, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	rec := m.byID[runID]
	if rec == nil {
		return nil, errors.New("run not found")
	}
	cpy := *rec
	return &cpy, nil
}
func (m *memIdemStore) AppendEvent(_ context.Context, runID string, seq int64, ev stream.StreamEvent) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.events[runID] = append(m.events[runID], StoredEvent{Seq: seq, Event: ev})
	return nil
}
func (m *memIdemStore) ListEvents(_ context.Context, runID string, afterSeq int64) ([]StoredEvent, error) {
	return nil, nil
}
func (m *memIdemStore) AcquireLease(context.Context, string, string, time.Duration) error {
	return nil
}
func (m *memIdemStore) HeartbeatLease(context.Context, string, string, time.Duration) error {
	return nil
}
func (m *memIdemStore) ReleaseLease(context.Context, string, string) error { return nil }
func (m *memIdemStore) ResetForResume(context.Context, string, int) error  { return nil }
func (m *memIdemStore) MarkOrphanedRunning(context.Context) (int, error)   { return 0, nil }

func TestClaimRun_ConcurrentSameKey_OneOwner(t *testing.T) {
	store := newMemIdemStore()
	mgr := NewManager(nil, nil)
	mgr.SetStore(store)

	const n = 16
	var owners atomic.Int32
	var losers atomic.Int32
	var wg sync.WaitGroup
	key := BuildTaskRunIdempotencyKey("sess-conc", "turn-conc-1")
	runIDs := make(chan string, n)

	wg.Add(n)
	for i := 0; i < n; i++ {
		go func(i int) {
			defer wg.Done()
			rec := &RunRecord{
				RunID:          "",
				SessionID:      "sess-conc",
				ClientTurnID:   "turn-conc-1",
				IdempotencyKey: key,
				UserID:         1,
				TaskKey:        "intake",
				Status:         "pending",
			}
			claimed, created, err := mgr.ClaimRun(context.Background(), rec)
			if err != nil {
				t.Errorf("claim: %v", err)
				return
			}
			if created {
				owners.Add(1)
				runIDs <- claimed.RunID
			} else {
				losers.Add(1)
				runIDs <- claimed.RunID
			}
		}(i)
	}
	wg.Wait()
	close(runIDs)

	if owners.Load() != 1 {
		t.Fatalf("want exactly 1 claim owner, got %d", owners.Load())
	}
	if losers.Load() != n-1 {
		t.Fatalf("want %d losers, got %d", n-1, losers.Load())
	}
	first := ""
	for id := range runIDs {
		if first == "" {
			first = id
		} else if id != first {
			t.Fatalf("all claimants must share same run_id, got %s vs %s", first, id)
		}
	}
	if len(store.byID) != 1 {
		t.Fatalf("store must have 1 run, got %d", len(store.byID))
	}
}

func TestClaimThenStart_Concurrent_OneExecutor(t *testing.T) {
	store := newMemIdemStore()
	mgr := NewManager(nil, nil)
	mgr.SetStore(store)

	var execCalls atomic.Int32
	mgr.SetExecutor(func(ctx context.Context, runID string, req bizdev.TaskRunRequest, st *stream.Streamer) (*bizdev.TaskRunResult, error) {
		execCalls.Add(1)
		time.Sleep(20 * time.Millisecond)
		return &bizdev.TaskRunResult{Success: true, WorkDir: t.TempDir(), TaskKey: "intake"}, nil
	}, nil, nil)

	const n = 12
	var chatWrites atomic.Int32
	var starts atomic.Int32
	key := BuildTaskRunIdempotencyKey("sess-exec", "turn-exec-1")
	var wg sync.WaitGroup
	wg.Add(n)
	for i := 0; i < n; i++ {
		go func() {
			defer wg.Done()
			rec := &RunRecord{
				SessionID:      "sess-exec",
				ClientTurnID:   "turn-exec-1",
				IdempotencyKey: key,
				UserID:         7,
				TaskKey:        "intake",
				Status:         "pending",
			}
			claimed, created, err := mgr.ClaimRun(context.Background(), rec)
			if err != nil {
				t.Errorf("claim: %v", err)
				return
			}
			if !created {
				return
			}
			// 模拟 handler：仅 owner 写 1 user + 1 assistant
			chatWrites.Add(2)
			claimed.AssistantMsgID = 42
			_ = mgr.BindAssistantMsgID(context.Background(), claimed.RunID, 42)
			starts.Add(1)
			_, err = mgr.StartClaimedRun(context.Background(), claimed, bizdev.TaskRunRequest{TaskKey: "intake"}, bizdev.StageRunRequest{}, "")
			if err != nil {
				t.Errorf("start: %v", err)
			}
		}()
	}
	wg.Wait()

	deadline := time.Now().Add(3 * time.Second)
	for time.Now().Before(deadline) {
		if execCalls.Load() >= 1 {
			break
		}
		time.Sleep(20 * time.Millisecond)
	}
	if starts.Load() != 1 || chatWrites.Load() != 2 {
		t.Fatalf("want 1 start + 2 chat writes, got starts=%d chats=%d", starts.Load(), chatWrites.Load())
	}
	if execCalls.Load() != 1 {
		t.Fatalf("want 1 executor call, got %d", execCalls.Load())
	}
	if len(store.byID) != 1 {
		t.Fatalf("want 1 persisted run, got %d", len(store.byID))
	}
}

type failUpdateStore struct {
	*memIdemStore
	updateErr error
}

func (f *failUpdateStore) UpdateRun(ctx context.Context, rec *RunRecord) error {
	if f.updateErr != nil {
		return f.updateErr
	}
	return f.memIdemStore.UpdateRun(ctx, rec)
}

func TestStartClaimedRun_UpdateRunFailureDoesNotStartWorker(t *testing.T) {
	base := newMemIdemStore()
	store := &failUpdateStore{memIdemStore: base, updateErr: errors.New("persist boom")}
	mgr := NewManager(nil, nil)
	mgr.SetStore(store)

	rec := &RunRecord{
		SessionID:      "sess-upd-fail",
		ClientTurnID:   "turn-upd-fail",
		IdempotencyKey: BuildTaskRunIdempotencyKey("sess-upd-fail", "turn-upd-fail"),
		UserID:         1,
		TaskKey:        "intake",
		Status:         "pending",
	}
	claimed, created, err := mgr.ClaimRun(context.Background(), rec)
	if err != nil || !created {
		t.Fatalf("claim: created=%v err=%v", created, err)
	}
	_, err = mgr.StartClaimedRun(context.Background(), claimed, bizdev.TaskRunRequest{TaskKey: "intake"}, bizdev.StageRunRequest{}, "")
	if err == nil {
		t.Fatal("expected StartClaimedRun to fail when UpdateRun fails")
	}
	mgr.mu.Lock()
	_, live := mgr.runs[claimed.RunID]
	mgr.mu.Unlock()
	if live {
		t.Fatal("worker must not start when UpdateRun persistence fails")
	}
}