Why: Buffering a whole attempt hid live output, overflow detection missed common provider errors, and a cut-off tool call could still run. Evidence: GOWORK=off go test -race -count=1 ./runtime/llm/ ./runtime/ passed. Impact: WithRetry forwards events immediately and retries only when nothing has been forwarded. Length or an unclean end drops unfinished tool calls. Loop usage keeps a billing sum and the last request separately. Refs: feature/llm-robustness Co-authored-by: Cursor <cursoragent@cursor.com>
232 lines
7.1 KiB
Go
232 lines
7.1 KiB
Go
package llm
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestParseModelRefAndThinking(t *testing.T) {
|
|
ref, err := ParseModelRef("anthropic/claude-sonnet:high")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if ref.Provider != "anthropic" || ref.Model != "claude-sonnet" || ref.Thinking != "high" {
|
|
t.Fatalf("%+v", ref)
|
|
}
|
|
if _, err := ParseModelRef("openai/gpt:nope"); err == nil {
|
|
t.Fatal("unknown thinking level must fail")
|
|
}
|
|
if AnthropicThinkingBudget("high") != 8192 || OpenAIReasoningEffort("off") != "" {
|
|
t.Fatal("thinking map")
|
|
}
|
|
p, ok := Lookup("groq", nil)
|
|
if !ok || !strings.Contains(p.BaseURL, "groq") || p.Protocol != ProtocolOpenAIChat {
|
|
t.Fatalf("preset: %+v", p)
|
|
}
|
|
}
|
|
|
|
func TestModelsJSONAndBodies(t *testing.T) {
|
|
dir := t.TempDir()
|
|
path := filepath.Join(dir, "models.json")
|
|
raw := []byte(`{"providers":[{"name":"local","protocol":"openai-chat","baseURL":"http://127.0.0.1:9","apiKeyEnv":"UNUSED_KEY"}]}`)
|
|
if err := os.WriteFile(path, raw, 0o644); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
got, err := LoadModelsJSON(path)
|
|
if err != nil || len(got) != 1 || got[0].Name != "local" {
|
|
t.Fatalf("%+v %v", got, err)
|
|
}
|
|
msgs := []Message{{
|
|
Role: RoleUser, Content: "see",
|
|
Images: []Image{{MIME: "image/png", Data: "aaaa"}},
|
|
}}
|
|
body, err := BuildBody(ProtocolAnthropic, "claude", append([]Message{{Role: RoleSystem, Content: "sys"}}, msgs...), nil, StreamOptions{ThinkingLevel: "low"})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if !strings.Contains(string(body), "budget_tokens") || !strings.Contains(string(body), "base64") {
|
|
t.Fatalf("anthropic body: %s", body)
|
|
}
|
|
gbody, err := BuildBody(ProtocolGoogle, "gemini", msgs, nil, StreamOptions{ThinkingLevel: "max"})
|
|
if err != nil || !strings.Contains(string(gbody), "inlineData") || !strings.Contains(string(gbody), "thinkingBudget") {
|
|
t.Fatalf("google: %s %v", gbody, err)
|
|
}
|
|
rbody, err := BuildBody(ProtocolOpenAIResponses, "gpt", msgs, nil, StreamOptions{ThinkingLevel: "medium", MaxTokens: 20})
|
|
if err != nil || !strings.Contains(string(rbody), "max_output_tokens") || !strings.Contains(string(rbody), "image_url") {
|
|
t.Fatalf("responses: %s %v", rbody, err)
|
|
}
|
|
var decoded map[string]any
|
|
if err := json.Unmarshal(rbody, &decoded); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
|
|
func TestPresetsAndClientForChatCompletions(t *testing.T) {
|
|
want := map[string]string{
|
|
"anthropic": "https://api.anthropic.com/v1",
|
|
"google": "https://generativelanguage.googleapis.com/v1beta/openai",
|
|
}
|
|
for _, p := range Presets() {
|
|
base, ok := want[p.Name]
|
|
if !ok {
|
|
continue
|
|
}
|
|
if p.Protocol != ProtocolOpenAIChat || p.BaseURL != base {
|
|
t.Fatalf("%s preset: %+v", p.Name, p)
|
|
}
|
|
}
|
|
for _, name := range []string{"openai-responses", "azure"} {
|
|
p, ok := Lookup(name, nil)
|
|
if !ok {
|
|
t.Fatal(name)
|
|
}
|
|
if _, err := ClientFor(p); err == nil {
|
|
t.Fatalf("%s should fail closed", name)
|
|
}
|
|
}
|
|
var path string
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
path = r.URL.Path
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
_, _ = w.Write([]byte("data: [DONE]\n\n"))
|
|
}))
|
|
defer srv.Close()
|
|
client, err := ClientFor(Provider{Protocol: ProtocolOpenAIChat, BaseURL: srv.URL, APIKey: "k"})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
for ev := range client.Stream(context.Background(), "m", []Message{{Role: RoleUser, Content: "hi"}}, nil, StreamOptions{}) {
|
|
if ev.Kind == KindError {
|
|
t.Fatal(ev.Err)
|
|
}
|
|
}
|
|
if path != "/chat/completions" {
|
|
t.Fatalf("path %q", path)
|
|
}
|
|
}
|
|
|
|
func TestRetryThenSuccess(t *testing.T) {
|
|
n := 0
|
|
base := StreamFunc(func(ctx context.Context, model string, messages []Message, tools []ToolDecl, opts StreamOptions) <-chan Event {
|
|
out := make(chan Event, 2)
|
|
n++
|
|
if n == 1 {
|
|
out <- Event{Kind: KindError, Retryable: true, ErrorMessage: "429"}
|
|
} else {
|
|
out <- Event{Kind: KindTextDelta, Text: "ok"}
|
|
out <- Event{Kind: KindMessageEnd, Message: Message{Role: RoleAssistant, Content: "ok"}}
|
|
}
|
|
close(out)
|
|
return out
|
|
})
|
|
var events []string
|
|
wrapped := WithRetry(base, 2, func(kind, detail string) { events = append(events, kind) })
|
|
var texts []string
|
|
for ev := range wrapped(context.Background(), "m", nil, nil, StreamOptions{}) {
|
|
if ev.Kind == KindTextDelta {
|
|
texts = append(texts, ev.Text)
|
|
}
|
|
}
|
|
if len(texts) != 1 || texts[0] != "ok" {
|
|
t.Fatalf("texts %v calls %d", texts, n)
|
|
}
|
|
if len(events) == 0 || events[0] != "auto_retry_start" {
|
|
t.Fatalf("events %v", events)
|
|
}
|
|
}
|
|
|
|
func TestContextOverflowPatterns(t *testing.T) {
|
|
hits := []string{
|
|
"provider status 413: payload too large",
|
|
"context_length_exceeded",
|
|
"prompt is too long",
|
|
"input length exceeds the maximum",
|
|
"maximum context length exceeded",
|
|
"context length overflow",
|
|
"too many tokens",
|
|
"输入过长,请缩短上下文",
|
|
"请求过长",
|
|
"超出最大长度",
|
|
}
|
|
for _, msg := range hits {
|
|
if !ContextOverflow(msg) {
|
|
t.Fatalf("want overflow: %s", msg)
|
|
}
|
|
}
|
|
if ContextOverflow("context cancelled") || ContextOverflow("temporary network error") {
|
|
t.Fatal("unrelated context text must not count as overflow")
|
|
}
|
|
}
|
|
|
|
func TestParseRetryAfter(t *testing.T) {
|
|
now := time.Date(2026, 10, 3, 12, 0, 0, 0, time.UTC)
|
|
if d, ok := ParseRetryAfter("2", now); !ok || d != 2*time.Second {
|
|
t.Fatalf("seconds: %v %v", d, ok)
|
|
}
|
|
if d, ok := ParseRetryAfter("not-a-date", now); ok || d != 0 {
|
|
t.Fatalf("garbage should fail: %v %v", d, ok)
|
|
}
|
|
if _, ok := ParseRetryAfter("", now); ok {
|
|
t.Fatal("empty header")
|
|
}
|
|
far := now.Add(2 * time.Hour).Format(http.TimeFormat)
|
|
if d, ok := ParseRetryAfter(far, now); !ok || d != maxRetryAfter {
|
|
t.Fatalf("date cap: %v %v", d, ok)
|
|
}
|
|
past := now.Add(-time.Minute).Format(http.TimeFormat)
|
|
if d, ok := ParseRetryAfter(past, now); !ok || d != 0 {
|
|
t.Fatalf("past date: %v %v", d, ok)
|
|
}
|
|
}
|
|
|
|
func TestRetryDoesNotRepeatDeltas(t *testing.T) {
|
|
n := 0
|
|
base := StreamFunc(func(ctx context.Context, model string, messages []Message, tools []ToolDecl, opts StreamOptions) <-chan Event {
|
|
out := make(chan Event, 2)
|
|
n++
|
|
out <- Event{Kind: KindTextDelta, Text: "a"}
|
|
out <- Event{Kind: KindError, Retryable: true, ErrorMessage: "overloaded"}
|
|
close(out)
|
|
return out
|
|
})
|
|
var texts []string
|
|
var errors int
|
|
for ev := range WithRetry(base, 2, nil)(context.Background(), "m", nil, nil, StreamOptions{}) {
|
|
if ev.Kind == KindTextDelta {
|
|
texts = append(texts, ev.Text)
|
|
}
|
|
if ev.Kind == KindError {
|
|
errors++
|
|
}
|
|
}
|
|
if n != 1 || len(texts) != 1 || texts[0] != "a" || errors != 1 {
|
|
t.Fatalf("calls=%d texts=%v errors=%d", n, texts, errors)
|
|
}
|
|
}
|
|
|
|
func TestContextOverflowIsNotRetried(t *testing.T) {
|
|
if !ContextOverflow("maximum context length exceeded") {
|
|
t.Fatal("expected overflow")
|
|
}
|
|
n := 0
|
|
base := StreamFunc(func(ctx context.Context, model string, messages []Message, tools []ToolDecl, opts StreamOptions) <-chan Event {
|
|
out := make(chan Event, 1)
|
|
n++
|
|
out <- Event{Kind: KindError, Retryable: true, ErrorMessage: "context length overflow"}
|
|
close(out)
|
|
return out
|
|
})
|
|
for range WithRetry(base, 3, nil)(context.Background(), "m", nil, nil, StreamOptions{}) {
|
|
}
|
|
if n != 1 {
|
|
t.Fatalf("overflow must not retry, calls=%d", n)
|
|
}
|
|
}
|