diff --git a/runtime/llm/client.go b/runtime/llm/client.go index e07db59..08e1469 100644 --- a/runtime/llm/client.go +++ b/runtime/llm/client.go @@ -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) { diff --git a/runtime/llm/providers.go b/runtime/llm/providers.go index 6d9a048..456f55a 100644 --- a/runtime/llm/providers.go +++ b/runtime/llm/providers.go @@ -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< 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< 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 diff --git a/runtime/llm/providers_test.go b/runtime/llm/providers_test.go index 27a184d..eeb6c56 100644 --- a/runtime/llm/providers_test.go +++ b/runtime/llm/providers_test.go @@ -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") diff --git a/runtime/llm/sse.go b/runtime/llm/sse.go index 0bd3f32..92a367b 100644 --- a/runtime/llm/sse.go +++ b/runtime/llm/sse.go @@ -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 } diff --git a/runtime/llm/sse_test.go b/runtime/llm/sse_test.go index 3a4f801..5b943bc 100644 --- a/runtime/llm/sse_test.go +++ b/runtime/llm/sse_test.go @@ -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 diff --git a/runtime/llm/types.go b/runtime/llm/types.go index abf7035..b000cac 100644 --- a/runtime/llm/types.go +++ b/runtime/llm/types.go @@ -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. diff --git a/runtime/loop.go b/runtime/loop.go index cd9c9e4..74f9530 100644 --- a/runtime/loop.go +++ b/runtime/loop.go @@ -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) diff --git a/runtime/loop_budget_test.go b/runtime/loop_budget_test.go index 5da9702..f2a09a7 100644 --- a/runtime/loop_budget_test.go +++ b/runtime/loop_budget_test.go @@ -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",