First Commit

db/intent_control_test.go

@@ -0,0 +1,285 @@
+package db
+
+import (
+	"fmt"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+)
+
+func TestCompareAndSetIntentStateAllowsSingleWinner(t *testing.T) {
+	d, err := Open(testDSN(t))
+	if err != nil {
+		t.Skipf("postgres unavailable (%v) - skipping", err)
+	}
+	defer d.Close()
+	expID, err := d.CreateExploration("intent CAS", "only one controller wins")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer d.Exec(`DELETE FROM explorations WHERE id=$1`, expID)
+	store := d.Exploration(expID)
+	intentID, err := store.AddIntent(map[string]any{"summary": "controlled"}, 1, nil, "planner")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if claimed, err := store.ClaimIntent(intentID, "worker"); err != nil || !claimed {
+		t.Fatalf("claim: claimed=%v err=%v", claimed, err)
+	}
+
+	const contenders = 12
+	var winners atomic.Int32
+	var wg sync.WaitGroup
+	start := make(chan struct{})
+	for range contenders {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			<-start
+			changed, transitionErr := store.CompareAndSetIntentState(intentID, "running", "paused")
+			if transitionErr != nil {
+				t.Errorf("transition: %v", transitionErr)
+				return
+			}
+			if changed {
+				winners.Add(1)
+			}
+		}()
+	}
+	close(start)
+	wg.Wait()
+	if got := winners.Load(); got != 1 {
+		t.Fatalf("CAS winners=%d, want 1", got)
+	}
+	node, err := store.GetNode(intentID)
+	if err != nil || node == nil || node.State != "paused" {
+		t.Fatalf("node=%+v err=%v, want paused", node, err)
+	}
+}
+
+func TestCancelIntentPreservesTokenMeteringWithoutDoubleCount(t *testing.T) {
+	d, err := Open(testDSN(t))
+	if err != nil {
+		t.Skipf("postgres unavailable (%v) - skipping", err)
+	}
+	defer d.Close()
+	expID, err := d.CreateExploration("cancel token rollup", "preserve consumed tokens")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer d.Exec(`DELETE FROM explorations WHERE id=$1`, expID)
+	store := d.Exploration(expID)
+	intentID, err := store.AddIntent(map[string]any{"summary": "cancelled"}, 1, nil, "planner")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if claimed, err := store.ClaimIntent(intentID, "work#1"); err != nil || !claimed {
+		t.Fatalf("claim: claimed=%v err=%v", claimed, err)
+	}
+
+	baselineDaily, err := d.TokenDailyAll(30)
+	if err != nil {
+		t.Fatal(err)
+	}
+	baseline := dailyTokenBuckets(baselineDaily)
+
+	appendUsage := func(kind string, at time.Time, input, output, read, write *int) {
+		t.Helper()
+		activityID, appendErr := store.AppendActivity(Activity{
+			NodeID: &intentID, Worker: "work#1", Kind: kind,
+			InputTokens: input, OutputTokens: output,
+			CacheReadTokens: read, CacheWriteTokens: write,
+		})
+		if appendErr != nil {
+			t.Fatal(appendErr)
+		}
+		if _, updateErr := d.Exec(`UPDATE activity SET created_at=$1 WHERE id=$2`, at, activityID); updateErr != nil {
+			t.Fatal(updateErr)
+		}
+	}
+	values := func(input, output, read, write int) (*int, *int, *int, *int) {
+		return &input, &output, &read, &write
+	}
+
+	now := time.Now().UTC()
+	atUTCNoon := func(daysAgo int) time.Time {
+		at := now.AddDate(0, 0, -daysAgo)
+		return time.Date(at.Year(), at.Month(), at.Day(), 12, 0, 0, 0, time.UTC)
+	}
+	authoritativeDay := atUTCNoon(6)
+	fallbackDay := atUTCNoon(4)
+	trailingDay := atUTCNoon(2)
+
+	// Run 1 has an authoritative result; its preceding cumulative frame must not
+	// be counted a second time or move usage to the frame's earlier date.
+	i, o, r, w := values(100, 100, 100, 100)
+	appendUsage("usage", atUTCNoon(7), i, o, r, w)
+	i, o, r, w = values(10, 11, 12, 13)
+	appendUsage("result", authoritativeDay, i, o, r, w)
+	// Run 2 models a legacy failed result that omitted usage. Its fallback belongs
+	// to the result's date, not the cumulative frame's date.
+	i, o, r, w = values(20, 21, 22, 23)
+	appendUsage("usage", atUTCNoon(5), i, o, r, w)
+	appendUsage("result", fallbackDay, nil, nil, nil, nil)
+	// Run 3 was interrupted before a terminal result was persisted.
+	i, o, r, w = values(30, 31, 32, 33)
+	appendUsage("usage", trailingDay, i, o, r, w)
+
+	beforeCancelDaily, err := d.TokenDailyAll(30)
+	if err != nil {
+		t.Fatal(err)
+	}
+	beforeCancel := dailyTokenBuckets(beforeCancelDaily)
+	assertDailyTokenDelta(t, beforeCancel, baseline, authoritativeDay, 10, 11, 12)
+	assertDailyTokenDelta(t, beforeCancel, baseline, fallbackDay, 0, 0, 0)
+	assertDailyTokenDelta(t, beforeCancel, baseline, trailingDay, 0, 0, 0)
+
+	if err := store.SetIntentState(intentID, "paused"); err != nil {
+		t.Fatal(err)
+	}
+	cleanup, err := store.CancelIntent(intentID)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if cleanup.Activities != 5 {
+		t.Fatalf("deleted activities=%d, want 5", cleanup.Activities)
+	}
+	total, err := store.TokenTotal()
+	if err != nil {
+		t.Fatal(err)
+	}
+	assertTokenUsage(t, total, 60, 63, 66, 69)
+	afterCancelDaily, err := d.TokenDailyAll(30)
+	if err != nil {
+		t.Fatal(err)
+	}
+	afterCancel := dailyTokenBuckets(afterCancelDaily)
+	assertDailyTokenDelta(t, afterCancel, baseline, authoritativeDay, 10, 11, 12)
+	assertDailyTokenDelta(t, afterCancel, baseline, fallbackDay, 20, 21, 22)
+	assertDailyTokenDelta(t, afterCancel, baseline, trailingDay, 30, 31, 32)
+	assertDailyTokenDelta(t, afterCancel, baseline, now, 0, 0, 0)
+
+	sessions, err := store.TokenStatsBySession()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(sessions) != 0 {
+		t.Fatalf("cancelled intent leaked into sessions: %+v", sessions)
+	}
+	var rollups, datedRollups int
+	if err := d.QueryRow(`SELECT COUNT(*), COUNT(*) FILTER (
+			WHERE metadata->>'token_day'=TO_CHAR(created_at AT TIME ZONE 'UTC', 'YYYY-MM-DD')
+		) FROM activity
+		WHERE exploration_id=$1 AND worker='token-ledger' AND kind='result'
+		  AND metadata->>'cancelled_intent_id'=$2`, expID, fmt.Sprint(intentID)).Scan(&rollups, &datedRollups); err != nil {
+		t.Fatal(err)
+	}
+	if rollups != 3 || datedRollups != rollups {
+		t.Fatalf("token rollups=%d dated=%d, want three correctly dated rows", rollups, datedRollups)
+	}
+	if _, err := store.CancelIntent(intentID); err == nil {
+		t.Fatal("second cancellation unexpectedly succeeded")
+	}
+	afterRetry, err := store.TokenTotal()
+	if err != nil {
+		t.Fatal(err)
+	}
+	assertTokenUsage(t, afterRetry, 60, 63, 66, 69)
+	afterRetryDaily, err := d.TokenDailyAll(30)
+	if err != nil {
+		t.Fatal(err)
+	}
+	assertDailyTokenBucketsEqual(t, dailyTokenBuckets(afterRetryDaily), afterCancel)
+}
+
+func dailyTokenBuckets(items []DailyTokenBucket) map[string]DailyTokenBucket {
+	out := make(map[string]DailyTokenBucket, len(items))
+	for _, item := range items {
+		out[item.Day] = item
+	}
+	return out
+}
+
+func assertDailyTokenDelta(t *testing.T, got, baseline map[string]DailyTokenBucket, at time.Time, input, output, read int) {
+	t.Helper()
+	day := at.UTC().Format(time.DateOnly)
+	actual, initial := got[day], baseline[day]
+	if actual.InputTokens-initial.InputTokens != input ||
+		actual.OutputTokens-initial.OutputTokens != output ||
+		actual.CacheReadTokens-initial.CacheReadTokens != read {
+		t.Fatalf("token delta for %s = input:%d output:%d cache-read:%d, want %d/%d/%d",
+			day, actual.InputTokens-initial.InputTokens, actual.OutputTokens-initial.OutputTokens,
+			actual.CacheReadTokens-initial.CacheReadTokens, input, output, read)
+	}
+}
+
+func assertDailyTokenBucketsEqual(t *testing.T, got, want map[string]DailyTokenBucket) {
+	t.Helper()
+	if len(got) != len(want) {
+		t.Fatalf("daily token buckets changed after retry: got=%+v want=%+v", got, want)
+	}
+	for day, expected := range want {
+		if actual, ok := got[day]; !ok || actual != expected {
+			t.Fatalf("daily token bucket %s changed after retry: got=%+v want=%+v", day, actual, expected)
+		}
+	}
+}
+
+func TestCancelIntentPreservesOutputsYieldedByAnotherIntent(t *testing.T) {
+	d, err := Open(testDSN(t))
+	if err != nil {
+		t.Skipf("postgres unavailable (%v) - skipping", err)
+	}
+	defer d.Close()
+	expID, err := d.CreateExploration("shared intent output", "preserve shared facts and findings")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer d.Exec(`DELETE FROM explorations WHERE id=$1`, expID)
+	store := d.Exploration(expID)
+	first, err := store.AddIntent(map[string]any{"summary": "first"}, 1, nil, "planner")
+	if err != nil {
+		t.Fatal(err)
+	}
+	second, err := store.AddIntent(map[string]any{"summary": "second"}, 1, nil, "planner")
+	if err != nil {
+		t.Fatal(err)
+	}
+	fact, err := store.AddNode(KindFact, map[string]any{"text": "shared fact"}, 1, "confirmed", "worker", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	finding, err := store.AddNode(KindFinding, map[string]any{"summary": "shared finding"}, 1, "confirmed", "worker", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	for _, intentID := range []int64{first, second} {
+		if err := store.Link(intentID, RelYields, fact); err != nil {
+			t.Fatal(err)
+		}
+		if err := store.Link(intentID, RelYields, finding); err != nil {
+			t.Fatal(err)
+		}
+	}
+	if claimed, err := store.ClaimIntent(first, "work#1"); err != nil || !claimed {
+		t.Fatalf("claim: claimed=%v err=%v", claimed, err)
+	}
+	if err := store.SetIntentState(first, "paused"); err != nil {
+		t.Fatal(err)
+	}
+
+	cleanup, err := store.CancelIntent(first)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if cleanup.Intents != 1 || cleanup.Facts != 0 || cleanup.Findings != 0 {
+		t.Fatalf("cleanup=%+v, want only the cancelled intent", cleanup)
+	}
+	for _, nodeID := range []int64{fact, finding} {
+		node, getErr := store.GetNode(nodeID)
+		if getErr != nil || node == nil {
+			t.Fatalf("shared node %d was removed: node=%+v err=%v", nodeID, node, getErr)
+		}
+	}
+}