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