1package provider
2
3import (
4 "context"
5 "crypto/tls"
6 "encoding/json"
7 "fmt"
8 "net/http"
9 "net/url"
10 "reflect"
11 "strings"
12 "sync"
13
14 "golang.org/x/net/websocket"
15)
16
17// responsesWebSocketClient uses the public Responses WebSocket mode for one
18// logical session at a time. It deliberately wraps only the public OpenAI
19// Responses client: ChatGPT Codex and OpenAI-compatible endpoints have
20// different authentication and cache contracts and must retain their HTTP/SSE
21// transport until they independently declare WebSocket support.
22type responsesWebSocketClient struct {
23 http *codexClient
24
25 mu sync.Mutex
26 sessions map[string]*responsesWebSocketSession
27 dial func(context.Context) (*websocket.Conn, error)
28}
29
30type responsesWebSocketSession struct {
31 mu sync.Mutex
32
33 turn chan struct{}
34
35 conn *websocket.Conn
36 raw chan sseEvent
37 readerDone chan struct{}
38
39 lastRequest []Message
40 lastResponseID string
41}
42
43func newResponsesWebSocketClient(httpClient *codexClient) Client {
44 c := &responsesWebSocketClient{
45 http: httpClient,
46 sessions: make(map[string]*responsesWebSocketSession),
47 }
48 c.dial = c.dialResponses
49 return c
50}
51
52func (c *responsesWebSocketClient) Name() string { return c.http.Name() }
53
54// Close releases all persistent sockets held by this client. Callers that
55// create a client for a bounded runtime (such as a resident child) must call
56// this when that runtime stops; otherwise a warm session can outlive its
57// owner until the process exits.
58func (c *responsesWebSocketClient) Close() error {
59 c.mu.Lock()
60 sessions := c.sessions
61 c.sessions = make(map[string]*responsesWebSocketSession)
62 c.mu.Unlock()
63 for _, session := range sessions {
64 session.invalidate()
65 }
66 return nil
67}
68
69func (c *responsesWebSocketClient) session(id string) *responsesWebSocketSession {
70 c.mu.Lock()
71 defer c.mu.Unlock()
72 if session := c.sessions[id]; session != nil {
73 return session
74 }
75 session := &responsesWebSocketSession{turn: make(chan struct{}, 1)}
76 session.turn <- struct{}{}
77 c.sessions[id] = session
78 return session
79}
80
81func (c *responsesWebSocketClient) dialResponses(ctx context.Context) (*websocket.Conn, error) {
82 endpoint, err := responsesWebSocketURL(c.http.baseURL)
83 if err != nil {
84 return nil, err
85 }
86 config, err := websocket.NewConfig(endpoint, "https://api.openai.com")
87 if err != nil {
88 return nil, err
89 }
90 config.Header = make(http.Header)
91 config.Header.Set("authorization", "Bearer "+c.http.token)
92 if tlsConfig := responsesWebSocketTLSConfig(c.http.http); tlsConfig != nil {
93 config.TlsConfig = tlsConfig
94 }
95 return config.DialContext(ctx)
96}
97
98// responsesWebSocketTLSConfig carries the explicit TLS policy supplied by
99// WithHTTPClient into the WebSocket dial.
100func responsesWebSocketTLSConfig(client *http.Client) *tls.Config {
101 if client == nil {
102 return nil
103 }
104 transport := client.Transport
105 if public, ok := transport.(*openaiResponsesTransport); ok {
106 transport = public.inner
107 }
108 if transport, ok := transport.(*http.Transport); ok && transport.TLSClientConfig != nil {
109 return transport.TLSClientConfig.Clone()
110 }
111 return nil
112}
113
114// responsesWebSocketHTTPCompatible refuses to bypass an injected HTTP
115// transport. A raw WebSocket dial can faithfully carry a TLS policy but not a
116// custom RoundTripper or configured proxy, so those configurations retain the
117// established HTTP/SSE transport instead.
118func responsesWebSocketHTTPCompatible(client *http.Client) bool {
119 if client == nil || client.Transport == nil {
120 return true
121 }
122 transport := client.Transport
123 if public, ok := transport.(*openaiResponsesTransport); ok {
124 transport = public.inner
125 }
126 base, ok := transport.(*http.Transport)
127 if !ok {
128 return false
129 }
130 if base.Proxy == nil {
131 return true
132 }
133 request, _ := http.NewRequest(http.MethodGet, "https://api.openai.com", nil)
134 proxy, err := base.Proxy(request)
135 return err == nil && proxy == nil
136}
137
138func responsesWebSocketURL(endpoint string) (string, error) {
139 u, err := url.Parse(endpoint)
140 if err != nil {
141 return "", fmt.Errorf("parse Responses endpoint: %w", err)
142 }
143 switch u.Scheme {
144 case "https":
145 u.Scheme = "wss"
146 case "http":
147 u.Scheme = "ws"
148 default:
149 return "", fmt.Errorf("unsupported Responses endpoint scheme %q", u.Scheme)
150 }
151 return u.String(), nil
152}
153
154func (c *responsesWebSocketClient) Stream(ctx context.Context, req Request) (<-chan Event, error) {
155 // WebSocket continuation state is keyed by the durable core session ID.
156 // Requests without that identity retain HTTP/SSE semantics rather than
157 // accidentally pooling unrelated conversations.
158 sessionID := strings.TrimSpace(req.Context.SessionID)
159 if sessionID == "" {
160 return c.http.Stream(ctx, req)
161 }
162
163 session := c.session(sessionID)
164 select {
165 case <-ctx.Done():
166 return nil, ctx.Err()
167 case <-session.turn:
168 }
169
170 release := true
171 defer func() {
172 if release {
173 session.turn <- struct{}{}
174 }
175 }()
176
177 streamReq, previousResponseID := session.incrementalRequest(req)
178 payload, err := c.websocketPayload(streamReq, sessionID, previousResponseID)
179 if err != nil {
180 return nil, err
181 }
182
183 if !responsesWebSocketHTTPCompatible(c.http.http) {
184 return c.http.Stream(ctx, req)
185 }
186 reportRequestAttempt(req.Lifecycle, 1, 2)
187 if err := session.ensureConnected(ctx, c.dial); err != nil {
188 // A connection that never became usable has no server-side turn. The
189 // SSE fallback therefore receives the complete request and preserves
190 // the same session-derived cache key without sharing WebSocket state.
191 session.invalidate()
192 return c.http.Stream(ctx, req)
193 }
194 conn, raw := session.connection()
195 if err := websocket.Message.Send(conn, string(payload)); err != nil {
196 session.invalidate()
197 return c.http.Stream(ctx, req)
198 }
199
200 out := make(chan Event, 16)
201 release = false
202 go func() {
203 defer func() { session.turn <- struct{}{} }()
204 defer close(out)
205
206 first, responseID, ok := awaitResponsesWebSocketResponseStart(ctx, raw)
207 if !ok && ctx.Err() != nil {
208 session.invalidateConnection(conn)
209 c.http.runResponseEvents(ctx, req, out, make(chan sseEvent), nil)
210 return
211 }
212 if !ok || responseWebSocketRecoveryEvent(first) {
213 // The server may discard a response ID or close an idle connection.
214 // Reconnect once with the full transcript before falling back to
215 // HTTP/SSE. That keeps the stable cache key and logical IDs intact
216 // without ever reusing the invalid previous_response_id.
217 session.invalidate()
218 reconnected, reconnectRaw, reconnectFirst, reconnectResponseID, err := c.reconnectFull(ctx, session, req, sessionID)
219 if err != nil {
220 c.forwardResponsesHTTP(ctx, req, out)
221 return
222 }
223 conn, raw, first, responseID = reconnected, reconnectRaw, reconnectFirst, reconnectResponseID
224 }
225 var stateMu sync.Mutex
226 terminal := false
227 finished := make(chan struct{})
228 go func(conn *websocket.Conn) {
229 select {
230 case <-ctx.Done():
231 stateMu.Lock()
232 if !terminal {
233 session.invalidateConnection(conn)
234 }
235 stateMu.Unlock()
236 case <-finished:
237 }
238 }(conn)
239
240 completed := false
241 turnRaw := filterResponsesWebSocketEvents(raw, responseID, finished)
242 c.http.runResponseEventsWithFirst(ctx, req, out, turnRaw, &first, func(responseID string) {
243 stateMu.Lock()
244 terminal = true
245 stateMu.Unlock()
246 completed = true
247 session.recordCompleted(req.Messages, responseID)
248 })
249 close(finished)
250 if !completed {
251 session.invalidate()
252 }
253 }()
254 return out, nil
255}
256
257func (c *responsesWebSocketClient) websocketPayload(req Request, sessionID, previousResponseID string) ([]byte, error) {
258 wire, err := c.http.buildRequest(req)
259 if err != nil {
260 return nil, err
261 }
262 if c.http.modelName != nil {
263 wire.Model = c.http.modelName(wire.Model)
264 }
265 if c.http.capabilities.StablePromptCacheKey {
266 wire.PromptCacheKey = sessionID
267 }
268 wire.Stream = false // WebSocket mode does not send HTTP stream controls.
269 wire.PreviousResponseID = previousResponseID
270 return json.Marshal(struct {
271 Type string `json:"type"`
272 *codexRequest
273 }{Type: "response.create", codexRequest: wire})
274}
275
276// reconnectFull establishes one replacement WebSocket turn after a cached
277// continuation failed. A fresh connection cannot rely on the previous server
278// response, so it always sends the complete local transcript.
279func (c *responsesWebSocketClient) reconnectFull(ctx context.Context, session *responsesWebSocketSession, req Request, sessionID string) (*websocket.Conn, <-chan sseEvent, sseEvent, string, error) {
280 payload, err := c.websocketPayload(req, sessionID, "")
281 if err != nil {
282 return nil, nil, sseEvent{}, "", err
283 }
284 reportRequestAttempt(req.Lifecycle, 2, 2)
285 if err := session.ensureConnected(ctx, c.dial); err != nil {
286 return nil, nil, sseEvent{}, "", err
287 }
288 conn, raw := session.connection()
289 if err := websocket.Message.Send(conn, string(payload)); err != nil {
290 session.invalidateConnection(conn)
291 return nil, nil, sseEvent{}, "", err
292 }
293 first, responseID, ok := awaitResponsesWebSocketResponseStart(ctx, raw)
294 if !ok || responseWebSocketRecoveryEvent(first) {
295 session.invalidateConnection(conn)
296 if first.Err != nil {
297 return nil, nil, sseEvent{}, "", first.Err
298 }
299 return nil, nil, sseEvent{}, "", fmt.Errorf("responses WebSocket reconnect did not produce a usable event")
300 }
301 return conn, raw, first, responseID, nil
302}
303
304func awaitResponsesWebSocketEvent(ctx context.Context, raw <-chan sseEvent) (sseEvent, bool) {
305 select {
306 case <-ctx.Done():
307 return sseEvent{}, false
308 case event, ok := <-raw:
309 return event, ok
310 }
311}
312
313// awaitResponsesWebSocketResponseStart discards frames left behind by an
314// earlier response until the Responses protocol announces the response this
315// request created. The API runs responses sequentially on a socket, but this
316// guard keeps a late terminal frame from completing the next local turn.
317func awaitResponsesWebSocketResponseStart(ctx context.Context, raw <-chan sseEvent) (sseEvent, string, bool) {
318 for {
319 event, ok := awaitResponsesWebSocketEvent(ctx, raw)
320 if !ok || responseWebSocketRecoveryEvent(event) {
321 return event, "", ok
322 }
323 var payload struct {
324 Type string `json:"type"`
325 Response struct {
326 ID string `json:"id"`
327 } `json:"response"`
328 }
329 if err := json.Unmarshal([]byte(event.Data), &payload); err != nil {
330 continue
331 }
332 if payload.Type != "response.created" {
333 continue
334 }
335 if strings.TrimSpace(payload.Response.ID) == "" {
336 return sseEvent{Err: fmt.Errorf("responses WebSocket created event has no response ID")}, "", false
337 }
338 return event, payload.Response.ID, true
339 }
340}
341
342// filterResponsesWebSocketEvents retains only frames for the response that
343// started this turn when the protocol supplies an identity. Regular Responses
344// streaming events do not all repeat response_id; the preceding created-event
345// barrier is therefore also required by the sequential WebSocket contract.
346func filterResponsesWebSocketEvents(raw <-chan sseEvent, responseID string, done <-chan struct{}) <-chan sseEvent {
347 out := make(chan sseEvent, 16)
348 go func() {
349 defer close(out)
350 for {
351 select {
352 case <-done:
353 return
354 case event, ok := <-raw:
355 if !ok {
356 return
357 }
358 if !responseWebSocketEventMatches(event, responseID) {
359 continue
360 }
361 select {
362 case <-done:
363 return
364 case out <- event:
365 }
366 }
367 }
368 }()
369 return out
370}
371
372func responseWebSocketEventMatches(event sseEvent, responseID string) bool {
373 if event.Err != nil {
374 return true
375 }
376 var payload struct {
377 ResponseID string `json:"response_id"`
378 Response struct {
379 ID string `json:"id"`
380 } `json:"response"`
381 }
382 if json.Unmarshal([]byte(event.Data), &payload) != nil {
383 return true
384 }
385 if payload.ResponseID != "" && payload.ResponseID != responseID {
386 return false
387 }
388 return payload.Response.ID == "" || payload.Response.ID == responseID
389}
390
391func responseWebSocketRecoveryEvent(event sseEvent) bool {
392 if event.Err != nil {
393 return true
394 }
395 var payload struct {
396 Type string `json:"type"`
397 }
398 if json.Unmarshal([]byte(event.Data), &payload) != nil || payload.Type != "error" {
399 return false
400 }
401 return true
402}
403
404func (c *responsesWebSocketClient) forwardResponsesHTTP(ctx context.Context, req Request, out chan<- Event) {
405 stream, err := c.http.Stream(ctx, req)
406 if err != nil {
407 providerName := c.http.providerName
408 if providerName == "" {
409 providerName = "openai"
410 }
411 out <- EventStart{Model: req.Model, Provider: providerName}
412 out <- EventDone{Stop: StopError, Err: err}
413 return
414 }
415 for event := range stream {
416 out <- event
417 }
418}
419
420func (s *responsesWebSocketSession) incrementalRequest(req Request) (Request, string) {
421 s.mu.Lock()
422 defer s.mu.Unlock()
423 if s.lastResponseID == "" || len(req.Messages) < len(s.lastRequest) || !reflect.DeepEqual(req.Messages[:len(s.lastRequest)], s.lastRequest) {
424 return req, ""
425 }
426 // The cached response already contains the assistant output from the last
427 // generation. Send only the new tool results and/or user message.
428 req.Messages = append([]Message(nil), req.Messages[len(s.lastRequest):]...)
429 for len(req.Messages) > 0 && req.Messages[0].Role == RoleAssistant {
430 req.Messages = req.Messages[1:]
431 }
432 return req, s.lastResponseID
433}
434
435func (s *responsesWebSocketSession) ensureConnected(ctx context.Context, dial func(context.Context) (*websocket.Conn, error)) error {
436 s.mu.Lock()
437 if s.conn != nil && s.raw != nil {
438 s.mu.Unlock()
439 return nil
440 }
441 s.mu.Unlock()
442 conn, err := dial(ctx)
443 if err != nil {
444 return err
445 }
446 raw := make(chan sseEvent, 16)
447 readerDone := make(chan struct{})
448 s.mu.Lock()
449 s.conn = conn
450 s.raw = raw
451 s.readerDone = readerDone
452 s.mu.Unlock()
453 go readResponsesWebSocket(conn, raw, readerDone)
454 return nil
455}
456
457func (s *responsesWebSocketSession) connection() (*websocket.Conn, <-chan sseEvent) {
458 s.mu.Lock()
459 defer s.mu.Unlock()
460 return s.conn, s.raw
461}
462
463func (s *responsesWebSocketSession) recordCompleted(messages []Message, responseID string) {
464 s.mu.Lock()
465 defer s.mu.Unlock()
466 s.lastRequest = append(s.lastRequest[:0], messages...)
467 s.lastResponseID = responseID
468}
469
470func (s *responsesWebSocketSession) invalidateConnection(conn *websocket.Conn) {
471 s.mu.Lock()
472 defer s.mu.Unlock()
473 if s.conn != conn {
474 return
475 }
476 s.invalidateLocked()
477}
478
479func (s *responsesWebSocketSession) invalidate() {
480 s.mu.Lock()
481 defer s.mu.Unlock()
482 s.invalidateLocked()
483}
484
485func (s *responsesWebSocketSession) invalidateLocked() {
486 if s.readerDone != nil {
487 close(s.readerDone)
488 }
489 if s.conn != nil {
490 _ = s.conn.Close()
491 }
492 s.conn = nil
493 s.raw = nil
494 s.readerDone = nil
495 s.lastRequest = nil
496 s.lastResponseID = ""
497}
498
499func readResponsesWebSocket(conn *websocket.Conn, out chan<- sseEvent, done <-chan struct{}) {
500 defer close(out)
501 for {
502 var data json.RawMessage
503 if err := websocket.JSON.Receive(conn, &data); err != nil {
504 select {
505 case out <- sseEvent{Err: err}:
506 case <-done:
507 }
508 return
509 }
510 select {
511 case out <- sseEvent{Data: string(data)}:
512 case <-done:
513 return
514 }
515 }
516}