diff --git a/runtime/loop.go b/runtime/loop.go index c4f5dbb..b6c8d6c 100644 --- a/runtime/loop.go +++ b/runtime/loop.go @@ -46,6 +46,12 @@ type Event struct { PromptTokens int CompletionTokens int TotalTokens int + // 以下在 type=usage 时有意义(WP6:上一请求与缓存命中拆分)。 + LastPromptTokens int + LastCompletionTokens int + LastTotalTokens int + CacheReadTokens int + CacheWriteTokens int } // Config drives one agent session. @@ -131,6 +137,10 @@ type Config struct { // MemoryScope 是 Recall 的 scope 过滤("" = 不限,含 user/project; // AgentSpec.MemoryScope 的消费点)。 MemoryScope string + // RetryAttempts 是流式错误的默认重试次数(WP6 下沉为库默认): + // 0 值 = 2 次(默认);负值 = 关闭;正数 = 次数。 + // 只在还没向调用方发出增量时重试(WithRetry 既有语义)。 + RetryAttempts int // traceCtx 是本次 Run 的 span 关联(Run 启动时建立;指针随 Config // 副本传播到各轮与各工具)。 traceCtx *traceIDs @@ -197,6 +207,18 @@ func Run(ctx context.Context, cfg Config, user llm.Message) (Result, error) { if cfg.Memory != nil && cfg.Tools != nil { memory.RegisterWrite(cfg.Tools, cfg.Memory, cfg.MemoryScope) } + // 流式错误默认重试(WP6):库默认 2 次、负值关闭;重试事件进事件面。 + if cfg.Stream != nil { + attempts := cfg.RetryAttempts + if attempts == 0 { + attempts = 2 + } + if attempts > 0 { + cfg.Stream = llm.WithRetry(cfg.Stream, attempts, func(kind, detail string) { + emit(cfg, Event{Type: kind, Text: detail}) + }) + } + } cfg.traceCtx = newTraceContext() endRun := startSpan(cfg, "invoke_agent "+agentName(cfg), map[string]string{ "gen_ai.system": providerName(cfg), @@ -313,12 +335,16 @@ func Run(ctx context.Context, cfg Config, user llm.Message) (Result, error) { end, errEv := consume(ctx, cfg, request) if errEv != nil { cfg.logger().Error("runtime turn failed", errEv, "retryable", end.Retryable) + emitUsage(cfg, usage) return Result{Text: end.Message.Content, Messages: messages, StopReason: end.Message.StopReason, Err: errEv, Retryable: end.Retryable}, nil } last = end.Message usage.PromptTokens += end.Usage.PromptTokens usage.CompletionTokens += end.Usage.CompletionTokens usage.TotalTokens += end.Usage.TotalTokens + // cache 只入累计位(WP6 决策:不加 Last* cache 字段,避免字段膨胀)。 + usage.CacheReadTokens += end.Usage.CacheReadTokens + usage.CacheWriteTokens += end.Usage.CacheWriteTokens usage.LastPromptTokens = end.Usage.PromptTokens usage.LastCompletionTokens = end.Usage.CompletionTokens usage.LastTotalTokens = end.Usage.TotalTokens @@ -1074,7 +1100,12 @@ func drain(fn func() []llm.Message) []llm.Message { } func emitUsage(cfg Config, usage llm.Usage) { - emit(cfg, Event{Type: "usage", PromptTokens: usage.PromptTokens, CompletionTokens: usage.CompletionTokens, TotalTokens: usage.TotalTokens}) + emit(cfg, Event{ + Type: "usage", + PromptTokens: usage.PromptTokens, CompletionTokens: usage.CompletionTokens, TotalTokens: usage.TotalTokens, + LastPromptTokens: usage.LastPromptTokens, LastCompletionTokens: usage.LastCompletionTokens, LastTotalTokens: usage.LastTotalTokens, + CacheReadTokens: usage.CacheReadTokens, CacheWriteTokens: usage.CacheWriteTokens, + }) } func emit(cfg Config, ev Event) { @@ -1107,6 +1138,7 @@ func budgetExceeded(b *Budget, usage llm.Usage, toolCalls int, now time.Time) st // budgetStop 以预算超限收尾:发事件、留日志、按 StopError 返回(不可重试)。 func budgetStop(cfg Config, last llm.Message, messages []llm.Message, usage llm.Usage, reason string) Result { emit(cfg, Event{Type: "budget_exceeded", Text: reason}) + emitUsage(cfg, usage) err := errString(reason) cfg.logger().Error("runtime budget exceeded", err) return Result{Text: last.Content, Messages: messages, StopReason: llm.StopError, Usage: usage, Err: err} diff --git a/runtime/usage_retry_n3_test.go b/runtime/usage_retry_n3_test.go new file mode 100644 index 0000000..5dcaecf --- /dev/null +++ b/runtime/usage_retry_n3_test.go @@ -0,0 +1,203 @@ +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