Agent-Native/runtime/loop_budget_test.go

package runtime

import (
	"context"
	"encoding/json"
	"strings"
	"testing"
	"time"

	"agent-native/runtime/llm"
	"agent-native/runtime/tools"
)

// usageScript 产出带 usage 的固定回合，用于预算测试。
func usageScript(usage llm.Usage, turns ...llm.Message) llm.StreamFunc {
	var i int
	return func(ctx context.Context, model string, messages []llm.Message, decls []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event {
		ch := make(chan llm.Event, 4)
		go func() {
			defer close(ch)
			if i >= len(turns) {
				ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{Role: llm.RoleAssistant, Content: "done", StopReason: llm.StopEnd}, Usage: usage}
				return
			}
			msg := turns[i]
			i++
			ch <- llm.Event{Kind: llm.KindMessageEnd, Message: msg, Usage: usage}
		}()
		return ch
	}
}

func TestLoopBudgetToolCallCap(t *testing.T) {
	dir := t.TempDir()
	reg := tools.WorkspaceTools(dir)
	sawBudgetEvent := false
	res, err := Run(context.Background(), Config{
		Model:  "fake",
		Tools:  reg,
		Budget: &Budget{MaxToolCalls: 1},
		OnEvent: func(ev Event) {
			if ev.Type == "budget_exceeded" {
				sawBudgetEvent = true
			}
		},
		Stream: usageScript(llm.Usage{},
			llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{
				ID: "c1", Name: "list_dir", Arguments: `{}`,
			}}},
			llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{
				{ID: "c2", Name: "list_dir", Arguments: `{}`},
				{ID: "c3", Name: "list_dir", Arguments: `{}`},
			}},
			llm.Message{Role: llm.RoleAssistant, Content: "never", StopReason: llm.StopEnd},
		),
	}, llm.Message{Role: llm.RoleUser, Content: "loop tools"})
	if err != nil {
		t.Fatal(err)
	}
	if res.StopReason != llm.StopError || res.Err == nil || !strings.Contains(res.Err.Error(), "budget_exceeded") {
		t.Fatalf("want budget_exceeded stop, got stop=%s err=%v", res.StopReason, res.Err)
	}
	if !strings.Contains(res.Err.Error(), "tool calls") {
		t.Fatalf("want tool-call budget reason, got %v", res.Err)
	}
	if !sawBudgetEvent {
		t.Fatal("expected budget_exceeded event")
	}
	// 第二批工具调用不该被执行（预算在执行前拦截）：transcript 只含一次 tool 结果。
	toolResults := 0
	for _, m := range res.Messages {
		if m.Role == llm.RoleTool {
			toolResults++
		}
	}
	if toolResults != 1 {
		t.Fatalf("expected exactly 1 tool result, got %d", toolResults)
	}
}

func TestLoopBudgetPromptTokens(t *testing.T) {
	res, err := Run(context.Background(), Config{
		Model:  "fake",
		Budget: &Budget{MaxPromptTokens: 10},
		Stream: usageScript(llm.Usage{PromptTokens: 500, TotalTokens: 600},
			llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{
				ID: "c1", Name: "list_dir", Arguments: `{}`,
			}}},
			llm.Message{Role: llm.RoleAssistant, Content: "never", StopReason: llm.StopEnd},
		),
	}, llm.Message{Role: llm.RoleUser, Content: "go"})
	if err != nil {
		t.Fatal(err)
	}
	if res.Err == nil || !strings.Contains(res.Err.Error(), "prompt tokens") {
		t.Fatalf("want prompt-token budget stop, got err=%v", res.Err)
	}
	if res.Usage.PromptTokens != 500 {
		t.Fatalf("usage should be reported on budget stop, got %d", res.Usage.PromptTokens)
	}
}

func TestLoopBudgetDeadline(t *testing.T) {
	res, err := Run(context.Background(), Config{
		Model:  "fake",
		Budget: &Budget{Deadline: time.Now().Add(-time.Second)},
		Stream: usageScript(llm.Usage{}, llm.Message{Role: llm.RoleAssistant, Content: "never", StopReason: llm.StopEnd}),
	}, llm.Message{Role: llm.RoleUser, Content: "go"})
	if err != nil {
		t.Fatal(err)
	}
	if res.Err == nil || !strings.Contains(res.Err.Error(), "deadline") {
		t.Fatalf("want deadline budget stop, got err=%v", res.Err)
	}
}

func TestLoopBudgetNotHitCompletesNormally(t *testing.T) {
	res, err := Run(context.Background(), Config{
		Model:  "fake",
		Budget: &Budget{MaxPromptTokens: 10_000, MaxToolCalls: 10},
		Stream: usageScript(llm.Usage{PromptTokens: 50},
			llm.Message{Role: llm.RoleAssistant, Content: "ok", StopReason: llm.StopEnd}),
	}, llm.Message{Role: llm.RoleUser, Content: "go"})
	if err != nil {
		t.Fatal(err)
	}
	if res.Err != nil || res.Text != "ok" {
		t.Fatalf("clean run should finish, got err=%v text=%q", res.Err, res.Text)
	}
}

// blockingTool 阻塞到 ctx 取消，验证取消传播到工具执行层。
type blockingTool struct{ canceled chan struct{} }

func (b *blockingTool) Name() string        { return "block_forever" }
func (b *blockingTool) Description() string { return "blocks until context cancel" }
func (b *blockingTool) Parameters() json.RawMessage {
	return json.RawMessage(`{"type":"object","properties":{}}`)
}

func (b *blockingTool) Execute(ctx context.Context, args json.RawMessage) (tools.Result, error) {
	<-ctx.Done()
	close(b.canceled)
	return tools.Result{Content: "canceled", IsError: true}, ctx.Err()
}

func TestLoopCancelPropagatesToTool(t *testing.T) {
	bt := &blockingTool{canceled: make(chan struct{})}
	reg := tools.NewRegistry(t.TempDir())
	reg.Register(bt)
	ctx, cancel := context.WithCancel(context.Background())

	firstCall := true
	stream := func(ctx context.Context, model string, messages []llm.Message, decls []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event {
		ch := make(chan llm.Event, 4)
		go func() {
			defer close(ch)
			if firstCall {
				firstCall = false
				ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{
					Role: llm.RoleAssistant, StopReason: llm.StopToolCalls,
					ToolCalls: []llm.ToolCall{{ID: "c1", Name: "block_forever", Arguments: `{}`}},
				}}
				return
			}
			// 工具返回后的下一轮：等取消或超时兜底。
			select {
			case <-ctx.Done():
				ch <- llm.Event{Kind: llm.KindError, Err: ctx.Err(), Message: llm.Message{StopReason: llm.StopAborted}}
			case <-time.After(3 * time.Second):
				ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{Role: llm.RoleAssistant, Content: "done", StopReason: llm.StopEnd}}
			}
		}()
		return ch
	}

	go func() {
		time.Sleep(100 * time.Millisecond)
		cancel()
	}()
	res, err := Run(ctx, Config{Model: "fake", Tools: reg, Stream: stream},
		llm.Message{Role: llm.RoleUser, Content: "block"})
	if err != nil {
		t.Fatal(err)
	}
	select {
	case <-bt.canceled:
	case <-time.After(2 * time.Second):
		t.Fatal("tool context was not canceled")
	}
	if res.StopReason != llm.StopAborted && res.Err == nil {
		t.Fatalf("expected aborted stop, got %+v", res)
	}
}