Files
Agent-Native/runtime/llm/providers_test.go
T
weichenandCursor e7ab48e8ea Add native provider protocols and retry
Why: The runtime could only speak OpenAI chat completions, so thinking levels, images, and the other vendor wires had nowhere to go.
Evidence: go test -count=1 ./runtime/llm/
Impact: Existing chat-completions calls stay the same unless a thinking level or image is set. Overflow errors are not retried.
Refs: feature/pi-headless N2
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-10-01 10:10:02 +00:00

116 lines
3.8 KiB
Go

package llm
import (
"context"
"encoding/json"
"os"
"path/filepath"
"strings"
"testing"
)
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 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 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)
}
}