Retry model streams only before the first delta
ci / vet-and-test (pull_request) Successful in 47s
ci / docker-checkout-hostgateway (pull_request) Successful in 1s

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>
This commit is contained in:
2026-10-03 15:52:08 +00:00
co-authored by Cursor
parent e51f2e6d00
commit 9f5751c205
8 changed files with 401 additions and 35 deletions
+13 -2
View File
@@ -89,8 +89,19 @@ func (c *Client) run(ctx context.Context, model string, messages []Message, tool
defer resp.Body.Close()
if resp.StatusCode >= 400 {
slurp, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
retryable := resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode >= 500
emitError(out, "", retryable, fmt.Sprintf("provider status %d: %s", resp.StatusCode, strings.TrimSpace(string(slurp))))
msg := fmt.Sprintf("provider status %d: %s", resp.StatusCode, strings.TrimSpace(string(slurp)))
ev := Event{
Kind: KindError,
Retryable: StatusRetryable(resp.StatusCode) && !ContextOverflow(msg),
ErrorMessage: msg,
Err: &StreamError{Message: msg, Retryable: StatusRetryable(resp.StatusCode) && !ContextOverflow(msg)},
Message: Message{Role: RoleAssistant, StopReason: StopError},
}
if d, ok := ParseRetryAfter(resp.Header.Get("Retry-After"), time.Now()); ok {
ev.HasRetryAfter = true
ev.RetryAfter = d
}
out <- ev
return
}
for ev := range ParseSSE(ctx, resp.Body) {
+151 -23
View File
@@ -4,7 +4,10 @@ import (
"context"
"encoding/json"
"fmt"
"math/rand"
"net/http"
"os"
"strconv"
"strings"
"time"
)
@@ -382,31 +385,136 @@ func firstNonEmpty(a, b string) string {
return b
}
const (
retryBaseDelay = 200 * time.Millisecond
maxRetryAfter = 30 * time.Second
)
// overflowNeedles are substrings that mean the prompt does not fit.
// Matching is case-insensitive. A bare "context" is not enough.
var overflowNeedles = []string{
"status 413",
"payload too large",
"context_length_exceeded",
"context length",
"maximum context",
"context window",
"prompt is too long",
"exceeds the maximum",
"exceeds max",
"too many tokens",
"input is too long",
"input length",
"tokens exceed",
"上下文",
"超出最大",
"输入过长",
"提示过长",
"请求过长",
"请求体过大",
}
// 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"))
for _, needle := range overflowNeedles {
if strings.Contains(m, strings.ToLower(needle)) {
return true
}
}
return false
}
// CapacityLimited reports an in-stream overload that is worth retrying
// before any delta has been forwarded.
func CapacityLimited(message string) bool {
m := strings.ToLower(message)
return strings.Contains(m, "overloaded") || strings.Contains(m, "at capacity") || strings.Contains(m, "too many requests") || strings.Contains(m, "rate limit")
}
// StatusRetryable is true for timeouts, rate limits, server errors, and 529.
func StatusRetryable(code int) bool {
return code == http.StatusRequestTimeout || code == http.StatusTooManyRequests || code == 529 || code >= 500
}
// ParseRetryAfter reads a delay in seconds or an HTTP date.
// The returned duration is capped. ok is false when the header is missing or garbage.
func ParseRetryAfter(raw string, now time.Time) (time.Duration, bool) {
raw = strings.TrimSpace(raw)
if raw == "" {
return 0, false
}
if secs, err := strconv.Atoi(raw); err == nil {
if secs < 0 {
return 0, false
}
return capRetryAfter(time.Duration(secs) * time.Second), true
}
when, err := http.ParseTime(raw)
if err != nil {
return 0, false
}
d := when.Sub(now)
if d < 0 {
d = 0
}
return capRetryAfter(d), true
}
func capRetryAfter(d time.Duration) time.Duration {
if d > maxRetryAfter {
return maxRetryAfter
}
return d
}
func retryBackoff(extraTry int) time.Duration {
if extraTry < 1 {
extraTry = 1
}
shift := extraTry - 1
if shift > 4 {
shift = 4
}
delay := retryBaseDelay * time.Duration(1<<uint(shift))
jitter := time.Duration(rand.Int63n(int64(delay/2) + 1))
return delay + jitter
}
func forwardedDelta(ev Event) bool {
switch ev.Kind {
case KindTextDelta, KindReasoningDelta, KindToolCallDelta:
return true
default:
return false
}
}
// 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.
// extra tries. Events are forwarded as they arrive. A retry happens only when
// no delta has been forwarded yet. onEvent receives auto_retry_start and
// auto_retry_end. Context-overflow errors are returned immediately.
// attempts <= 0 disables extra tries.
func WithRetry(base StreamFunc, attempts int, onEvent func(kind, detail string)) StreamFunc {
if attempts < 1 {
attempts = 1
if attempts < 0 {
attempts = 0
}
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
headerWait := time.Duration(0)
useHeader := false
for try := 0; try <= attempts; try++ {
if try > 0 {
delay := retryBackoff(try)
if useHeader {
delay = headerWait
}
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())
@@ -414,32 +522,52 @@ func WithRetry(base StreamFunc, attempts int, onEvent func(kind, detail string))
case <-time.After(delay):
}
}
last = nil
useHeader = false
forwarded := false
var terminal *Event
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 terminal != nil {
continue
}
if ev.Kind == KindError && ev.Retryable {
retryable = true
if ev.Kind == KindError {
cp := ev
terminal = &cp
continue
}
if forwardedDelta(ev) {
forwarded = true
}
select {
case <-ctx.Done():
return
case out <- ev:
}
if ev.Kind == KindMessageEnd {
if try > 0 && onEvent != nil {
onEvent("auto_retry_end", fmt.Sprintf("attempt %d", try+1))
}
return
}
}
if overflow || !retryable || try == attempts {
if terminal == nil {
return
}
overflow := ContextOverflow(terminal.ErrorMessage)
if forwarded || overflow || !terminal.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:
}
select {
case <-ctx.Done():
return
case out <- *terminal:
}
return
}
if terminal.HasRetryAfter {
useHeader = true
headerWait = terminal.RetryAfter
}
}
}()
return out
+70
View File
@@ -9,6 +9,7 @@ import (
"path/filepath"
"strings"
"testing"
"time"
)
func TestParseModelRefAndThinking(t *testing.T) {
@@ -141,6 +142,75 @@ func TestRetryThenSuccess(t *testing.T) {
}
}
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")
+29 -2
View File
@@ -56,6 +56,7 @@ func emitParsed(ctx context.Context, r io.Reader, out chan<- Event) {
order := []int{}
var usage Usage
finished := ""
sawDone := false
sc := bufio.NewScanner(r)
sc.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
reader := &sseEventReader{sc: sc}
@@ -69,6 +70,7 @@ func emitParsed(ctx context.Context, r io.Reader, out chan<- Event) {
return
}
if data == "[DONE]" {
sawDone = true
break
}
var chunk sseChunk
@@ -77,7 +79,9 @@ func emitParsed(ctx context.Context, r io.Reader, out chan<- Event) {
return
}
if chunk.Error != nil && chunk.Error.Message != "" {
emitError(out, content.String(), false, chunk.Error.Message)
msg := chunk.Error.Message
retryable := CapacityLimited(msg) && !ContextOverflow(msg)
emitError(out, content.String(), retryable, msg)
return
}
if chunk.Usage != nil {
@@ -157,6 +161,10 @@ func emitParsed(ctx context.Context, r io.Reader, out chan<- Event) {
emitAbort(out, content.String(), ctx.Err())
return
}
// length、空结束原因,或没有 [DONE] 的中断,都不把半截工具调用交给执行层。
if !toolFinishExecutable(finished, sawDone) {
order = nil
}
msg := Message{Role: RoleAssistant, Content: content.String(), StopReason: mapFinish(finished, len(order) > 0)}
for _, idx := range order {
msg.ToolCalls = append(msg.ToolCalls, *calls[idx])
@@ -222,13 +230,32 @@ func (r *sseEventReader) next() (string, bool) {
return "", false
}
// toolFinishExecutable 只在供应商明确以工具调用或正常 stop 结束时执行。
// length、缺失的结束原因,以及没有 [DONE] 的断开,都视为未完成。
func toolFinishExecutable(reason string, sawDone bool) bool {
switch reason {
case "tool_calls", "function_call":
return true
case "stop":
return true
case "length":
return false
case "":
return false
default:
return sawDone && reason != "error" && reason != "content_filter"
}
}
func mapFinish(reason string, hasTools bool) StopReason {
switch reason {
case "tool_calls", "function_call":
return StopToolCalls
case "content_filter", "error":
return StopError
case "stop", "length", "":
case "length":
return StopEnd
case "stop", "":
if hasTools {
return StopToolCalls
}
+45
View File
@@ -93,6 +93,51 @@ func TestParseSSEToolCallIncrements(t *testing.T) {
}
}
func TestParseSSELengthDropsUnfinishedToolCall(t *testing.T) {
body := strings.Join([]string{
`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"read_file","arguments":"{\"path\":"}}]}}]}`,
`data: {"choices":[{"delta":{},"finish_reason":"length"}]}`,
`data: [DONE]`,
"",
}, "\n")
var end Event
for ev := range ParseSSE(context.Background(), strings.NewReader(body)) {
if ev.Kind == KindMessageEnd {
end = ev
}
}
if end.Message.StopReason != StopEnd || len(end.Message.ToolCalls) != 0 {
t.Fatalf("truncated tool call must not be executable: %#v", end.Message)
}
if end.FinishReason != "length" {
t.Fatalf("finish %q", end.FinishReason)
}
}
func TestParseSSEUncleanEndDropsToolCall(t *testing.T) {
body := `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"read_file","arguments":"{\"path\":"}}]}}]}`
var end Event
for ev := range ParseSSE(context.Background(), strings.NewReader(body)) {
if ev.Kind == KindMessageEnd {
end = ev
}
}
if end.Kind != KindMessageEnd || end.Message.StopReason != StopEnd || len(end.Message.ToolCalls) != 0 {
t.Fatalf("unclean stream must not execute tools: %#v", end)
}
}
func TestParseSSECapacityErrorIsRetryable(t *testing.T) {
body := "data: {\"error\":{\"message\":\"server is at capacity\"}}\n\n"
var got Event
for ev := range ParseSSE(context.Background(), strings.NewReader(body)) {
got = ev
}
if got.Kind != KindError || !got.Retryable || !CapacityLimited(got.ErrorMessage) {
t.Fatalf("event %#v", got)
}
}
func TestParseSSEMalformedIsError(t *testing.T) {
body := "data: {not-json}\n\n"
var got Event
+17 -6
View File
@@ -3,6 +3,7 @@ package llm
import (
"context"
"encoding/json"
"time"
)
// Role is a transcript role carried into the provider request.
@@ -69,13 +70,19 @@ type Image struct {
Data string `json:"data,omitempty"`
}
// Usage counts tokens reported by the provider. Zero means the provider omitted usage.
// Usage counts tokens. On one provider response, PromptTokens, CompletionTokens
// and TotalTokens are that request. On a loop Result they are the billing sum
// across requests; Last* is only the latest request, for context observation.
// 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"`
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"`
LastPromptTokens int `json:"last_prompt_tokens,omitempty"`
LastCompletionTokens int `json:"last_completion_tokens,omitempty"`
LastTotalTokens int `json:"last_total_tokens,omitempty"`
}
// Event is one streamed update. Failures are KindError plus a final assistant
@@ -96,6 +103,10 @@ type Event struct {
Err error
Retryable bool
ErrorMessage string
// RetryAfter is the provider's requested wait. HasRetryAfter distinguishes
// a real zero from a missing or unparseable header.
RetryAfter time.Duration
HasRetryAfter bool
}
// StreamOptions tunes one provider request.
+25 -2
View File
@@ -192,6 +192,9 @@ func Run(ctx context.Context, cfg Config, user llm.Message) (Result, error) {
usage.PromptTokens += end.Usage.PromptTokens
usage.CompletionTokens += end.Usage.CompletionTokens
usage.TotalTokens += end.Usage.TotalTokens
usage.LastPromptTokens = end.Usage.PromptTokens
usage.LastCompletionTokens = end.Usage.CompletionTokens
usage.LastTotalTokens = end.Usage.TotalTokens
if err := persist(cfg, last); err != nil {
return Result{}, err
}
@@ -464,10 +467,30 @@ func cancelledToolMessage(call llm.ToolCall, err error) llm.Message {
}
}
func rejectIncompleteToolCall(call llm.ToolCall) string {
if strings.TrimSpace(call.Name) == "" {
return "工具调用未执行:名称不完整"
}
args := strings.TrimSpace(call.Arguments)
if !json.Valid([]byte(args)) || !strings.HasPrefix(args, "{") {
return "工具调用未执行:参数不是完整的 JSON 对象"
}
return ""
}
func runOne(ctx context.Context, cfg Config, call llm.ToolCall) llm.Message {
if err := ctx.Err(); err != nil {
return cancelledToolMessage(call, err)
}
if why := rejectIncompleteToolCall(call); why != "" {
msg := llm.Message{Role: llm.RoleTool, Name: call.Name, ToolCallID: call.ID, Content: why, IsError: true}
emit(cfg, Event{Type: "tool_execution_start", Tool: call.Name, Detail: call.Arguments})
emit(cfg, Event{Type: "tool_execution_end", Tool: call.Name, Text: why})
if cfg.AfterTool != nil {
cfg.AfterTool(call, msg)
}
return msg
}
emit(cfg, Event{Type: "tool_execution_start", Tool: call.Name, Detail: call.Arguments})
if cfg.BeforeTool != nil {
if err := cfg.BeforeTool(call); err != nil {
@@ -616,8 +639,8 @@ func budgetExceeded(b *Budget, usage llm.Usage, toolCalls int, now time.Time) st
if !b.Deadline.IsZero() && now.After(b.Deadline) {
return fmt.Sprintf("budget_exceeded: deadline %s passed", b.Deadline.Format(time.RFC3339))
}
if b.MaxPromptTokens > 0 && usage.PromptTokens > b.MaxPromptTokens {
return fmt.Sprintf("budget_exceeded: prompt tokens %d > max %d", usage.PromptTokens, b.MaxPromptTokens)
if b.MaxPromptTokens > 0 && usage.LastPromptTokens > b.MaxPromptTokens {
return fmt.Sprintf("budget_exceeded: prompt tokens %d > max %d", usage.LastPromptTokens, b.MaxPromptTokens)
}
if b.MaxCompletionTokens > 0 && usage.CompletionTokens > b.MaxCompletionTokens {
return fmt.Sprintf("budget_exceeded: completion tokens %d > max %d", usage.CompletionTokens, b.MaxCompletionTokens)
+51
View File
@@ -194,6 +194,57 @@ func TestLoopBudgetPromptTokens(t *testing.T) {
}
}
func TestLoopUsageSplitsLastRequestFromCumulative(t *testing.T) {
res, err := Run(context.Background(), Config{
Model: "fake",
Tools: tools.WorkspaceTools(t.TempDir()),
Budget: &Budget{MaxPromptTokens: 10},
Stream: usageScript(llm.Usage{PromptTokens: 6, CompletionTokens: 1, TotalTokens: 7},
llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{
ID: "c1", Name: "list_dir", Arguments: `{}`,
}}},
llm.Message{Role: llm.RoleAssistant, Content: "ok", StopReason: llm.StopEnd},
),
}, llm.Message{Role: llm.RoleUser, Content: "go"})
if err != nil {
t.Fatal(err)
}
if res.Err != nil || res.Text != "ok" {
t.Fatalf("last request stays under the window: err=%v text=%q", res.Err, res.Text)
}
if res.Usage.PromptTokens != 12 || res.Usage.LastPromptTokens != 6 {
t.Fatalf("billing sum %d last %d", res.Usage.PromptTokens, res.Usage.LastPromptTokens)
}
}
func TestLoopRejectsIncompleteToolCall(t *testing.T) {
tool := &countingTool{name: "list_dir"}
reg := tools.NewRegistry(t.TempDir())
reg.Register(tool)
res, err := Run(context.Background(), Config{
Model: "fake",
Tools: reg,
Stream: script(llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{
ID: "c1", Name: "list_dir", Arguments: `{"path":`,
}}}, llm.Message{Role: llm.RoleAssistant, Content: "stopped", StopReason: llm.StopEnd}),
}, llm.Message{Role: llm.RoleUser, Content: "go"})
if err != nil {
t.Fatal(err)
}
if tool.calls != 0 {
t.Fatalf("incomplete call executed %d times", tool.calls)
}
var paired bool
for _, m := range res.Messages {
if m.Role == llm.RoleTool && m.ToolCallID == "c1" && m.IsError {
paired = true
}
}
if !paired {
t.Fatal("expected a paired error result")
}
}
func TestLoopBudgetDeadline(t *testing.T) {
res, err := Run(context.Background(), Config{
Model: "fake",