Retry transient provider stream errors

commit: 16b95cb97448b69ec861edacf96a335931bf5313
author: patriceckhart · 4mo ago · 2026-05-30 13:25 UTC
parent: dfd2501
signature: unsigned
refs: tag: v0.2.5

2 files · +213 -7

M packages/core/agent.go +102 -7


@@ -46,6 +46,13 @@ type Agent struct {
46	46		// model can still see what it said in subsequent turns).
47	47		BeforeAssistantMessage func(text string) (allowed bool, reason, replacement string)
48	48
	49	+	// MaxRetries controls agent-level retries for transient provider
	50	+	// failures that arrive after the HTTP stream opens (for example
	51	+	// Anthropic overloaded_error). Zero disables this retry layer.
	52	+	// RetryBaseDelay is doubled for each attempt; zero uses 2s.
	53	+	MaxRetries int
	54	+	RetryBaseDelay time.Duration
	55	+
49	56		// OnEvent, if set, mirrors every AgentEvent the loop emits to
50	57		// this callback in addition to the per-Prompt sink. Used by the
51	58		// extension manager to fan events out to subscribed extensions

@@ -94,11 +101,13 @@ type Agent struct {
94	101		// NewAgent returns an Agent with sensible defaults.
95	102		func NewAgent(client provider.Client, model, system string, tools Registry) *Agent {
96	103		return &Agent{
97		-	Client: client,
98		-	Model: model,
99		-	System: system,
100		-	Tools: tools,
101		-	MaxSteps: 0, // 0 = unlimited
	104	+	Client: client,
	105	+	Model: model,
	106	+	System: system,
	107	+	Tools: tools,
	108	+	MaxSteps: 0, // 0 = unlimited
	109	+	MaxRetries: 3,
	110	+	RetryBaseDelay: 2 * time.Second,
102	111		}
103	112		}
104	113

@@ -337,8 +346,23 @@ func (a *Agent) runLoop(ctx context.Context, sink func(AgentEvent)) error {
337	346		return nil
338	347		}
339	348		}
340		-	stop, assistantMsg, err := a.oneTurn(ctx, sink)
341		-	sink(EvTurnEnd{Stop: stop, Err: err})
	349	+
	350	+	var (
	351	+	stop provider.StopReason
	352	+	assistantMsg provider.Message
	353	+	err error
	354	+	)
	355	+	for attempt := 0; ; attempt++ {
	356	+	stop, assistantMsg, err = a.oneTurn(ctx, sink)
	357	+	sink(EvTurnEnd{Stop: stop, Err: err})
	358	+	if err == nil \|\| !a.canRetryError(err, attempt) {
	359	+	break
	360	+	}
	361	+	a.dropLastAssistantMessage()
	362	+	if sleepErr := sleepRetry(ctx, a.retryDelay(attempt)); sleepErr != nil {
	363	+	return sleepErr
	364	+	}
	365	+	}
342	366		if err != nil {
343	367		return err
344	368		}

@@ -403,6 +427,77 @@ func (a *Agent) runLoop(ctx context.Context, sink func(AgentEvent)) error {
403	427		return nil
404	428		}
405	429
	430	+	func (a *Agent) canRetryError(err error, attempt int) bool {
	431	+	if err == nil \|\| a.MaxRetries <= 0 \|\| attempt >= a.MaxRetries {
	432	+	return false
	433	+	}
	434	+	if errors.Is(err, context.Canceled) \|\| errors.Is(err, context.DeadlineExceeded) {
	435	+	return false
	436	+	}
	437	+	msg := strings.ToLower(err.Error())
	438	+	if msg == "" \|\| isNonRetryableProviderLimit(msg) {
	439	+	return false
	440	+	}
	441	+	needles := []string{
	442	+	"overloaded", "provider returned error", "rate limit", "ratelimit", "too many requests",
	443	+	"429", "http 429", "500", "http 500", "502", "http 502", "503", "http 503", "504", "http 504",
	444	+	"service unavailable", "server error", "internal error", "network error", "connection error",
	445	+	"connection refused", "connection lost", "fetch failed", "upstream connect", "reset before headers",
	446	+	"socket hang up", "ended without", "stream ended before", "did not get a response", "timed out",
	447	+	"timeout", "terminated", "unexpected eof", "transport failure",
	448	+	}
	449	+	for _, needle := range needles {
	450	+	if strings.Contains(msg, needle) {
	451	+	return true
	452	+	}
	453	+	}
	454	+	return false
	455	+	}
	456	+
	457	+	func isNonRetryableProviderLimit(msg string) bool {
	458	+	needles := []string{
	459	+	"usage limit", "monthly usage limit", "freeusagelimit", "gousagelimit",
	460	+	"available balance", "insufficient_quota", "out of budget", "quota exceeded", "billing",
	461	+	}
	462	+	for _, needle := range needles {
	463	+	if strings.Contains(msg, needle) {
	464	+	return true
	465	+	}
	466	+	}
	467	+	return false
	468	+	}
	469	+
	470	+	func (a *Agent) retryDelay(attempt int) time.Duration {
	471	+	base := a.RetryBaseDelay
	472	+	if base <= 0 {
	473	+	base = 2 * time.Second
	474	+	}
	475	+	return base * time.Duration(1<<attempt)
	476	+	}
	477	+
	478	+	func sleepRetry(ctx context.Context, d time.Duration) error {
	479	+	if d <= 0 {
	480	+	return ctx.Err()
	481	+	}
	482	+	t := time.NewTimer(d)
	483	+	defer t.Stop()
	484	+	select {
	485	+	case <-ctx.Done():
	486	+	return ctx.Err()
	487	+	case <-t.C:
	488	+	return nil
	489	+	}
	490	+	}
	491	+
	492	+	func (a *Agent) dropLastAssistantMessage() {
	493	+	a.mu.Lock()
	494	+	defer a.mu.Unlock()
	495	+	if n := len(a.messages); n > 0 && a.messages[n-1].Role == provider.RoleAssistant {
	496	+	a.messages = a.messages[:n-1]
	497	+	a.rev++
	498	+	}
	499	+	}
	500	+
406	501		// oneTurn calls the LLM once, forwards events, returns the stop reason
407	502		// and the assembled assistant message (already appended to the transcript).
408	503		func (a *Agent) oneTurn(ctx context.Context, sink func(AgentEvent)) (provider.StopReason, provider.Message, error) {


@@ -46,6 +46,13 @@ type Agent struct {
46		// model can still see what it said in subsequent turns).	46		// model can still see what it said in subsequent turns).
47		BeforeAssistantMessage func(text string) (allowed bool, reason, replacement string)	47		BeforeAssistantMessage func(text string) (allowed bool, reason, replacement string)
48			48
			49	+	// MaxRetries controls agent-level retries for transient provider
			50	+	// failures that arrive after the HTTP stream opens (for example
			51	+	// Anthropic overloaded_error). Zero disables this retry layer.
			52	+	// RetryBaseDelay is doubled for each attempt; zero uses 2s.
			53	+	MaxRetries int
			54	+	RetryBaseDelay time.Duration
			55	+
49		// OnEvent, if set, mirrors every AgentEvent the loop emits to	56		// OnEvent, if set, mirrors every AgentEvent the loop emits to
50		// this callback in addition to the per-Prompt sink. Used by the	57		// this callback in addition to the per-Prompt sink. Used by the
51		// extension manager to fan events out to subscribed extensions	58		// extension manager to fan events out to subscribed extensions

@@ -94,11 +101,13 @@ type Agent struct {
94		// NewAgent returns an Agent with sensible defaults.	101		// NewAgent returns an Agent with sensible defaults.
95		func NewAgent(client provider.Client, model, system string, tools Registry) *Agent {	102		func NewAgent(client provider.Client, model, system string, tools Registry) *Agent {
96		return &Agent{	103		return &Agent{
97	-	Client: client,	104	+	Client: client,
98	-	Model: model,	105	+	Model: model,
99	-	System: system,	106	+	System: system,
100	-	Tools: tools,	107	+	Tools: tools,
101	-	MaxSteps: 0, // 0 = unlimited	108	+	MaxSteps: 0, // 0 = unlimited
			109	+	MaxRetries: 3,
			110	+	RetryBaseDelay: 2 * time.Second,
102		}	111		}
103		}	112		}
104			113

@@ -337,8 +346,23 @@ func (a *Agent) runLoop(ctx context.Context, sink func(AgentEvent)) error {
337		return nil	346		return nil
338		}	347		}
339		}	348		}
340	-	stop, assistantMsg, err := a.oneTurn(ctx, sink)	349	+
341	-	sink(EvTurnEnd{Stop: stop, Err: err})	350	+	var (
			351	+	stop provider.StopReason
			352	+	assistantMsg provider.Message
			353	+	err error
			354	+	)
			355	+	for attempt := 0; ; attempt++ {
			356	+	stop, assistantMsg, err = a.oneTurn(ctx, sink)
			357	+	sink(EvTurnEnd{Stop: stop, Err: err})
			358	+	if err == nil \|\| !a.canRetryError(err, attempt) {
			359	+	break
			360	+	}
			361	+	a.dropLastAssistantMessage()
			362	+	if sleepErr := sleepRetry(ctx, a.retryDelay(attempt)); sleepErr != nil {
			363	+	return sleepErr
			364	+	}
			365	+	}
342		if err != nil {	366		if err != nil {
343		return err	367		return err
344		}	368		}

@@ -403,6 +427,77 @@ func (a *Agent) runLoop(ctx context.Context, sink func(AgentEvent)) error {
403		return nil	427		return nil
404		}	428		}
405			429
			430	+	func (a *Agent) canRetryError(err error, attempt int) bool {
			431	+	if err == nil \|\| a.MaxRetries <= 0 \|\| attempt >= a.MaxRetries {
			432	+	return false
			433	+	}
			434	+	if errors.Is(err, context.Canceled) \|\| errors.Is(err, context.DeadlineExceeded) {
			435	+	return false
			436	+	}
			437	+	msg := strings.ToLower(err.Error())
			438	+	if msg == "" \|\| isNonRetryableProviderLimit(msg) {
			439	+	return false
			440	+	}
			441	+	needles := []string{
			442	+	"overloaded", "provider returned error", "rate limit", "ratelimit", "too many requests",
			443	+	"429", "http 429", "500", "http 500", "502", "http 502", "503", "http 503", "504", "http 504",
			444	+	"service unavailable", "server error", "internal error", "network error", "connection error",
			445	+	"connection refused", "connection lost", "fetch failed", "upstream connect", "reset before headers",
			446	+	"socket hang up", "ended without", "stream ended before", "did not get a response", "timed out",
			447	+	"timeout", "terminated", "unexpected eof", "transport failure",
			448	+	}
			449	+	for _, needle := range needles {
			450	+	if strings.Contains(msg, needle) {
			451	+	return true
			452	+	}
			453	+	}
			454	+	return false
			455	+	}
			456	+
			457	+	func isNonRetryableProviderLimit(msg string) bool {
			458	+	needles := []string{
			459	+	"usage limit", "monthly usage limit", "freeusagelimit", "gousagelimit",
			460	+	"available balance", "insufficient_quota", "out of budget", "quota exceeded", "billing",
			461	+	}
			462	+	for _, needle := range needles {
			463	+	if strings.Contains(msg, needle) {
			464	+	return true
			465	+	}
			466	+	}
			467	+	return false
			468	+	}
			469	+
			470	+	func (a *Agent) retryDelay(attempt int) time.Duration {
			471	+	base := a.RetryBaseDelay
			472	+	if base <= 0 {
			473	+	base = 2 * time.Second
			474	+	}
			475	+	return base * time.Duration(1<<attempt)
			476	+	}
			477	+
			478	+	func sleepRetry(ctx context.Context, d time.Duration) error {
			479	+	if d <= 0 {
			480	+	return ctx.Err()
			481	+	}
			482	+	t := time.NewTimer(d)
			483	+	defer t.Stop()
			484	+	select {
			485	+	case <-ctx.Done():
			486	+	return ctx.Err()
			487	+	case <-t.C:
			488	+	return nil
			489	+	}
			490	+	}
			491	+
			492	+	func (a *Agent) dropLastAssistantMessage() {
			493	+	a.mu.Lock()
			494	+	defer a.mu.Unlock()
			495	+	if n := len(a.messages); n > 0 && a.messages[n-1].Role == provider.RoleAssistant {
			496	+	a.messages = a.messages[:n-1]
			497	+	a.rev++
			498	+	}
			499	+	}
			500	+
406		// oneTurn calls the LLM once, forwards events, returns the stop reason	501		// oneTurn calls the LLM once, forwards events, returns the stop reason
407		// and the assembled assistant message (already appended to the transcript).	502		// and the assembled assistant message (already appended to the transcript).
408		func (a *Agent) oneTurn(ctx context.Context, sink func(AgentEvent)) (provider.StopReason, provider.Message, error) {	503		func (a *Agent) oneTurn(ctx context.Context, sink func(AgentEvent)) (provider.StopReason, provider.Message, error) {

A packages/core/agent_retry_test.go +111 -0

@@ -0,0 +1,111 @@
	1	+	package core
	2	+
	3	+	import (
	4	+	"context"
	5	+	"fmt"
	6	+	"strings"
	7	+	"sync/atomic"
	8	+	"testing"
	9	+	"time"
	10	+
	11	+	"github.com/patriceckhart/zot/packages/provider"
	12	+	)
	13	+
	14	+	type retryFakeClient struct {
	15	+	calls int32
	16	+	}
	17	+
	18	+	func (c *retryFakeClient) Name() string { return "retry-fake" }
	19	+
	20	+	func (c *retryFakeClient) Stream(ctx context.Context, req provider.Request) (<-chan provider.Event, error) {
	21	+	call := atomic.AddInt32(&c.calls, 1)
	22	+	out := make(chan provider.Event, 4)
	23	+	go func() {
	24	+	defer close(out)
	25	+	out <- provider.EventStart{Provider: "retry-fake", Model: req.Model}
	26	+	if call == 1 {
	27	+	out <- provider.EventDone{Stop: provider.StopError, Err: fmt.Errorf("anthropic overloaded_error: Overloaded")}
	28	+	return
	29	+	}
	30	+	out <- provider.EventTextDelta{Delta: "ok"}
	31	+	out <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{
	32	+	Role: provider.RoleAssistant,
	33	+	Content: []provider.Content{provider.TextBlock{Text: "ok"}},
	34	+	}}
	35	+	}()
	36	+	return out, nil
	37	+	}
	38	+
	39	+	func TestAgentRetriesOverloadedStreamError(t *testing.T) {
	40	+	client := &retryFakeClient{}
	41	+	a := NewAgent(client, "fake-model", "system", Registry{})
	42	+	a.RetryBaseDelay = time.Millisecond
	43	+
	44	+	var turnErrs []string
	45	+	err := a.Prompt(context.Background(), "hello", nil, func(ev AgentEvent) {
	46	+	if e, ok := ev.(EvTurnEnd); ok && e.Err != nil {
	47	+	turnErrs = append(turnErrs, e.Err.Error())
	48	+	}
	49	+	})
	50	+	if err != nil {
	51	+	t.Fatalf("Prompt returned %v", err)
	52	+	}
	53	+	if got := atomic.LoadInt32(&client.calls); got != 2 {
	54	+	t.Fatalf("Stream calls = %d; want 2", got)
	55	+	}
	56	+	if len(turnErrs) != 1 \|\| !strings.Contains(turnErrs[0], "overloaded_error") {
	57	+	t.Fatalf("turn errors = %v; want one overloaded error before retry", turnErrs)
	58	+	}
	59	+	msgs := a.Messages()
	60	+	if len(msgs) != 2 {
	61	+	t.Fatalf("message count = %d; want user + final assistant", len(msgs))
	62	+	}
	63	+	if got := extractText(msgs[1]); got != "ok" {
	64	+	t.Fatalf("final assistant text = %q; want ok", got)
	65	+	}
	66	+	}
	67	+
	68	+	type partialRetryFakeClient struct {
	69	+	calls int32
	70	+	}
	71	+
	72	+	func (c *partialRetryFakeClient) Name() string { return "partial-retry-fake" }
	73	+
	74	+	func (c *partialRetryFakeClient) Stream(ctx context.Context, req provider.Request) (<-chan provider.Event, error) {
	75	+	call := atomic.AddInt32(&c.calls, 1)
	76	+	out := make(chan provider.Event, 4)
	77	+	go func() {
	78	+	defer close(out)
	79	+	out <- provider.EventStart{Provider: "partial-retry-fake", Model: req.Model}
	80	+	if call == 1 {
	81	+	out <- provider.EventTextDelta{Delta: "partial"}
	82	+	out <- provider.EventDone{Stop: provider.StopError, Err: fmt.Errorf("provider returned error: 503"), Message: provider.Message{
	83	+	Role: provider.RoleAssistant,
	84	+	Content: []provider.Content{provider.TextBlock{Text: "partial"}},
	85	+	}}
	86	+	return
	87	+	}
	88	+	out <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{
	89	+	Role: provider.RoleAssistant,
	90	+	Content: []provider.Content{provider.TextBlock{Text: "recovered"}},
	91	+	}}
	92	+	}()
	93	+	return out, nil
	94	+	}
	95	+
	96	+	func TestAgentDropsPartialAssistantBeforeRetry(t *testing.T) {
	97	+	client := &partialRetryFakeClient{}
	98	+	a := NewAgent(client, "fake-model", "system", Registry{})
	99	+	a.RetryBaseDelay = time.Millisecond
	100	+
	101	+	if err := a.Prompt(context.Background(), "hello", nil, nil); err != nil {
	102	+	t.Fatalf("Prompt returned %v", err)
	103	+	}
	104	+	msgs := a.Messages()
	105	+	if len(msgs) != 2 {
	106	+	t.Fatalf("message count = %d; want user + recovered assistant", len(msgs))
	107	+	}
	108	+	if got := extractText(msgs[1]); got != "recovered" {
	109	+	t.Fatalf("final assistant text = %q; want recovered", got)
	110	+	}
	111	+	}

@@ -0,0 +1,111 @@
	1	+	package core
	2	+
	3	+	import (
	4	+	"context"
	5	+	"fmt"
	6	+	"strings"
	7	+	"sync/atomic"
	8	+	"testing"
	9	+	"time"
	10	+
	11	+	"github.com/patriceckhart/zot/packages/provider"
	12	+	)
	13	+
	14	+	type retryFakeClient struct {
	15	+	calls int32
	16	+	}
	17	+
	18	+	func (c *retryFakeClient) Name() string { return "retry-fake" }
	19	+
	20	+	func (c *retryFakeClient) Stream(ctx context.Context, req provider.Request) (<-chan provider.Event, error) {
	21	+	call := atomic.AddInt32(&c.calls, 1)
	22	+	out := make(chan provider.Event, 4)
	23	+	go func() {
	24	+	defer close(out)
	25	+	out <- provider.EventStart{Provider: "retry-fake", Model: req.Model}
	26	+	if call == 1 {
	27	+	out <- provider.EventDone{Stop: provider.StopError, Err: fmt.Errorf("anthropic overloaded_error: Overloaded")}
	28	+	return
	29	+	}
	30	+	out <- provider.EventTextDelta{Delta: "ok"}
	31	+	out <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{
	32	+	Role: provider.RoleAssistant,
	33	+	Content: []provider.Content{provider.TextBlock{Text: "ok"}},
	34	+	}}
	35	+	}()
	36	+	return out, nil
	37	+	}
	38	+
	39	+	func TestAgentRetriesOverloadedStreamError(t *testing.T) {
	40	+	client := &retryFakeClient{}
	41	+	a := NewAgent(client, "fake-model", "system", Registry{})
	42	+	a.RetryBaseDelay = time.Millisecond
	43	+
	44	+	var turnErrs []string
	45	+	err := a.Prompt(context.Background(), "hello", nil, func(ev AgentEvent) {
	46	+	if e, ok := ev.(EvTurnEnd); ok && e.Err != nil {
	47	+	turnErrs = append(turnErrs, e.Err.Error())
	48	+	}
	49	+	})
	50	+	if err != nil {
	51	+	t.Fatalf("Prompt returned %v", err)
	52	+	}
	53	+	if got := atomic.LoadInt32(&client.calls); got != 2 {
	54	+	t.Fatalf("Stream calls = %d; want 2", got)
	55	+	}
	56	+	if len(turnErrs) != 1 \|\| !strings.Contains(turnErrs[0], "overloaded_error") {
	57	+	t.Fatalf("turn errors = %v; want one overloaded error before retry", turnErrs)
	58	+	}
	59	+	msgs := a.Messages()
	60	+	if len(msgs) != 2 {
	61	+	t.Fatalf("message count = %d; want user + final assistant", len(msgs))
	62	+	}
	63	+	if got := extractText(msgs[1]); got != "ok" {
	64	+	t.Fatalf("final assistant text = %q; want ok", got)
	65	+	}
	66	+	}
	67	+
	68	+	type partialRetryFakeClient struct {
	69	+	calls int32
	70	+	}
	71	+
	72	+	func (c *partialRetryFakeClient) Name() string { return "partial-retry-fake" }
	73	+
	74	+	func (c *partialRetryFakeClient) Stream(ctx context.Context, req provider.Request) (<-chan provider.Event, error) {
	75	+	call := atomic.AddInt32(&c.calls, 1)
	76	+	out := make(chan provider.Event, 4)
	77	+	go func() {
	78	+	defer close(out)
	79	+	out <- provider.EventStart{Provider: "partial-retry-fake", Model: req.Model}
	80	+	if call == 1 {
	81	+	out <- provider.EventTextDelta{Delta: "partial"}
	82	+	out <- provider.EventDone{Stop: provider.StopError, Err: fmt.Errorf("provider returned error: 503"), Message: provider.Message{
	83	+	Role: provider.RoleAssistant,
	84	+	Content: []provider.Content{provider.TextBlock{Text: "partial"}},
	85	+	}}
	86	+	return
	87	+	}
	88	+	out <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{
	89	+	Role: provider.RoleAssistant,
	90	+	Content: []provider.Content{provider.TextBlock{Text: "recovered"}},
	91	+	}}
	92	+	}()
	93	+	return out, nil
	94	+	}
	95	+
	96	+	func TestAgentDropsPartialAssistantBeforeRetry(t *testing.T) {
	97	+	client := &partialRetryFakeClient{}
	98	+	a := NewAgent(client, "fake-model", "system", Registry{})
	99	+	a.RetryBaseDelay = time.Millisecond
	100	+
	101	+	if err := a.Prompt(context.Background(), "hello", nil, nil); err != nil {
	102	+	t.Fatalf("Prompt returned %v", err)
	103	+	}
	104	+	msgs := a.Messages()
	105	+	if len(msgs) != 2 {
	106	+	t.Fatalf("message count = %d; want user + recovered assistant", len(msgs))
	107	+	}
	108	+	if got := extractText(msgs[1]); got != "recovered" {
	109	+	t.Fatalf("final assistant text = %q; want recovered", got)
	110	+	}
	111	+	}