Agent-System-Backend/internal/runtime/usage_n3_test.go

package runtime

import (
	"context"
	"errors"
	"testing"

	agentllm "agent-native/runtime/llm"
	core "agent-system-core"
)

// ---- N3 WP6：计费拆分与重试收口（backend 侧）----

// retryAttemptsOf 映射：设置 0=不重试→-1；N→N；默认 2→2。
func TestRetryAttemptsMapping(t *testing.T) {
	if got := retryAttemptsOf(core.EffectiveSettings{}); got != 2 {
		t.Fatalf("默认应 2，得到 %d", got)
	}
	if got := retryAttemptsOf(core.EffectiveSettings{"runtime": map[string]any{"llmRetries": float64(0)}}); got != -1 {
		t.Fatalf("设置 0 应映射为关闭（-1），得到 %d", got)
	}
	if got := retryAttemptsOf(core.EffectiveSettings{"runtime": map[string]any{"llmRetries": float64(5)}}); got != 5 {
		t.Fatalf("设置 5 应映射为 5，得到 %d", got)
	}
}

// 直连 run 的用量拆分：message_end 到达后 Last* 覆盖、cache 累计入账。
func TestDirectRunUsageSplitOnRecord(t *testing.T) {
	rc := &fakeRunCtx{ctx: context.Background(), prof: core.AgentProfile{Model: core.ModelRoute{Chat: "test-model"}}}
	stream := func(ctx context.Context, model string, msgs []agentllm.Message, _ []agentllm.ToolDecl, _ agentllm.StreamOptions) <-chan agentllm.Event {
		ch := make(chan agentllm.Event, 1)
		ch <- agentllm.Event{Kind: agentllm.KindMessageEnd,
			Message: agentllm.Message{Role: agentllm.RoleAssistant, Content: "x", StopReason: agentllm.StopEnd},
			Usage:   agentllm.Usage{PromptTokens: 120, CompletionTokens: 12, TotalTokens: 132, CacheReadTokens: 80, CacheWriteTokens: 6},
		}
		close(ch)
		return ch
	}
	r := &LLMRunner{Stream: stream}
	if err := r.Execute(rc); err != nil {
		t.Fatalf("Execute: %v", err)
	}
	u := rc.usageTotal
	if u.PromptTokens != 120 || u.CompletionTokens != 12 ||
		u.LastPromptTokens != 120 || u.LastCompletionTokens != 12 ||
		u.CacheReadTokens != 80 || u.CacheWriteTokens != 6 {
		t.Fatalf("用量拆分未入账：%+v", u)
	}
}

// 子代理记账：subRun 的 message_end 累计 token 并经 AddUsage 入账。
func TestSubRunAccountsTokens(t *testing.T) {
	rc := &fakeRunCtx{ctx: context.Background(), prof: core.AgentProfile{Model: core.ModelRoute{Chat: "test-model"}}}
	stream := func(ctx context.Context, model string, msgs []agentllm.Message, _ []agentllm.ToolDecl, _ agentllm.StreamOptions) <-chan agentllm.Event {
		ch := make(chan agentllm.Event, 1)
		ch <- agentllm.Event{Kind: agentllm.KindMessageEnd,
			Message: agentllm.Message{Role: agentllm.RoleAssistant, Content: "子结果", StopReason: agentllm.StopEnd},
			Usage:   agentllm.Usage{PromptTokens: 66, CompletionTokens: 6, TotalTokens: 72, CacheReadTokens: 30},
		}
		close(ch)
		return ch
	}
	r := &SubLLMRunner{Stream: stream}
	if err := r.Execute(rc); err != nil {
		t.Fatalf("Execute: %v", err)
	}
	u := rc.usageTotal
	if u.PromptTokens != 66 || u.CompletionTokens != 6 ||
		u.LastPromptTokens != 66 || u.LastCompletionTokens != 6 || u.CacheReadTokens != 30 {
		t.Fatalf("子代理用量未入账：%+v", u)
	}
}

func TestSubRunAccountsPromptAndCompletionWhenTotalIsMissing(t *testing.T) {
	rc := &fakeRunCtx{ctx: context.Background(), prof: core.AgentProfile{Model: core.ModelRoute{Chat: "test-model"}}}
	stream := func(ctx context.Context, model string, msgs []agentllm.Message, _ []agentllm.ToolDecl, _ agentllm.StreamOptions) <-chan agentllm.Event {
		ch := make(chan agentllm.Event, 1)
		ch <- agentllm.Event{Kind: agentllm.KindMessageEnd, Message: agentllm.Message{Role: agentllm.RoleAssistant, Content: "子结果", StopReason: agentllm.StopEnd}, Usage: agentllm.Usage{PromptTokens: 66, CompletionTokens: 6}}
		close(ch)
		return ch
	}
	if err := (&SubLLMRunner{Stream: stream}).Execute(rc); err != nil {
		t.Fatal(err)
	}
	if rc.usageTotal.TotalTokens != 72 {
		t.Fatalf("sub-run lost reported prompt/completion usage: %+v", rc.usageTotal)
	}
}

func TestSubRunAccountsFailedTerminalUsage(t *testing.T) {
	rc := &fakeRunCtx{ctx: context.Background(), prof: core.AgentProfile{Model: core.ModelRoute{Chat: "test-model"}}}
	stream := func(context.Context, string, []agentllm.Message, []agentllm.ToolDecl, agentllm.StreamOptions) <-chan agentllm.Event {
		out := make(chan agentllm.Event, 1)
		out <- agentllm.Event{Kind: agentllm.KindMessageEnd,
			Message: agentllm.Message{Role: agentllm.RoleAssistant, StopReason: agentllm.StopError},
			Usage:   agentllm.Usage{PromptTokens: 12, CompletionTokens: 5, TotalTokens: 17, CacheReadTokens: 3},
			Err:     errors.New("tool stream ended before [DONE]"),
		}
		close(out)
		return out
	}
	if err := (&SubLLMRunner{Stream: stream}).Execute(rc); err == nil {
		t.Fatal("failed terminal was accepted")
	}
	u := rc.usageTotal
	if u.PromptTokens != 12 || u.CompletionTokens != 5 || u.TotalTokens != 17 ||
		u.LastPromptTokens != 12 || u.LastCompletionTokens != 5 || u.CacheReadTokens != 3 {
		t.Fatalf("failed terminal returned Usage but child did not account for it: %+v", u)
	}
}

func TestDirectAndSubRunsPassStableProviderRunID(t *testing.T) {
	for _, tc := range []struct {
		name      string
		newRunner func(agentllm.StreamFunc) Runner
	}{
		{"direct", func(stream agentllm.StreamFunc) Runner { return &LLMRunner{Stream: stream} }},
		{"subagent", func(stream agentllm.StreamFunc) Runner { return &SubLLMRunner{Stream: stream} }},
	} {
		t.Run(tc.name, func(t *testing.T) {
			rc := &fakeRunCtx{ctx: context.Background(), prof: core.AgentProfile{Model: core.ModelRoute{Chat: "test-model", Lightweight: "test-model"}}}
			observed := ""
			stream := func(ctx context.Context, model string, msgs []agentllm.Message, _ []agentllm.ToolDecl, _ agentllm.StreamOptions) <-chan agentllm.Event {
				observed = ProviderRunID(ctx)
				out := make(chan agentllm.Event, 1)
				out <- agentllm.Event{Kind: agentllm.KindMessageEnd, Message: agentllm.Message{Role: agentllm.RoleAssistant, Content: "完成", StopReason: agentllm.StopEnd}, Usage: agentllm.Usage{PromptTokens: 2, CompletionTokens: 1, TotalTokens: 3}}
				close(out)
				return out
			}
			if err := tc.newRunner(stream).Execute(rc); err != nil {
				t.Fatal(err)
			}
			if observed != string(rc.Run().ID) {
				t.Fatalf("provider run identity changed or missing: got %q want %q", observed, rc.Run().ID)
			}
		})
	}
}