Agent-Native/runtime/usage_retry_n3_test.go

package runtime

import (
	"context"
	"errors"
	"strings"
	"sync"
	"testing"
	"time"

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

// ---- N3 WP6：计费拆分与默认重试 ----

// cache 累计 + usage 事件带 Last*/cache；终态路径（budgetStop）也发 usage。
func TestUsageEventSplitsLastAndCache(t *testing.T) {
	var mu sync.Mutex
	var usageEvents []Event
	cfg := Config{
		Model: "fake",
		Stream: func() llm.StreamFunc {
			n := 0
			return func(ctx context.Context, model string, msgs []llm.Message, _ []llm.ToolDecl, _ llm.StreamOptions) <-chan llm.Event {
				n++
				ch := make(chan llm.Event, 1)
				ch <- llm.Event{Kind: llm.KindMessageEnd,
					Message: llm.Message{Role: llm.RoleAssistant, Content: "x", StopReason: llm.StopEnd},
					Usage:   llm.Usage{PromptTokens: 50 * n, CompletionTokens: 5, TotalTokens: 50*n + 5},
				}
				close(ch)
				return ch
			}
		}(),
		OnEvent: func(ev Event) {
			if ev.Type == "usage" {
				mu.Lock()
				usageEvents = append(usageEvents, ev)
				mu.Unlock()
			}
		},
	}
	res, err := Run(context.Background(), cfg, llm.Message{Role: llm.RoleUser, Content: "go"})
	if err != nil {
		t.Fatal(err)
	}
	_ = res
	if len(usageEvents) != 1 {
		t.Fatalf("usage 事件数 = %d", len(usageEvents))
	}
	ev := usageEvents[0]
	if ev.PromptTokens == 0 || ev.LastPromptTokens == 0 {
		t.Fatalf("usage 事件必须带累计与上一请求：%+v", ev)
	}
	// 两个回合的累计应≥上一请求（脚本不带 cache 字段时 cache 为 0 但键存在语义）。
	if ev.LastPromptTokens > ev.PromptTokens {
		t.Fatalf("上一请求不得大于累计：%+v", ev)
	}
}

// cache tokens 进入累计（脚本带 cache 字段时）。
func TestUsageAccumulatesCacheTokens(t *testing.T) {
	var ev Event
	stream := func(ctx context.Context, model string, msgs []llm.Message, _ []llm.ToolDecl, _ llm.StreamOptions) <-chan llm.Event {
		ch := make(chan llm.Event, 1)
		ch <- llm.Event{Kind: llm.KindMessageEnd,
			Message: llm.Message{Role: llm.RoleAssistant, Content: "x", StopReason: llm.StopEnd},
			Usage:   llm.Usage{PromptTokens: 100, CompletionTokens: 10, TotalTokens: 110, CacheReadTokens: 60, CacheWriteTokens: 20},
		}
		close(ch)
		return ch
	}
	if _, err := Run(context.Background(), Config{
		Model: "fake", Stream: stream,
		OnEvent: func(e Event) {
			if e.Type == "usage" {
				ev = e
			}
		},
	}, llm.Message{Role: llm.RoleUser, Content: "go"}); err != nil {
		t.Fatal(err)
	}
	if ev.CacheReadTokens != 60 || ev.CacheWriteTokens != 20 {
		t.Fatalf("cache 未累计：%+v", ev)
	}
}

// 预算终态路径也发终态 usage。
func TestBudgetStopEmitsTerminalUsage(t *testing.T) {
	var sawUsage bool
	regB := tools.NewRegistry(t.TempDir())
	regB.Register(&countingTool{name: "read_file"})
	resB, err := Run(context.Background(), Config{
		Model: "fake", Budget: &Budget{MaxToolCalls: 1}, Tools: regB,
		Stream: script(
			llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls,
				ToolCalls: []llm.ToolCall{{ID: "c1", Name: "read_file", Arguments: `{}`}}},
			llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls,
				ToolCalls: []llm.ToolCall{{ID: "c2", Name: "read_file", Arguments: `{}`}}},
		),
		OnEvent: func(e Event) {
			if e.Type == "usage" {
				sawUsage = true
			}
		},
	}, llm.Message{Role: llm.RoleUser, Content: "go"})
	if err != nil || resB.Err == nil {
		t.Fatalf("预算必须触发 budgetStop（Result.Err 非空）：err=%v res.Err=%v", err, resB.Err)
	}
	if !sawUsage {
		t.Fatal("budgetStop 必须发终态 usage")
	}
}

// 默认重试：0 值 = 2 次；事件进事件面；负值关闭。
func TestDefaultRetryAttemptsTwo(t *testing.T) {
	var calls int
	var events []string
	stream := func(ctx context.Context, model string, msgs []llm.Message, _ []llm.ToolDecl, _ llm.StreamOptions) <-chan llm.Event {
		calls++
		ch := make(chan llm.Event, 3)
		if calls < 3 {
			ch <- llm.Event{Kind: llm.KindError, Retryable: true,
				Err: &llm.StreamError{Message: "boom", Retryable: true},
				Message: llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopError}}
		} else {
			ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{Role: llm.RoleAssistant, Content: "成了", StopReason: llm.StopEnd}}
		}
		close(ch)
		return ch
	}
	res, err := Run(context.Background(), Config{
		Model: "fake", Stream: stream,
		OnEvent: func(e Event) { events = append(events, e.Type) },
	}, llm.Message{Role: llm.RoleUser, Content: "go"})
	if err != nil {
		t.Fatalf("第 3 次应成功：%v", err)
	}
	if res.Err != nil {
		t.Fatalf("第 3 次应成功：%v", res.Err)
	}
	if calls != 3 {
		t.Fatalf("默认 2 次重试 + 首发 = 3 次调用，得到 %d", calls)
	}
	var starts int
	for _, t2 := range events {
		if t2 == "auto_retry_start" {
			starts++
		}
	}
	if starts != 2 {
		t.Fatalf("auto_retry_start 事件数 = %d：%v", starts, events)
	}
}

func TestNegativeRetryAttemptsDisables(t *testing.T) {
	var calls int
	stream := func(ctx context.Context, model string, msgs []llm.Message, _ []llm.ToolDecl, _ llm.StreamOptions) <-chan llm.Event {
		calls++
		ch := make(chan llm.Event, 1)
		ch <- llm.Event{Kind: llm.KindError, Retryable: true,
			Err: &llm.StreamError{Message: "boom", Retryable: true},
			Message: llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopError}}
		close(ch)
		return ch
	}
	res2, err := Run(context.Background(), Config{
		Model: "fake", Stream: stream, RetryAttempts: -1,
	}, llm.Message{Role: llm.RoleUser, Content: "go"})
	if err != nil || res2.Err == nil {
		t.Fatalf("必须失败（Result.Err 非空）：err=%v res.Err=%v", err, res2.Err)
	}
	if calls != 1 {
		t.Fatalf("负值必须关闭重试（1 次调用），得到 %d", calls)
	}
}

// 已出增量后不再重试（既有语义保持）。
func TestRetryNeverDuplicatesForwardedDeltas(t *testing.T) {
	var calls int
	stream := func(ctx context.Context, model string, msgs []llm.Message, _ []llm.ToolDecl, _ llm.StreamOptions) <-chan llm.Event {
		calls++
		ch := make(chan llm.Event, 3)
		ch <- llm.Event{Kind: llm.KindTextDelta, Text: "半截"}
		ch <- llm.Event{Kind: llm.KindError, Retryable: true,
			Err: &llm.StreamError{Message: "boom", Retryable: true},
			Message: llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopError}}
		close(ch)
		return ch
	}
	res3, err := Run(context.Background(), Config{Model: "fake", Stream: stream}, llm.Message{Role: llm.RoleUser, Content: "go"})
	if err != nil || res3.Err == nil {
		t.Fatalf("出增量后错误不得重试但须失败：err=%v res.Err=%v", err, res3.Err)
	}
	if calls != 1 {
		t.Fatalf("出增量后不得重试（1 次调用），得到 %d", calls)
	}
}

var _ = errors.New
var _ = strings.Contains
var _ = time.Now