agent: retry on transient backend errors and mid-stream drops
Previously only HTTP 429 triggered a retry. Network failures, 5xx responses, and mid-stream drops all hard-failed the turn. - Add TransientError for 5xx and network-level failures; streamRequest wraps client.Do errors and >=500 responses in it - Merge pre-stream and mid-stream retry into one loop (maxTransientRetries=3) covering rate limits, transient errors, and pure stream drops - On stream drop with ctx cancelled: record partial state and return (existing pause behaviour preserved) - On stream drop without cancellation: close any open <think> block, emit retry countdown, restart ChatStream from the same history Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Levi Neely
parent
9fecf5f60f
commit
625257dea1
|
|
@ -11,7 +11,7 @@ import (
|
|||
"ollie/pkg/backend"
|
||||
)
|
||||
|
||||
const maxRateLimitRetries = 3
|
||||
const maxTransientRetries = 3
|
||||
|
||||
type toolExecutor func(ctx context.Context, name string, args json.RawMessage) (string, error)
|
||||
|
||||
|
|
@ -38,100 +38,117 @@ func run(ctx context.Context, cfg agentConfig, state state) error {
|
|||
history = append([]backend.Message{{Role: "system", Content: cfg.preamble}}, history...)
|
||||
}
|
||||
|
||||
// Stream the assistant's response, retrying on HTTP 429.
|
||||
var ch <-chan backend.StreamEvent
|
||||
for attempt := range maxRateLimitRetries + 1 {
|
||||
var err error
|
||||
ch, err = cfg.Backend.ChatStream(ctx, history, cfg.Tools, cfg.GenerationParams)
|
||||
if err == nil {
|
||||
break
|
||||
}
|
||||
if ctx.Err() != nil {
|
||||
return ctx.Err()
|
||||
}
|
||||
var rlErr *backend.RateLimitError
|
||||
if !errors.As(err, &rlErr) || attempt >= maxRateLimitRetries {
|
||||
return fmt.Errorf("step %d: %w", step, err)
|
||||
}
|
||||
wait := rlErr.RetryAfter
|
||||
if wait == 0 {
|
||||
wait = time.Duration(5<<attempt) * time.Second
|
||||
}
|
||||
if err := retryCountdown(ctx, cfg, wait); err != nil {
|
||||
return fmt.Errorf("step %d: %w", step, err)
|
||||
}
|
||||
}
|
||||
|
||||
// Stream the assistant's response, retrying on rate limits, transient
|
||||
// backend errors (5xx, network), and mid-stream drops.
|
||||
var content strings.Builder
|
||||
var toolCalls []backend.ToolCall
|
||||
var stopReason string
|
||||
var done bool
|
||||
var hadReasoning bool
|
||||
|
||||
for ev := range ch {
|
||||
if ev.Reasoning != "" {
|
||||
if !hadReasoning {
|
||||
emit(cfg, Event{Role: "reasoning", Content: "<think>\n"})
|
||||
hadReasoning = true
|
||||
for attempt := range maxTransientRetries + 1 {
|
||||
content.Reset()
|
||||
toolCalls = nil
|
||||
|
||||
ch, err := cfg.Backend.ChatStream(ctx, history, cfg.Tools, cfg.GenerationParams)
|
||||
if err != nil {
|
||||
if ctx.Err() != nil {
|
||||
return ctx.Err()
|
||||
}
|
||||
emit(cfg, Event{Role: "reasoning", Content: ev.Reasoning})
|
||||
wait, retryable := transientWait(err, attempt)
|
||||
if !retryable || attempt >= maxTransientRetries {
|
||||
return fmt.Errorf("step %d: %w", step, err)
|
||||
}
|
||||
if err := retryCountdown(ctx, cfg, wait); err != nil {
|
||||
return fmt.Errorf("step %d: %w", step, err)
|
||||
}
|
||||
continue
|
||||
}
|
||||
if ev.Content != "" {
|
||||
if hadReasoning {
|
||||
emit(cfg, Event{Role: "reasoning", Content: "\n</think>\n"})
|
||||
hadReasoning = false
|
||||
}
|
||||
content.WriteString(ev.Content)
|
||||
emit(cfg, Event{Role: "assistant", Content: ev.Content})
|
||||
}
|
||||
toolCalls = append(toolCalls, ev.ToolCalls...)
|
||||
if ev.Done {
|
||||
if hadReasoning {
|
||||
emit(cfg, Event{Role: "reasoning", Content: "\n</think>\n"})
|
||||
hadReasoning = false
|
||||
}
|
||||
stopReason = ev.StopReason
|
||||
done = true
|
||||
if ev.Usage.InputTokens > 0 || ev.Usage.OutputTokens > 0 {
|
||||
emit(cfg, Event{
|
||||
Role: "usage",
|
||||
Content: fmt.Sprintf("%d %d 0", ev.Usage.InputTokens, ev.Usage.OutputTokens),
|
||||
})
|
||||
} else {
|
||||
// Backend didn't report usage; estimate from content.
|
||||
inChars := 0
|
||||
for _, m := range history {
|
||||
inChars += len(m.Content)
|
||||
for _, tc := range m.ToolCalls {
|
||||
inChars += len(tc.Name) + len(tc.Arguments)
|
||||
}
|
||||
|
||||
var done bool
|
||||
var hadReasoning bool
|
||||
for ev := range ch {
|
||||
if ev.Reasoning != "" {
|
||||
if !hadReasoning {
|
||||
emit(cfg, Event{Role: "reasoning", Content: "<think>\n"})
|
||||
hadReasoning = true
|
||||
}
|
||||
emit(cfg, Event{
|
||||
Role: "usage",
|
||||
Content: fmt.Sprintf("%d %d 1", inChars/4, content.Len()/4),
|
||||
})
|
||||
emit(cfg, Event{Role: "reasoning", Content: ev.Reasoning})
|
||||
}
|
||||
if ev.Content != "" {
|
||||
if hadReasoning {
|
||||
emit(cfg, Event{Role: "reasoning", Content: "\n</think>\n"})
|
||||
hadReasoning = false
|
||||
}
|
||||
content.WriteString(ev.Content)
|
||||
emit(cfg, Event{Role: "assistant", Content: ev.Content})
|
||||
}
|
||||
toolCalls = append(toolCalls, ev.ToolCalls...)
|
||||
if ev.Done {
|
||||
if hadReasoning {
|
||||
emit(cfg, Event{Role: "reasoning", Content: "\n</think>\n"})
|
||||
hadReasoning = false
|
||||
}
|
||||
stopReason = ev.StopReason
|
||||
done = true
|
||||
if ev.Usage.InputTokens > 0 || ev.Usage.OutputTokens > 0 {
|
||||
emit(cfg, Event{
|
||||
Role: "usage",
|
||||
Content: fmt.Sprintf("%d %d 0", ev.Usage.InputTokens, ev.Usage.OutputTokens),
|
||||
})
|
||||
} else {
|
||||
// Backend didn't report usage; estimate from content.
|
||||
inChars := 0
|
||||
for _, m := range history {
|
||||
inChars += len(m.Content)
|
||||
for _, tc := range m.ToolCalls {
|
||||
inChars += len(tc.Name) + len(tc.Arguments)
|
||||
}
|
||||
}
|
||||
emit(cfg, Event{
|
||||
Role: "usage",
|
||||
Content: fmt.Sprintf("%d %d 1", inChars/4, content.Len()/4),
|
||||
})
|
||||
}
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if done {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if !done {
|
||||
// Stream interrupted. Record whatever we got.
|
||||
msg := backend.Message{Role: "assistant", Content: content.String(), ToolCalls: toolCalls}
|
||||
var results []toolResult
|
||||
for _, tc := range toolCalls {
|
||||
results = append(results, toolResult{
|
||||
ToolCallID: tc.ID,
|
||||
Name: tc.Name,
|
||||
Content: `{"status":"cancelled","error":"stream interrupted"}`,
|
||||
IsError: true,
|
||||
})
|
||||
if ctx.Err() != nil {
|
||||
// User pause: record partial state and return.
|
||||
if hadReasoning {
|
||||
emit(cfg, Event{Role: "reasoning", Content: "\n</think>\n"})
|
||||
}
|
||||
msg := backend.Message{Role: "assistant", Content: content.String(), ToolCalls: toolCalls}
|
||||
var results []toolResult
|
||||
for _, tc := range toolCalls {
|
||||
results = append(results, toolResult{
|
||||
ToolCallID: tc.ID,
|
||||
Name: tc.Name,
|
||||
Content: `{"status":"cancelled","error":"stream interrupted"}`,
|
||||
IsError: true,
|
||||
})
|
||||
}
|
||||
state.update(msg, results)
|
||||
if cfg.SaveSession != nil {
|
||||
cfg.SaveSession()
|
||||
}
|
||||
return ctx.Err()
|
||||
}
|
||||
state.update(msg, results)
|
||||
if cfg.SaveSession != nil {
|
||||
cfg.SaveSession()
|
||||
|
||||
// Pure stream drop — retry if attempts remain.
|
||||
if attempt >= maxTransientRetries {
|
||||
return fmt.Errorf("step %d: stream dropped (no more retries)", step)
|
||||
}
|
||||
if hadReasoning {
|
||||
emit(cfg, Event{Role: "reasoning", Content: "\n</think>\n"})
|
||||
}
|
||||
wait := time.Duration(2<<attempt) * time.Second
|
||||
if err := retryCountdown(ctx, cfg, wait); err != nil {
|
||||
return fmt.Errorf("step %d: %w", step, err)
|
||||
}
|
||||
return fmt.Errorf("step %d: stream interrupted: %w", step, ctx.Err())
|
||||
}
|
||||
|
||||
switch stopReason {
|
||||
|
|
@ -285,6 +302,24 @@ func emit(cfg agentConfig, msg Event) {
|
|||
}
|
||||
}
|
||||
|
||||
// transientWait returns the retry wait for a retryable error and whether it is
|
||||
// retryable. Rate limits use longer waits; transient/network errors use shorter.
|
||||
func transientWait(err error, attempt int) (time.Duration, bool) {
|
||||
var rlErr *backend.RateLimitError
|
||||
if errors.As(err, &rlErr) {
|
||||
wait := rlErr.RetryAfter
|
||||
if wait == 0 {
|
||||
wait = time.Duration(5<<attempt) * time.Second
|
||||
}
|
||||
return wait, true
|
||||
}
|
||||
var tErr *backend.TransientError
|
||||
if errors.As(err, &tErr) {
|
||||
return time.Duration(2<<attempt) * time.Second, true
|
||||
}
|
||||
return 0, false
|
||||
}
|
||||
|
||||
func retryCountdown(ctx context.Context, cfg agentConfig, wait time.Duration) error {
|
||||
deadline := time.Now().Add(wait)
|
||||
for {
|
||||
|
|
|
|||
|
|
@ -68,6 +68,14 @@ func (e *RateLimitError) Error() string {
|
|||
return fmt.Sprintf("rate limited: %s", e.Message)
|
||||
}
|
||||
|
||||
// TransientError is returned for retryable backend failures: 5xx HTTP responses
|
||||
// and network-level errors (connection reset, DNS failure, etc.).
|
||||
type TransientError struct {
|
||||
Message string
|
||||
}
|
||||
|
||||
func (e *TransientError) Error() string { return e.Message }
|
||||
|
||||
// GenerationParams controls sampling behaviour for a single ChatStream call.
|
||||
// Zero values mean "use the API default".
|
||||
type GenerationParams struct {
|
||||
|
|
@ -98,19 +106,26 @@ type Backend interface {
|
|||
Models(ctx context.Context) []string
|
||||
}
|
||||
|
||||
// streamRequest executes req via client, handles 429/non-200 errors, and
|
||||
// spawns a goroutine that feeds resp.Body through parseFn into the returned
|
||||
// channel. label is used in error messages (e.g. "openai", "ollama").
|
||||
// streamRequest executes req via client, handles status errors, and spawns a
|
||||
// goroutine that feeds resp.Body through parseFn into the returned channel.
|
||||
// label is used in error messages (e.g. "openai", "ollama").
|
||||
// Returns *RateLimitError for 429, *TransientError for 5xx and network errors,
|
||||
// and a plain error for other non-200 responses.
|
||||
func streamRequest(client *http.Client, req *http.Request, label string, parseFn func(io.Reader, chan<- StreamEvent)) (<-chan StreamEvent, error) {
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
return nil, &TransientError{Message: err.Error()}
|
||||
}
|
||||
if resp.StatusCode == http.StatusTooManyRequests {
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
resp.Body.Close()
|
||||
return nil, &RateLimitError{RetryAfter: parseRetryAfter(resp.Header.Get("Retry-After")), Message: string(body)}
|
||||
}
|
||||
if resp.StatusCode >= 500 {
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
resp.Body.Close()
|
||||
return nil, &TransientError{Message: fmt.Sprintf("%s HTTP %d: %s", label, resp.StatusCode, string(body))}
|
||||
}
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
resp.Body.Close()
|
||||
|
|
|
|||
|
|
@ -245,6 +245,21 @@ func checkErrorHTTP(t *testing.T, b backend.Backend) {
|
|||
}
|
||||
}
|
||||
|
||||
// checkErrorTransient verifies that 5xx responses return *TransientError.
|
||||
func checkErrorTransient(t *testing.T, b backend.Backend) {
|
||||
t.Helper()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
_, err := b.ChatStream(ctx, []backend.Message{{Role: "user", Content: "hi"}}, nil, backend.GenerationParams{})
|
||||
if err == nil {
|
||||
t.Fatal("expected error for 5xx")
|
||||
}
|
||||
if _, ok := err.(*backend.TransientError); !ok {
|
||||
t.Errorf("error type = %T; want *TransientError", err)
|
||||
}
|
||||
}
|
||||
|
||||
// checkErrorRateLimit verifies that 429 returns *RateLimitError.
|
||||
func checkErrorRateLimit(t *testing.T, b backend.Backend) {
|
||||
t.Helper()
|
||||
|
|
@ -434,7 +449,12 @@ func TestOpenAIErrorHTTP(t *testing.T) {
|
|||
t.Run(fmt.Sprintf("%d", code), func(t *testing.T) {
|
||||
srv := openAIErrorServer(code)
|
||||
defer srv.Close()
|
||||
checkErrorHTTP(t, mustNewOpenAI(t, "openai", srv.URL, "k"))
|
||||
b := mustNewOpenAI(t, "openai", srv.URL, "k")
|
||||
if code >= 500 {
|
||||
checkErrorTransient(t, b)
|
||||
} else {
|
||||
checkErrorHTTP(t, b)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
@ -494,7 +514,12 @@ func TestOllamaErrorHTTP(t *testing.T) {
|
|||
t.Run(fmt.Sprintf("%d", code), func(t *testing.T) {
|
||||
srv := ollamaErrorServer(code)
|
||||
defer srv.Close()
|
||||
checkErrorHTTP(t, mustNewOllama(t, srv.URL))
|
||||
b := mustNewOllama(t, srv.URL)
|
||||
if code >= 500 {
|
||||
checkErrorTransient(t, b)
|
||||
} else {
|
||||
checkErrorHTTP(t, b)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
@ -547,7 +572,11 @@ func TestAnthropicErrorHTTP(t *testing.T) {
|
|||
defer srv.Close()
|
||||
b := mustNewAnthropic(t, "k")
|
||||
b.BaseURL, _ = url.Parse(srv.URL)
|
||||
checkErrorHTTP(t, b)
|
||||
if code >= 500 {
|
||||
checkErrorTransient(t, b)
|
||||
} else {
|
||||
checkErrorHTTP(t, b)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Reference in New Issue