gadfly/cmd/gadfly/executus_test.go

package main

import (
	"context"
	"testing"

	llm "gitea.stevedudenhoeffer.com/steve/majordomo/llm"
)

func TestGadflyBudget(t *testing.T) {
	ctx := context.Background()

	// A nil budget never blocks and never panics.
	var nilB *gadflyBudget
	if err := nilB.Check(ctx, "pr"); err != nil {
		t.Errorf("nil budget Check should be nil, got %v", err)
	}
	nilB.Commit(ctx, "pr", 10)
	nilB.addUsage(llm.Usage{InputTokens: 5})

	// Token ceiling: passes until usage crosses it.
	b := &gadflyBudget{maxTokens: 100}
	if err := b.Check(ctx, "pr"); err != nil {
		t.Fatalf("fresh budget should pass, got %v", err)
	}
	b.addUsage(llm.Usage{InputTokens: 60, OutputTokens: 50}) // 110 >= 100
	if err := b.Check(ctx, "pr"); err == nil {
		t.Error("budget over the token cap should reject the next run")
	}

	// Seconds ceiling, accumulated via Commit.
	s := &gadflyBudget{maxSeconds: 30}
	s.Commit(ctx, "pr", 31)
	if err := s.Check(ctx, "pr"); err == nil {
		t.Error("budget over the time cap should reject the next run")
	}
}

func TestNewPRBudget(t *testing.T) {
	t.Setenv("GADFLY_PR_TOKEN_BUDGET", "")
	t.Setenv("GADFLY_PR_TIME_BUDGET_SECS", "")
	if newPRBudget() != nil {
		t.Error("no caps set should yield a nil (disabled) budget")
	}
	t.Setenv("GADFLY_PR_TOKEN_BUDGET", "1000")
	if b := newPRBudget(); b == nil || b.maxTokens != 1000 {
		t.Errorf("token cap should build a budget with maxTokens=1000, got %+v", b)
	}
}

func TestCompactRatio(t *testing.T) {
	t.Setenv("GADFLY_COMPACT_RATIO", "")
	if got := compactRatio(); got != defaultCompactRatio {
		t.Errorf("default ratio = %v, want %v", got, defaultCompactRatio)
	}
	t.Setenv("GADFLY_COMPACT_RATIO", "0.6")
	if got := compactRatio(); got != 0.6 {
		t.Errorf("ratio override = %v, want 0.6", got)
	}
	for _, bad := range []string{"0", "-1", "2", "nope"} {
		t.Setenv("GADFLY_COMPACT_RATIO", bad)
		if got := compactRatio(); got != defaultCompactRatio {
			t.Errorf("invalid ratio %q should fall back to the default, got %v", bad, got)
		}
	}
}

func TestResolveContextTokensOverride(t *testing.T) {
	// The explicit override short-circuits any model introspection (no network).
	t.Setenv("GADFLY_MODEL_CONTEXT_TOKENS", "123456")
	if got := resolveContextTokens("anything"); got != 123456 {
		t.Errorf("explicit context-token override = %d, want 123456", got)
	}
}

func TestPRCallerID(t *testing.T) {
	t.Setenv("GADFLY_REPO", "")
	t.Setenv("GADFLY_PR", "")
	if got := prCallerID(); got != "local" {
		t.Errorf("no repo/PR should be %q, got %q", "local", got)
	}
	t.Setenv("GADFLY_REPO", "steve/mort")
	t.Setenv("GADFLY_PR", "1367")
	if got := prCallerID(); got != "steve/mort#1367" {
		t.Errorf("callerID = %q, want steve/mort#1367", got)
	}
}

func TestCompactionEnabled(t *testing.T) {
	for _, v := range []string{"", "1", "true", "yes"} {
		t.Setenv("GADFLY_COMPACT", v)
		if !compactionEnabled() {
			t.Errorf("GADFLY_COMPACT=%q should be enabled", v)
		}
	}
	for _, v := range []string{"0", "false", "no", "off"} {
		t.Setenv("GADFLY_COMPACT", v)
		if compactionEnabled() {
			t.Errorf("GADFLY_COMPACT=%q should be disabled", v)
		}
	}
}

func TestGadflyToolRegistry(t *testing.T) {
	fs, err := newRepoFS(t.TempDir(), "diff")
	if err != nil {
		t.Fatal(err)
	}
	_, names, err := gadflyToolRegistry(fs)
	if err != nil {
		t.Fatalf("gadflyToolRegistry: %v", err)
	}
	want := map[string]bool{"read_file": true, "list_dir": true, "grep": true, "find_files": true, "get_diff": true}
	has := func(n string) bool {
		for _, g := range names {
			if g == n {
				return true
			}
		}
		return false
	}
	for n := range want {
		if !has(n) {
			t.Errorf("registry missing tool %q (got %v)", n, names)
		}
	}
	if has("delegate_investigation") {
		t.Error("delegate_investigation must be absent without a worker model")
	}

	// With a worker model the delegate tool is registered too.
	fs.worker = fakeModel(t, "x")
	_, names, err = gadflyToolRegistry(fs)
	if err != nil {
		t.Fatalf("gadflyToolRegistry with worker: %v", err)
	}
	if !has("delegate_investigation") {
		t.Errorf("delegate_investigation should be registered with a worker model, got %v", names)
	}
}