返回 DeepSeek-Reasonix
pi_recovery_test.go
根目录 / internal / agent / pi_recovery_test.go
1 package agent
2
3 import (
4 "context"
5 "errors"
6 "testing"
7 "time"
8
9 "reasonix/internal/agent/testutil"
10 "reasonix/internal/event"
11 "reasonix/internal/provider"
12 )
13
14 func TestCompatibleMissingReasoningDoesNotRegenerate(t *testing.T) {
15 mock := testutil.NewMock("compatible", testutil.Turn{ToolCalls: []provider.ToolCall{{ID: "one", Name: "echo", Arguments: `{"text":"hi"}`}}}, testutil.Turn{Text: "done"})
16 sink := &recordSink{}
17 a := New(toolCallReasoningRequiredProvider{mock}, echoRegistry(), NewSession(""), Options{MissingReasoningWarnStateDir: t.TempDir()}, sink)
18 if err := a.Run(withNoClosedLoop(context.Background()), "go"); err != nil {
19 t.Fatal(err)
20 }
21 if mock.CallCount() != 2 || len(sink.kinds(event.ToolResult)) != 1 || len(sink.kinds(event.Retrying)) != 0 {
22 t.Fatalf("calls=%d tools=%d retries=%d", mock.CallCount(), len(sink.kinds(event.ToolResult)), len(sink.kinds(event.Retrying)))
23 }
24 }
25
26 type transientHeaderProvider struct{ calls int }
27
28 func (*transientHeaderProvider) Name() string { return "transient" }
29 func (p *transientHeaderProvider) Stream(ctx context.Context, _ provider.Request) (<-chan provider.Chunk, error) {
30 p.calls++
31 return nil, &provider.APIError{Status: 503}
32 }
33 func TestMainWaitsAfterFiniteRetriesAndCancels(t *testing.T) {
34 p := &transientHeaderProvider{}
35 ctx, cancel := context.WithCancel(context.Background())
36 defer cancel()
37 old := recoverySleep
38 defer func() { recoverySleep = old }()
39 recoverySleep = func(ctx context.Context, d time.Duration) bool {
40 if d < time.Minute {
41 t.Errorf("wait=%s", d)
42 }
43 cancel()
44 return false
45 }
46 sink := &recordSink{}
47 a := New(p, echoRegistry(), NewSession(""), Options{}, sink)
48 err := a.Run(withNoClosedLoop(ctx), "go")
49 if !errors.Is(err, context.Canceled) || p.calls != 4 {
50 t.Fatalf("calls=%d err=%v", p.calls, err)
51 }
52 retries := sink.kinds(event.Retrying)
53 if len(retries) != 4 || retries[3].Recovery == nil || !retries[3].Recovery.Waiting {
54 t.Fatalf("retries=%+v", retries)
55 }
56 }
57 func TestSubagentStopsAfterFiniteRetries(t *testing.T) {
58 p := &transientHeaderProvider{}
59 a := New(p, echoRegistry(), NewSession(""), Options{}, event.Discard)
60 err := a.Run(withNoClosedLoop(WithSubagentDepth(context.Background(), 1)), "go")
61 if err == nil || p.calls != 4 {
62 t.Fatalf("calls=%d err=%v", p.calls, err)
63 }
64 }
65
66 func TestTruncatedArgumentsNeverExecute(t *testing.T) {
67 mock := testutil.NewMock("m", testutil.Turn{ToolCalls: []provider.ToolCall{{ID: "cut", Name: "echo", Arguments: `{"text":"partial"}`}}, Usage: &provider.Usage{FinishReason: "length"}}, testutil.Turn{Text: "recovered"})
68 sink := &recordSink{}
69 a := New(mock, echoRegistry(), NewSession(""), Options{}, sink)
70 if err := a.Run(withNoClosedLoop(context.Background()), "go"); err != nil {
71 t.Fatal(err)
72 }
73 results := sink.kinds(event.ToolResult)
74 if len(results) != 1 || results[0].Tool.Output == "echoed: partial" {
75 t.Fatalf("results=%+v", results)
76 }
77 for _, m := range a.Session().Snapshot() {
78 if m.ToolCallID == "cut" && m.ToolRunState != provider.ToolRunNotStarted {
79 t.Fatalf("state=%s", m.ToolRunState)
80 }
81 }
82 }
83
84 func TestPartialStreamNeverEntersContinuousWaiting(t *testing.T) {
85 turns := make([]testutil.Turn, maxSamplingAttempts)
86 for i := range turns {
87 turns[i] = testutil.Turn{Text: "partial", ChunkError: provider.StreamInterrupt(errors.New("closed"), provider.StreamInterruptPrematureEOF)}
88 }
89 mock := testutil.NewMock("m", turns...)
90 sink := &recordSink{}
91 a := New(mock, echoRegistry(), NewSession(""), Options{}, sink)
92 if err := a.Run(withNoClosedLoop(context.Background()), "go"); err == nil {
93 t.Fatal("partial stream accepted")
94 }
95 if mock.CallCount() != 4 {
96 t.Fatalf("calls=%d", mock.CallCount())
97 }
98 for _, e := range sink.kinds(event.Retrying) {
99 if e.Recovery != nil && e.Recovery.Waiting {
100 t.Fatal("partial generation waited indefinitely")
101 }
102 }
103 }
104
105 func TestKnownSpendStillStopsRecoveryWithUnknownUsage(t *testing.T) {
106 var b runBudget
107 b.observe(&provider.Usage{PromptTokens: 1000000, Unknown: true}, &provider.Pricing{Input: 2, Currency: "USD"})
108 if axis, _ := b.exceeded(TaskBudget{Cost: 1}); axis != "cost" || b.totals().Priced {
109 t.Fatalf("axis=%s totals=%+v", axis, b.totals())
110 }
111 }
112
113 func TestPiReferenceRetryScenarios(t *testing.T) {
114 for _, tc := range []struct {
115 name string
116 turns []testutil.Turn
117 calls int
118 success bool
119 }{
120 {"temporary_then_success", []testutil.Turn{{StreamError: &provider.APIError{Status: 503}}, {StreamError: &provider.APIError{Status: 503}}, {Text: "done"}}, 3, true},
121 {"quota", []testutil.Turn{{StreamError: &provider.APIError{Status: 429, Body: "insufficient_quota"}}}, 1, false},
122 } {
123 t.Run(tc.name, func(t *testing.T) {
124 p := testutil.NewMock("reference", tc.turns...)
125 a := New(p, echoRegistry(), NewSession(""), Options{}, event.Discard)
126 err := a.Run(withNoClosedLoop(context.Background()), "go")
127 if (err == nil) != tc.success || p.CallCount() != tc.calls {
128 t.Fatalf("calls=%d err=%v", p.CallCount(), err)
129 }
130 })
131 }
132 }
133
134 func TestMixedFailuresCannotRenewRecoveryBudget(t *testing.T) {
135 missing := testutil.Turn{ToolCalls: []provider.ToolCall{{ID: "missing", Name: "echo", Arguments: `{"text":"unsafe"}`}}}
136 p := testutil.NewMock("strict", testutil.Turn{StreamError: &provider.APIError{Status: 503}}, testutil.Turn{StreamError: &provider.APIError{Status: 503}}, missing, missing)
137 sink := &recordSink{}
138 a := New(strictToolCallReasoningProvider{p}, echoRegistry(), NewSession(""), Options{}, sink)
139 if err := a.Run(withNoClosedLoop(context.Background()), "go"); err == nil {
140 t.Fatal("invalid reasoning accepted")
141 }
142 if p.CallCount() != 4 || len(sink.kinds(event.ToolResult)) != 0 {
143 t.Fatalf("calls=%d tools=%d", p.CallCount(), len(sink.kinds(event.ToolResult)))
144 }
145 }
146
147 type canceledCompletionProvider struct{ cancel context.CancelFunc }
148
149 func (*canceledCompletionProvider) Name() string { return "late" }
150 func (p *canceledCompletionProvider) Stream(context.Context, provider.Request) (<-chan provider.Chunk, error) {
151 ch := make(chan provider.Chunk, 2)
152 ch <- provider.Chunk{Type: provider.ChunkToolCall, ToolCall: &provider.ToolCall{ID: "late", Name: "echo", Arguments: `{"text":"late"}`}}
153 ch <- provider.Chunk{Type: provider.ChunkDone}
154 close(ch)
155 p.cancel()
156 return ch, nil
157 }
158 func TestCanceledCompletionCannotStartTools(t *testing.T) {
159 ctx, cancel := context.WithCancel(context.Background())
160 defer cancel()
161 sink := &recordSink{}
162 a := New(&canceledCompletionProvider{cancel}, echoRegistry(), NewSession(""), Options{}, sink)
163 if err := a.Run(withNoClosedLoop(ctx), "go"); !errors.Is(err, context.Canceled) {
164 t.Fatalf("err=%v", err)
165 }
166 if len(sink.kinds(event.ToolResult)) != 0 {
167 t.Fatal("late completion executed a tool")
168 }
169 }
170
171 func TestPlannerAndNonRetryableFailuresNeverEnterContinuousWait(t *testing.T) {
172 a := New(&transientHeaderProvider{}, echoRegistry(), NewSession(""), Options{}, event.Discard)
173 failure := provider.ClassifyRecovery(&provider.APIError{Status: 503})
174 if a.canWaitSampling(context.Background(), &samplingRecoveryState{}, provider.ClassifyRecovery(&provider.APIError{Status: 409})) {
175 t.Fatal("conflict allowed endless waiting")
176 }
177 ctx := context.WithValue(context.Background(), turnContextRoleKey{}, turnContextPlanner)
178 if a.canWaitSampling(ctx, &samplingRecoveryState{}, failure) {
179 t.Fatal("planner allowed endless waiting")
180 }
181 }
182
182 lines GO