36 lines
1.1 KiB
Go
36 lines
1.1 KiB
Go
package runtime
|
|
|
|
import (
|
|
agentllm "agent-native/runtime/llm"
|
|
core "agent-system-core"
|
|
"context"
|
|
"testing"
|
|
)
|
|
|
|
func TestRunBudgetStreamStopsReportedTokenExhaustion(t *testing.T) {
|
|
calls := 0
|
|
cap := 10
|
|
var output int
|
|
base := func(ctx context.Context, m string, msgs []agentllm.Message, tools []agentllm.ToolDecl, opts agentllm.StreamOptions) <-chan agentllm.Event {
|
|
calls++
|
|
output = opts.MaxTokens
|
|
ch := make(chan agentllm.Event, 1)
|
|
ch <- agentllm.Event{Kind: agentllm.KindMessageEnd, Usage: agentllm.Usage{PromptTokens: 8, CompletionTokens: 2, TotalTokens: 10}}
|
|
close(ch)
|
|
return ch
|
|
}
|
|
wrapped := budgetedStream(base, core.RunBudget{MaxTokens: &cap}, 4)
|
|
for range wrapped(context.Background(), "m", nil, nil, agentllm.StreamOptions{}) {
|
|
}
|
|
if output != 4 {
|
|
t.Fatal("output cap ignored", output)
|
|
}
|
|
errorSeen := false
|
|
for e := range wrapped(context.Background(), "m", nil, nil, agentllm.StreamOptions{}) {
|
|
errorSeen = e.Kind == agentllm.KindError
|
|
}
|
|
if !errorSeen || calls != 1 {
|
|
t.Fatal("exhausted budget made another provider call", calls)
|
|
}
|
|
}
|