Agent-Native/runtime/loop_budget_test.go

package runtime

import (
	"context"
	"encoding/json"
	"strings"
	"sync"
	"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")
	}
	// 第一批执行 1 次；第二批超限，两条合成错误结果，不执行。
	toolResults := 0
	skipped := 0
	for _, m := range res.Messages {
		if m.Role != llm.RoleTool {
			continue
		}
		toolResults++
		if m.IsError && m.Content == toolBudgetSkipContent {
			skipped++
		}
	}
	if toolResults != 3 || skipped != 2 {
		t.Fatalf("expected 3 tool results including 2 skipped, got results=%d skipped=%d", toolResults, skipped)
	}
	assertToolCallsPaired(t, res.Messages)
}

func assertToolCallsPaired(t *testing.T, messages []llm.Message) {
	t.Helper()
	pending := map[string]string{}
	for _, m := range messages {
		if m.Role == llm.RoleAssistant {
			for _, call := range m.ToolCalls {
				if _, ok := pending[call.ID]; ok {
					t.Fatalf("duplicate tool call id %s", call.ID)
				}
				pending[call.ID] = call.Name
			}
		}
		if m.Role == llm.RoleTool {
			name, ok := pending[m.ToolCallID]
			if !ok {
				t.Fatalf("tool result %s has no matching call", m.ToolCallID)
			}
			if name != m.Name {
				t.Fatalf("tool result %s name %q != call %q", m.ToolCallID, m.Name, name)
			}
			delete(pending, m.ToolCallID)
		}
	}
	if len(pending) != 0 {
		t.Fatalf("unpaired tool calls: %v", pending)
	}
}

type pauseTool struct{}

func (pauseTool) Name() string        { return "pause" }
func (pauseTool) Description() string { return "slow tool" }
func (pauseTool) Parameters() json.RawMessage {
	return json.RawMessage(`{"type":"object","properties":{}}`)
}
func (pauseTool) Execute(context.Context, json.RawMessage) (tools.Result, error) {
	time.Sleep(40 * time.Millisecond)
	return tools.Result{Content: "ok"}, nil
}

func TestParallelHostCallbacksAreSerial(t *testing.T) {
	reg := tools.NewRegistry(t.TempDir())
	reg.Register(pauseTool{})
	var mu sync.Mutex
	inFlight, maxInFlight := 0, 0
	enter := func() {
		mu.Lock()
		inFlight++
		if inFlight > maxInFlight {
			maxInFlight = inFlight
		}
		mu.Unlock()
		time.Sleep(20 * time.Millisecond)
		mu.Lock()
		inFlight--
		mu.Unlock()
	}
	_, err := Run(context.Background(), Config{
		Model:         "fake",
		Tools:         reg,
		ToolExecution: ExecParallel,
		OnEvent: func(Event) {
			enter()
		},
		BeforeTool: func(llm.ToolCall) error {
			enter()
			return nil
		},
		AfterTool: func(llm.ToolCall, llm.Message) {
			enter()
		},
		Stream: usageScript(llm.Usage{},
			llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{
				{ID: "a", Name: "pause", Arguments: `{}`},
				{ID: "b", Name: "pause", Arguments: `{}`},
			}},
			llm.Message{Role: llm.RoleAssistant, Content: "done", StopReason: llm.StopEnd},
		),
	}, llm.Message{Role: llm.RoleUser, Content: "go"})
	if err != nil {
		t.Fatal(err)
	}
	if maxInFlight != 1 {
		t.Fatalf("host callbacks overlapped: max in flight %d", maxInFlight)
	}
}

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 TestLoopUsageSplitsLastRequestFromCumulative(t *testing.T) {
	res, err := Run(context.Background(), Config{
		Model:  "fake",
		Tools:  tools.WorkspaceTools(t.TempDir()),
		Budget: &Budget{MaxPromptTokens: 10},
		Stream: usageScript(llm.Usage{PromptTokens: 6, CompletionTokens: 1, TotalTokens: 7},
			llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{
				ID: "c1", Name: "list_dir", Arguments: `{}`,
			}}},
			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("last request stays under the window: err=%v text=%q", res.Err, res.Text)
	}
	if res.Usage.PromptTokens != 12 || res.Usage.LastPromptTokens != 6 {
		t.Fatalf("billing sum %d last %d", res.Usage.PromptTokens, res.Usage.LastPromptTokens)
	}
}

func TestLoopRejectsIncompleteToolCall(t *testing.T) {
	tool := &countingTool{name: "list_dir"}
	reg := tools.NewRegistry(t.TempDir())
	reg.Register(tool)
	res, err := Run(context.Background(), Config{
		Model: "fake",
		Tools: reg,
		Stream: script(llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{
			ID: "c1", Name: "list_dir", Arguments: `{"path":`,
		}}}, llm.Message{Role: llm.RoleAssistant, Content: "stopped", StopReason: llm.StopEnd}),
	}, llm.Message{Role: llm.RoleUser, Content: "go"})
	if err != nil {
		t.Fatal(err)
	}
	if tool.calls != 0 {
		t.Fatalf("incomplete call executed %d times", tool.calls)
	}
	var paired bool
	for _, m := range res.Messages {
		if m.Role == llm.RoleTool && m.ToolCallID == "c1" && m.IsError {
			paired = true
		}
	}
	if !paired {
		t.Fatal("expected a paired error result")
	}
}

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)
	}
}