1package core
2
3import (
4 "context"
5 "fmt"
6 "strings"
7 "sync/atomic"
8 "testing"
9 "time"
10
11 "github.com/patriceckhart/zot/packages/provider"
12)
13
14type retryFakeClient struct {
15 calls int32
16}
17
18func (c *retryFakeClient) Name() string { return "retry-fake" }
19
20func (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
39func 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
68type partialRetryFakeClient struct {
69 calls int32
70}
71
72func (c *partialRetryFakeClient) Name() string { return "partial-retry-fake" }
73
74func (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
96func 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}