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>
This commit is contained in:
2026-10-01 10:10:02 +00:00
co-authored by Cursor
parent 6bc1b08540
commit e7ab48e8ea
5 changed files with 575 additions and 3 deletions
+4 -1
View File
@@ -53,6 +53,9 @@ func (c *Client) run(ctx context.Context, model string, messages []Message, tool
if opts.MaxTokens > 0 {
body["max_tokens"] = opts.MaxTokens
}
if effort := OpenAIReasoningEffort(opts.ThinkingLevel); effort != "" {
body["reasoning_effort"] = effort
}
if len(tools) > 0 {
body["tools"] = toProviderTools(tools)
body["tool_choice"] = "auto"
@@ -125,7 +128,7 @@ func toProviderMessages(messages []Message) []map[string]any {
default:
out = append(out, map[string]any{
"role": string(m.Role),
"content": m.Content,
"content": messageContent(m),
})
}
}
+431
View File
@@ -0,0 +1,431 @@
package llm
import (
"context"
"encoding/json"
"fmt"
"os"
"strings"
"time"
)
// Thinking levels, from off to max. Unknown values are rejected by NormalizeThinking.
var ThinkingLevels = []string{"off", "minimal", "low", "medium", "high", "xhigh", "max"}
// NormalizeThinking returns the canonical level. Empty stays empty.
func NormalizeThinking(level string) (string, error) {
if level == "" {
return "", nil
}
level = strings.ToLower(strings.TrimSpace(level))
for _, known := range ThinkingLevels {
if level == known {
return level, nil
}
}
return "", fmt.Errorf("llm: unknown thinking level %q", level)
}
// OpenAIReasoningEffort maps a thinking level onto OpenAI reasoning_effort.
// off and empty produce no parameter.
func OpenAIReasoningEffort(level string) string {
switch level {
case "minimal":
return "minimal"
case "low":
return "low"
case "medium":
return "medium"
case "high", "xhigh", "max":
return "high"
default:
return ""
}
}
// AnthropicThinkingBudget maps a thinking level onto a token budget.
// off and empty disable thinking.
func AnthropicThinkingBudget(level string) int {
switch level {
case "minimal":
return 1024
case "low":
return 2048
case "medium":
return 4096
case "high":
return 8192
case "xhigh":
return 16384
case "max":
return 32768
default:
return 0
}
}
// GoogleThinkingBudget maps a thinking level onto Gemini thinkingBudget.
func GoogleThinkingBudget(level string) int {
switch level {
case "minimal", "low":
return 1024
case "medium":
return 4096
case "high":
return 8192
case "xhigh", "max":
return 24576
default:
return 0
}
}
// ModelRef is provider/model with an optional :thinking suffix.
type ModelRef struct {
Provider string
Model string
Thinking string
}
// ParseModelRef accepts "model", "provider/model", and "provider/model:level".
func ParseModelRef(spec string) (ModelRef, error) {
spec = strings.TrimSpace(spec)
if spec == "" {
return ModelRef{}, fmt.Errorf("llm: empty model ref")
}
thinking := ""
if i := strings.LastIndex(spec, ":"); i >= 0 {
thinking = spec[i+1:]
spec = spec[:i]
norm, err := NormalizeThinking(thinking)
if err != nil {
return ModelRef{}, err
}
thinking = norm
}
provider, model := "", spec
if i := strings.Index(spec, "/"); i >= 0 {
provider = spec[:i]
model = spec[i+1:]
}
if model == "" {
return ModelRef{}, fmt.Errorf("llm: model id missing in %q", spec)
}
return ModelRef{Provider: provider, Model: model, Thinking: thinking}, nil
}
// Protocol names the wire format a provider speaks.
type Protocol string
const (
ProtocolOpenAIChat Protocol = "openai-chat"
ProtocolOpenAIResponses Protocol = "openai-responses"
ProtocolAnthropic Protocol = "anthropic"
ProtocolGoogle Protocol = "google"
ProtocolAzure Protocol = "azure"
)
// Provider is one configured endpoint.
type Provider struct {
Name string `json:"name"`
Protocol Protocol `json:"protocol"`
BaseURL string `json:"baseURL"`
APIKeyEnv string `json:"apiKeyEnv,omitempty"`
APIKey string `json:"-"`
}
// Presets are OpenAI-compatible or native endpoints that only need an API key.
func Presets() []Provider {
return []Provider{
{Name: "anthropic", Protocol: ProtocolAnthropic, BaseURL: "https://api.anthropic.com", APIKeyEnv: "ANTHROPIC_API_KEY"},
{Name: "openai", Protocol: ProtocolOpenAIChat, BaseURL: "https://api.openai.com/v1", APIKeyEnv: "OPENAI_API_KEY"},
{Name: "openai-responses", Protocol: ProtocolOpenAIResponses, BaseURL: "https://api.openai.com/v1", APIKeyEnv: "OPENAI_API_KEY"},
{Name: "google", Protocol: ProtocolGoogle, BaseURL: "https://generativelanguage.googleapis.com/v1beta", APIKeyEnv: "GEMINI_API_KEY"},
{Name: "azure", Protocol: ProtocolAzure, BaseURL: "", APIKeyEnv: "AZURE_OPENAI_API_KEY"},
{Name: "deepseek", Protocol: ProtocolOpenAIChat, BaseURL: "https://api.deepseek.com/v1", APIKeyEnv: "DEEPSEEK_API_KEY"},
{Name: "minimax", Protocol: ProtocolOpenAIChat, BaseURL: "https://api.minimaxi.com/v1", APIKeyEnv: "MINIMAX_API_KEY"},
{Name: "openrouter", Protocol: ProtocolOpenAIChat, BaseURL: "https://openrouter.ai/api/v1", APIKeyEnv: "OPENROUTER_API_KEY"},
{Name: "groq", Protocol: ProtocolOpenAIChat, BaseURL: "https://api.groq.com/openai/v1", APIKeyEnv: "GROQ_API_KEY"},
}
}
// modelsFile is the on-disk shape of models.json.
type modelsFile struct {
Providers []Provider `json:"providers"`
}
// LoadModelsJSON reads custom providers. Missing file yields an empty list.
func LoadModelsJSON(path string) ([]Provider, error) {
raw, err := os.ReadFile(path)
if err != nil {
if os.IsNotExist(err) {
return nil, nil
}
return nil, err
}
var file modelsFile
if err := json.Unmarshal(raw, &file); err != nil {
return nil, err
}
for i := range file.Providers {
if file.Providers[i].APIKey == "" && file.Providers[i].APIKeyEnv != "" {
file.Providers[i].APIKey = os.Getenv(file.Providers[i].APIKeyEnv)
}
}
return file.Providers, nil
}
// Lookup finds a preset or a custom provider by name.
func Lookup(name string, custom []Provider) (Provider, bool) {
for _, p := range custom {
if p.Name == name {
return p, true
}
}
for _, p := range Presets() {
if p.Name == name {
if p.APIKey == "" {
p.APIKey = os.Getenv(p.APIKeyEnv)
}
return p, true
}
}
return Provider{}, false
}
func messageContent(m Message) any {
if len(m.Images) == 0 {
return m.Content
}
parts := []map[string]any{}
if m.Content != "" {
parts = append(parts, map[string]any{"type": "text", "text": m.Content})
}
for _, img := range m.Images {
url := img.URL
if url == "" && img.Data != "" {
mime := img.MIME
if mime == "" {
mime = "image/png"
}
url = "data:" + mime + ";base64," + img.Data
}
parts = append(parts, map[string]any{
"type": "image_url",
"image_url": map[string]any{"url": url},
})
}
return parts
}
// BuildBody returns the JSON body for one native protocol.
func BuildBody(protocol Protocol, model string, messages []Message, tools []ToolDecl, opts StreamOptions) ([]byte, error) {
switch protocol {
case ProtocolOpenAIChat, ProtocolAzure:
body := map[string]any{
"model": model,
"messages": toProviderMessages(messages),
"stream": true,
}
if opts.MaxTokens > 0 {
body["max_tokens"] = opts.MaxTokens
}
if effort := OpenAIReasoningEffort(opts.ThinkingLevel); effort != "" {
body["reasoning_effort"] = effort
}
if len(tools) > 0 {
body["tools"] = toProviderTools(tools)
}
return json.Marshal(body)
case ProtocolOpenAIResponses:
body := map[string]any{
"model": model,
"input": responsesInput(messages),
"stream": true,
}
if opts.MaxTokens > 0 {
body["max_output_tokens"] = opts.MaxTokens
}
if effort := OpenAIReasoningEffort(opts.ThinkingLevel); effort != "" {
body["reasoning"] = map[string]any{"effort": effort}
}
return json.Marshal(body)
case ProtocolAnthropic:
sys, msgs := anthropicMessages(messages)
body := map[string]any{
"model": model,
"messages": msgs,
"max_tokens": opts.MaxTokens,
"stream": true,
}
if body["max_tokens"] == 0 {
body["max_tokens"] = 4096
}
if sys != "" {
body["system"] = sys
}
if budget := AnthropicThinkingBudget(opts.ThinkingLevel); budget > 0 {
body["thinking"] = map[string]any{"type": "enabled", "budget_tokens": budget}
}
return json.Marshal(body)
case ProtocolGoogle:
body := map[string]any{
"contents": googleContents(messages),
}
if budget := GoogleThinkingBudget(opts.ThinkingLevel); budget > 0 {
body["generationConfig"] = map[string]any{
"thinkingConfig": map[string]any{"thinkingBudget": budget},
}
}
return json.Marshal(body)
default:
return nil, fmt.Errorf("llm: unknown protocol %q", protocol)
}
}
func responsesInput(messages []Message) []map[string]any {
out := make([]map[string]any, 0, len(messages))
for _, m := range messages {
role := string(m.Role)
if role == "" {
role = "user"
}
out = append(out, map[string]any{
"role": role,
"content": messageContent(m),
})
}
return out
}
func anthropicMessages(messages []Message) (string, []map[string]any) {
var system strings.Builder
var out []map[string]any
for _, m := range messages {
if m.Role == RoleSystem {
if system.Len() > 0 {
system.WriteString("\n")
}
system.WriteString(m.Content)
continue
}
role := "user"
if m.Role == RoleAssistant {
role = "assistant"
}
content := []map[string]any{}
if m.Content != "" {
content = append(content, map[string]any{"type": "text", "text": m.Content})
}
for _, img := range m.Images {
content = append(content, map[string]any{
"type": "image",
"source": map[string]any{
"type": "base64",
"media_type": firstNonEmpty(img.MIME, "image/png"),
"data": img.Data,
},
})
}
out = append(out, map[string]any{"role": role, "content": content})
}
return system.String(), out
}
func googleContents(messages []Message) []map[string]any {
var out []map[string]any
for _, m := range messages {
if m.Role == RoleSystem {
continue
}
role := "user"
if m.Role == RoleAssistant {
role = "model"
}
parts := []map[string]any{}
if m.Content != "" {
parts = append(parts, map[string]any{"text": m.Content})
}
for _, img := range m.Images {
parts = append(parts, map[string]any{
"inlineData": map[string]any{
"mimeType": firstNonEmpty(img.MIME, "image/png"),
"data": img.Data,
},
})
}
out = append(out, map[string]any{"role": role, "parts": parts})
}
return out
}
func firstNonEmpty(a, b string) string {
if a != "" {
return a
}
return b
}
// ContextOverflow reports provider errors that ask the caller to compact.
func ContextOverflow(message string) bool {
m := strings.ToLower(message)
return strings.Contains(m, "context") && (strings.Contains(m, "overflow") || strings.Contains(m, "too long") || strings.Contains(m, "length"))
}
// WithRetry re-invokes base on a retryable terminal error, up to attempts
// extra tries. onEvent receives auto_retry_start and auto_retry_end.
// Context-overflow errors are returned immediately so the caller can compact.
func WithRetry(base StreamFunc, attempts int, onEvent func(kind, detail string)) StreamFunc {
if attempts < 1 {
attempts = 1
}
return func(ctx context.Context, model string, messages []Message, tools []ToolDecl, opts StreamOptions) <-chan Event {
out := make(chan Event)
go func() {
defer close(out)
var last []Event
for try := 0; try <= attempts; try++ {
if try > 0 {
detail := fmt.Sprintf("attempt %d", try+1)
if onEvent != nil {
onEvent("auto_retry_start", detail)
}
delay := time.Duration(200*(1<<uint(try-1))) * time.Millisecond
select {
case <-ctx.Done():
emitAbort(out, "", ctx.Err())
return
case <-time.After(delay):
}
}
last = nil
for ev := range base(ctx, model, messages, tools, opts) {
last = append(last, ev)
}
overflow, retryable := false, false
for _, ev := range last {
if ev.Kind == KindError && ContextOverflow(ev.ErrorMessage) {
overflow = true
}
if ev.Kind == KindError && ev.Retryable {
retryable = true
}
}
if overflow || !retryable || try == attempts {
if try > 0 && onEvent != nil {
onEvent("auto_retry_end", fmt.Sprintf("attempt %d", try+1))
}
for _, ev := range last {
select {
case <-ctx.Done():
return
case out <- ev:
}
}
return
}
}
}()
return out
}
}
+115
View File
@@ -0,0 +1,115 @@
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)
}
}
+12 -2
View File
@@ -26,12 +26,19 @@ type sseChoice struct {
type sseChunk struct {
Choices []sseChoice `json:"choices"`
Usage *Usage `json:"usage"`
Usage *usageWire `json:"usage"`
Error *struct {
Message string `json:"message"`
} `json:"error"`
}
type usageWire struct {
Usage
Details struct {
Cached int `json:"cached_tokens"`
} `json:"prompt_tokens_details"`
}
// ParseSSE reads an OpenAI-compatible server-sent event body and emits
// deltas plus one terminal message_end or error event.
func ParseSSE(ctx context.Context, r io.Reader) <-chan Event {
@@ -77,7 +84,10 @@ func emitParsed(ctx context.Context, r io.Reader, out chan<- Event) {
return
}
if chunk.Usage != nil {
usage = *chunk.Usage
usage = chunk.Usage.Usage
if chunk.Usage.Details.Cached > 0 {
usage.CacheReadTokens = chunk.Usage.Details.Cached
}
}
if len(chunk.Choices) == 0 {
continue
+13
View File
@@ -48,6 +48,7 @@ type Message struct {
Role Role `json:"role"`
Content string `json:"content,omitempty"`
Name string `json:"name,omitempty"`
Images []Image `json:"images,omitempty"`
ToolCallID string `json:"tool_call_id,omitempty"`
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
StopReason StopReason `json:"stop_reason,omitempty"`
@@ -61,11 +62,20 @@ type ToolDecl struct {
Parameters json.RawMessage `json:"parameters"`
}
// Image is one picture attached to a user message.
type Image struct {
MIME string `json:"mime,omitempty"`
URL string `json:"url,omitempty"`
Data string `json:"data,omitempty"`
}
// Usage counts tokens reported by the provider. Zero means the provider omitted usage.
type Usage struct {
PromptTokens int `json:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens"`
TotalTokens int `json:"total_tokens"`
CacheReadTokens int `json:"cache_read_tokens,omitempty"`
CacheWriteTokens int `json:"cache_write_tokens,omitempty"`
}
// Event is one streamed update. Failures are KindError plus a final assistant
@@ -90,6 +100,9 @@ type StreamOptions struct {
MaxTokens int
APIBase string
APIKey string
// ThinkingLevel is off, minimal, low, medium, high, xhigh, or max.
// Empty leaves the provider default unchanged.
ThinkingLevel string
}
// StreamFunc is the provider boundary. It must return a channel that is closed