返回 DeepSeek-Reasonix
reconnect_test.go
根目录 / internal / provider / openai / reconnect_test.go
1 package openai
2
3 import (
4 "context"
5 "errors"
6 "io"
7 "net"
8 "net/http"
9 "net/http/httptest"
10 "strings"
11 "sync/atomic"
12 "testing"
13 "time"
14
15 "reasonix/internal/provider"
16 )
17
18 // rstAfter writes a 200 SSE head plus the given prelude, then forces a TCP RST
19 // (SetLinger(0) + Close) so the client read fails like a proxy that idle-drops
20 // the long-lived connection (wsarecv: forcibly closed), not a clean EOF.
21 func rstAfter(t *testing.T, w http.ResponseWriter, prelude string) {
22 t.Helper()
23 hj, ok := w.(http.Hijacker)
24 if !ok {
25 t.Fatal("ResponseWriter is not a Hijacker")
26 }
27 conn, buf, err := hj.Hijack()
28 if err != nil {
29 t.Fatalf("hijack: %v", err)
30 }
31 _, _ = buf.WriteString("HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\n\r\n")
32 _, _ = buf.WriteString(prelude)
33 _ = buf.Flush()
34 if tcp, ok := conn.(*net.TCPConn); ok {
35 _ = tcp.SetLinger(0)
36 }
37 _ = conn.Close()
38 }
39
40 // TestStreamSurfacesEarlyConnResetAsInterrupt moves body-phase replay to the
41 // Agent: a pre-output connection reset is StreamInterruptedError, not an
42 // in-provider transparent reconnect (avoids stacked retry budgets).
43 func TestStreamSurfacesEarlyConnResetAsInterrupt(t *testing.T) {
44 var reqs int
45 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
46 reqs++
47 rstAfter(t, w, ": keep-alive\n\n") // a comment line, zero model output
48 }))
49 defer srv.Close()
50
51 p, err := New(provider.Config{Name: "deepseek", BaseURL: srv.URL, Model: "deepseek-v4", APIKey: "k"})
52 if err != nil {
53 t.Fatalf("New: %v", err)
54 }
55 ch, err := p.Stream(context.Background(), provider.Request{Messages: []provider.Message{{Role: provider.RoleUser, Content: "hi"}}})
56 if err != nil {
57 t.Fatalf("Stream: %v", err)
58 }
59
60 var gotInterrupted bool
61 for chunk := range ch {
62 if chunk.Type == provider.ChunkError {
63 var interrupted *provider.StreamInterruptedError
64 gotInterrupted = errors.As(chunk.Err, &interrupted)
65 }
66 }
67 if !gotInterrupted {
68 t.Error("early conn reset must surface as StreamInterruptedError for Agent replay")
69 }
70 if reqs != 1 {
71 t.Errorf("server saw %d requests, want 1 (no provider body replay)", reqs)
72 }
73 }
74
75 func TestStreamCancelDoesNotReconnect(t *testing.T) {
76 var reqs atomic.Int32
77 ready := make(chan struct{})
78 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
79 first := reqs.Add(1) == 1
80 w.Header().Set("Content-Type", "text/event-stream")
81 _, _ = io.WriteString(w, ": keep-alive\n\n")
82 flush(w)
83 if first {
84 close(ready)
85 }
86 <-r.Context().Done()
87 }))
88 defer srv.Close()
89
90 p, err := New(provider.Config{Name: "deepseek", BaseURL: srv.URL, Model: "deepseek-v4", APIKey: "k"})
91 if err != nil {
92 t.Fatalf("New: %v", err)
93 }
94 ctx, cancel := context.WithCancel(context.Background())
95 ch, err := p.Stream(ctx, provider.Request{Messages: []provider.Message{{Role: provider.RoleUser, Content: "hi"}}})
96 if err != nil {
97 t.Fatalf("Stream: %v", err)
98 }
99 select {
100 case <-ready:
101 case <-time.After(2 * time.Second):
102 t.Fatal("server did not receive the streaming request")
103 }
104 cancel()
105
106 var got error
107 for chunk := range ch {
108 if chunk.Type == provider.ChunkError {
109 got = chunk.Err
110 }
111 }
112 // Depending on whether the server close or the client watchdog observes
113 // cancellation first, the stream may close silently or surface cancellation.
114 // The contract guarded here is that cancellation never triggers a replay.
115 if got != nil && !errors.Is(got, context.Canceled) {
116 t.Fatalf("stream error = %v, want nil or context.Canceled", got)
117 }
118 if reqs.Load() != 1 {
119 t.Fatalf("cancelled stream reconnected; server saw %d requests, want 1", reqs.Load())
120 }
121 }
122
123 // TestStreamTreatsCleanEOFWithoutDoneAsCut reproduces issue #3953: a proxy that
124 // idle-closes the SSE connection with a clean FIN ends the scan with no error,
125 // which used to commit the turn as complete. Body-phase cuts surface as
126 // StreamInterruptedError so the Agent can replay the frozen request.
127 func TestStreamTreatsCleanEOFWithoutDoneAsCut(t *testing.T) {
128 var reqs int
129 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
130 reqs++
131 w.Header().Set("Content-Type", "text/event-stream")
132 _, _ = io.WriteString(w, ": keep-alive\n\n") // clean close, no [DONE], no finish_reason
133 }))
134 defer srv.Close()
135
136 p, err := New(provider.Config{Name: "deepseek", BaseURL: srv.URL, Model: "deepseek-v4", APIKey: "k"})
137 if err != nil {
138 t.Fatalf("New: %v", err)
139 }
140 ch, err := p.Stream(context.Background(), provider.Request{Messages: []provider.Message{{Role: provider.RoleUser, Content: "hi"}}})
141 if err != nil {
142 t.Fatalf("Stream: %v", err)
143 }
144
145 var gotInterrupted bool
146 for chunk := range ch {
147 switch chunk.Type {
148 case provider.ChunkToolCall:
149 t.Fatalf("incomplete stream must not emit tool calls: %+v", chunk.ToolCall)
150 case provider.ChunkError:
151 var interrupted *provider.StreamInterruptedError
152 gotInterrupted = errors.As(chunk.Err, &interrupted)
153 }
154 }
155 if !gotInterrupted {
156 t.Error("clean EOF before terminal must surface as StreamInterruptedError")
157 }
158 if reqs != 1 {
159 t.Errorf("server saw %d requests, want 1 (no provider body replay)", reqs)
160 }
161 }
162
163 // TestStreamDropsPartialToolCallOnCleanEOF is the post-output half of #3953: the
164 // connection dies mid-tool-call after the call's start was forwarded. The partial
165 // arguments must never surface as a ChunkToolCall; the cut surfaces as a stream
166 // interruption so the agent's recovery path takes over.
167 func TestStreamDropsPartialToolCallOnCleanEOF(t *testing.T) {
168 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
169 w.Header().Set("Content-Type", "text/event-stream")
170 _, _ = io.WriteString(w, "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"c1\",\"function\":{\"name\":\"bash\",\"arguments\":\"{\"}}]}}]}\n\n")
171 }))
172 defer srv.Close()
173
174 p, err := New(provider.Config{Name: "deepseek", BaseURL: srv.URL, Model: "deepseek-v4", APIKey: "k"})
175 if err != nil {
176 t.Fatalf("New: %v", err)
177 }
178 ch, err := p.Stream(context.Background(), provider.Request{Messages: []provider.Message{{Role: provider.RoleUser, Content: "hi"}}})
179 if err != nil {
180 t.Fatalf("Stream: %v", err)
181 }
182
183 var gotInterrupted bool
184 for chunk := range ch {
185 switch chunk.Type {
186 case provider.ChunkToolCall:
187 t.Fatalf("partial tool call surfaced: %+v", chunk.ToolCall)
188 case provider.ChunkError:
189 var interrupted *provider.StreamInterruptedError
190 gotInterrupted = errors.As(chunk.Err, &interrupted)
191 }
192 }
193 if !gotInterrupted {
194 t.Error("a cut after the tool-call start should surface as a stream interruption")
195 }
196 }
197
198 // TestStreamAcceptsFinishReasonWithoutDone keeps gateways that omit the [DONE]
199 // sentinel working: a finish_reason marks the turn complete on its own.
200 func TestStreamAcceptsFinishReasonWithoutDone(t *testing.T) {
201 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
202 w.Header().Set("Content-Type", "text/event-stream")
203 _, _ = io.WriteString(w, "data: {\"choices\":[{\"delta\":{\"content\":\"hello\"},\"finish_reason\":\"stop\"}]}\n\n")
204 }))
205 defer srv.Close()
206
207 p, err := New(provider.Config{Name: "deepseek", BaseURL: srv.URL, Model: "deepseek-v4", APIKey: "k"})
208 if err != nil {
209 t.Fatalf("New: %v", err)
210 }
211 ch, err := p.Stream(context.Background(), provider.Request{Messages: []provider.Message{{Role: provider.RoleUser, Content: "hi"}}})
212 if err != nil {
213 t.Fatalf("Stream: %v", err)
214 }
215
216 var text strings.Builder
217 for chunk := range ch {
218 if chunk.Type == provider.ChunkError {
219 t.Fatalf("finish_reason without [DONE] should complete cleanly: %v", chunk.Err)
220 }
221 if chunk.Type == provider.ChunkText {
222 text.WriteString(chunk.Text)
223 }
224 }
225 if text.String() != "hello" {
226 t.Errorf("text = %q, want %q", text.String(), "hello")
227 }
228 }
229
230 // TestStreamDoesNotReplayAfterOutput guards against duplicated output: once a
231 // token has streamed, a mid-stream reset must surface as an error rather than
232 // replaying the request (which would re-emit the already-shown text).
233 func TestStreamDoesNotReplayAfterOutput(t *testing.T) {
234 var reqs int
235 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
236 reqs++
237 rstAfter(t, w, "data: {\"choices\":[{\"delta\":{\"content\":\"partial\"}}]}\n\n")
238 }))
239 defer srv.Close()
240
241 p, err := New(provider.Config{Name: "deepseek", BaseURL: srv.URL, Model: "deepseek-v4", APIKey: "k"})
242 if err != nil {
243 t.Fatalf("New: %v", err)
244 }
245 ch, err := p.Stream(context.Background(), provider.Request{Messages: []provider.Message{{Role: provider.RoleUser, Content: "hi"}}})
246 if err != nil {
247 t.Fatalf("Stream: %v", err)
248 }
249
250 var text strings.Builder
251 var gotErr bool
252 var gotInterrupted bool
253 for chunk := range ch {
254 switch chunk.Type {
255 case provider.ChunkText:
256 text.WriteString(chunk.Text)
257 case provider.ChunkError:
258 gotErr = true
259 var interrupted *provider.StreamInterruptedError
260 gotInterrupted = errors.As(chunk.Err, &interrupted)
261 }
262 }
263 if text.String() != "partial" {
264 t.Errorf("text = %q, want %q (the one delta that streamed)", text.String(), "partial")
265 }
266 if !gotErr {
267 t.Error("a reset after output should surface a ChunkError")
268 }
269 if !gotInterrupted {
270 t.Error("a reset after output should be marked as a stream interruption")
271 }
272 if reqs != 1 {
273 t.Errorf("server saw %d requests, want 1 (no replay after output)", reqs)
274 }
275 }
276
276 lines GO