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