openai_responses_websocket_test.go RawHistoryFolder

1package provider
2
3import (
4	"context"
5	"encoding/json"
6	"io"
7	"net/http"
8	"net/http/httptest"
9	"reflect"
10	"strings"
11	"sync"
12	"testing"
13	"time"
14
15	"golang.org/x/net/websocket"
16)
17
18func TestResponsesWebSocketReconnectsMissingPreviousResponseWithFullContext(t *testing.T) {
19	var (
20		mu       sync.Mutex
21		requests []map[string]any
22		visits   int
23	)
24	ws := websocket.Handler(func(conn *websocket.Conn) {
25		defer conn.Close()
26		mu.Lock()
27		visits++
28		visit := visits
29		mu.Unlock()
30		if visit == 1 {
31			var first map[string]any
32			if err := websocket.JSON.Receive(conn, &first); err != nil {
33				return
34			}
35			mu.Lock()
36			requests = append(requests, first)
37			mu.Unlock()
38			for _, event := range []map[string]any{
39				{"type": "response.created", "response": map[string]any{"id": "resp_1"}},
40				{"type": "response.output_item.added", "output_index": 0, "item": map[string]any{"type": "message"}},
41				{"type": "response.output_text.delta", "output_index": 0, "delta": "ok"},
42				{"type": "response.completed", "response": map[string]any{"id": "resp_1", "usage": map[string]any{}}},
43			} {
44				_ = websocket.JSON.Send(conn, event)
45			}
46			var second map[string]any
47			if err := websocket.JSON.Receive(conn, &second); err != nil {
48				return
49			}
50			mu.Lock()
51			requests = append(requests, second)
52			mu.Unlock()
53			_ = websocket.JSON.Send(conn, map[string]any{
54				"type": "error", "status": 400,
55				"error": map[string]any{"code": "previous_response_not_found", "message": "evicted"},
56			})
57			return
58		}
59		var recovered map[string]any
60		if err := websocket.JSON.Receive(conn, &recovered); err != nil {
61			return
62		}
63		mu.Lock()
64		requests = append(requests, recovered)
65		mu.Unlock()
66		_ = websocket.JSON.Send(conn, map[string]any{"type": "response.created", "response": map[string]any{"id": "resp_2"}})
67		_ = websocket.JSON.Send(conn, map[string]any{"type": "response.output_item.added", "output_index": 0, "item": map[string]any{"type": "message"}})
68		_ = websocket.JSON.Send(conn, map[string]any{"type": "response.output_text.delta", "output_index": 0, "delta": "reconnected"})
69		_ = websocket.JSON.Send(conn, map[string]any{"type": "response.completed", "response": map[string]any{"id": "resp_2", "usage": map[string]any{}}})
70	})
71	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
72		if r.Header.Get("Upgrade") == "websocket" {
73			ws.ServeHTTP(w, r)
74			return
75		}
76		t.Errorf("unexpected HTTP fallback: %s %s", r.Method, r.URL)
77		w.WriteHeader(http.StatusInternalServerError)
78	}))
79	defer server.Close()
80
81	client := newResponsesWebSocketClient(&codexClient{
82		token: "test-token", baseURL: server.URL + "/v1/responses", providerName: "openai",
83		capabilities: responsesCapabilities{StablePromptCacheKey: true}, http: &http.Client{},
84	})
85	first := Request{Model: "gpt-5.6-sol", Context: RequestContext{CacheSessionID: "cache-root-1", ThreadID: "thread-1"}, Messages: []Message{{Role: RoleUser, Content: []Content{TextBlock{Text: "first"}}}}}
86	assertCompletedResponse(t, client, first)
87	lifecycle := &recordingRequestLifecycle{}
88	second := Request{Model: "gpt-5.6-sol", Context: RequestContext{CacheSessionID: "cache-root-1", ThreadID: "thread-1"}, Lifecycle: lifecycle, Messages: append(append([]Message(nil), first.Messages...), Message{Role: RoleAssistant, Content: []Content{TextBlock{Text: "ok"}}}, Message{Role: RoleUser, Content: []Content{TextBlock{Text: "second"}}})}
89	assertCompletedResponseText(t, client, second, "reconnected")
90	mu.Lock()
91	defer mu.Unlock()
92	if visits != 2 || len(requests) != 3 {
93		t.Fatalf("WebSocket visits/requests = %d/%d, want 2/3", visits, len(requests))
94	}
95	if len(lifecycle.ids) != 2 || lifecycle.ids[0] == lifecycle.ids[1] {
96		t.Fatalf("reconnect request IDs = %#v, want two distinct IDs", lifecycle.ids)
97	}
98	if requests[2]["previous_response_id"] != nil {
99		t.Fatalf("reconnect previous_response_id = %#v, want absent", requests[2]["previous_response_id"])
100	}
101	recovered, _ := json.Marshal(requests[2]["input"])
102	if !strings.Contains(string(recovered), "first") || !strings.Contains(string(recovered), "second") {
103		t.Fatalf("reconnect did not receive complete context: %s", recovered)
104	}
105}
106
107func TestResponsesWebSocketReusesSessionAndSendsIncrementalInput(t *testing.T) {
108	var (
109		mu       sync.Mutex
110		headers  []map[string]string
111		requests []map[string]any
112	)
113	server := httptest.NewServer(websocket.Handler(func(conn *websocket.Conn) {
114		defer conn.Close()
115		for i := 0; i < 2; i++ {
116			var payload map[string]any
117			if err := websocket.JSON.Receive(conn, &payload); err != nil {
118				t.Errorf("receive websocket request: %v", err)
119				return
120			}
121			mu.Lock()
122			requests = append(requests, payload)
123			headers = append(headers, map[string]string{
124				"authorization":      conn.Request().Header.Get("authorization"),
125				"chatgpt-account-id": conn.Request().Header.Get("chatgpt-account-id"),
126				"openai-beta":        conn.Request().Header.Get("openai-beta"),
127			})
128			mu.Unlock()
129			_ = websocket.JSON.Send(conn, map[string]any{
130				"type": "response.created",
131				"response": map[string]any{
132					"id": "resp_" + string(rune('1'+i)),
133				},
134			})
135			_ = websocket.JSON.Send(conn, map[string]any{
136				"type":         "response.output_item.added",
137				"output_index": 0,
138				"item":         map[string]any{"type": "message"},
139			})
140			_ = websocket.JSON.Send(conn, map[string]any{
141				"type":         "response.output_text.delta",
142				"output_index": 0,
143				"delta":        "ok",
144			})
145			_ = websocket.JSON.Send(conn, map[string]any{
146				"type": "response.completed",
147				"response": map[string]any{
148					"id":    "resp_" + string(rune('1'+i)),
149					"usage": map[string]any{"input_tokens": 1, "output_tokens": 1},
150				},
151			})
152		}
153	}))
154	defer server.Close()
155
156	inner := &codexClient{
157		token:        "test-token",
158		baseURL:      server.URL + "/v1/responses",
159		providerName: "openai",
160		capabilities: responsesCapabilities{StablePromptCacheKey: true},
161	}
162	client := newResponsesWebSocketClient(inner)
163
164	lifecycle := &recordingRequestLifecycle{}
165	first := Request{
166		Model:     "gpt-5.6-sol",
167		FastMode:  true,
168		Context:   RequestContext{CacheSessionID: "cache-root-1", ThreadID: "thread-1", TurnID: "turn-1"},
169		Lifecycle: lifecycle,
170		Messages:  []Message{{Role: RoleUser, Content: []Content{TextBlock{Text: "first"}}}},
171	}
172	assertCompletedResponse(t, client, first)
173	second := Request{
174		Model:    "gpt-5.6-sol",
175		FastMode: true,
176		Context:  RequestContext{CacheSessionID: "cache-root-1", ThreadID: "thread-1", TurnID: "turn-2"},
177		Messages: append(append([]Message(nil), first.Messages...),
178			Message{Role: RoleAssistant, Content: []Content{TextBlock{Text: "ok"}}},
179			Message{Role: RoleUser, Content: []Content{TextBlock{Text: "second"}}},
180		),
181	}
182	assertCompletedResponse(t, client, second)
183	if len(lifecycle.ids) != 1 || lifecycle.ids[0] == "" {
184		t.Fatalf("WebSocket request IDs = %#v", lifecycle.ids)
185	}
186
187	mu.Lock()
188	defer mu.Unlock()
189	if len(requests) != 2 {
190		t.Fatalf("websocket requests = %d, want 2", len(requests))
191	}
192	if got := requests[0]["type"]; got != "response.create" {
193		t.Fatalf("first type = %#v", got)
194	}
195	if got := requests[0]["store"]; got != false {
196		t.Fatalf("store = %#v, want false", got)
197	}
198	if _, ok := requests[0]["stream"]; ok {
199		t.Fatalf("WebSocket payload unexpectedly carries stream: %#v", requests[0])
200	}
201	if got := requests[0]["prompt_cache_key"]; got != "cache-root-1" {
202		t.Fatalf("prompt_cache_key = %#v", got)
203	}
204	if got := requests[0]["service_tier"]; got != fastModeServiceTier {
205		t.Fatalf("service_tier = %#v", got)
206	}
207	if got := requests[1]["previous_response_id"]; got != "resp_1" {
208		t.Fatalf("previous_response_id = %#v", got)
209	}
210	input, ok := requests[1]["input"].([]any)
211	if !ok || len(input) != 1 {
212		t.Fatalf("incremental input = %#v", requests[1]["input"])
213	}
214	encoded, _ := json.Marshal(input[0])
215	if !strings.Contains(string(encoded), "second") || strings.Contains(string(encoded), "ok") {
216		t.Fatalf("incremental input = %s", encoded)
217	}
218	for _, got := range headers {
219		if got["authorization"] != "Bearer test-token" {
220			t.Fatalf("authorization = %q", got["authorization"])
221		}
222		if got["chatgpt-account-id"] != "" || got["openai-beta"] != "" {
223			t.Fatalf("public WebSocket leaked Codex headers: %#v", got)
224		}
225	}
226}
227
228func TestOpenAICodexUsesWebSocketWithSubscriptionHeaders(t *testing.T) {
229	headers := make(chan http.Header, 1)
230	requests := make(chan map[string]any, 1)
231	server := httptest.NewServer(websocket.Handler(func(conn *websocket.Conn) {
232		defer conn.Close()
233		headers <- conn.Request().Header.Clone()
234		var payload map[string]any
235		if err := websocket.JSON.Receive(conn, &payload); err != nil {
236			t.Errorf("receive websocket request: %v", err)
237			return
238		}
239		requests <- payload
240		for _, event := range []map[string]any{
241			{"type": "response.created", "response": map[string]any{"id": "resp_subscription"}},
242			{"type": "response.output_item.added", "output_index": 0, "item": map[string]any{"type": "message"}},
243			{"type": "response.output_text.delta", "output_index": 0, "delta": "ok"},
244			{"type": "response.completed", "response": map[string]any{"id": "resp_subscription", "usage": map[string]any{}}},
245		} {
246			_ = websocket.JSON.Send(conn, event)
247		}
248	}))
249	defer server.Close()
250
251	client := NewOpenAICodex("subscription-token", "account-1", server.URL)
252	websocketClient, ok := client.(*responsesWebSocketClient)
253	if !ok {
254		t.Fatalf("NewOpenAICodex returned %T, want WebSocket Responses client", client)
255	}
256	defer websocketClient.Close()
257
258	request := Request{
259		Model:    "gpt-5.6-luna",
260		Context:  RequestContext{CacheSessionID: "root-1", ThreadID: "thread-1", TurnID: "turn-1"},
261		Messages: []Message{{Role: RoleUser, Content: []Content{TextBlock{Text: "hello"}}}},
262	}
263	assertCompletedResponse(t, client, request)
264
265	gotHeaders := <-headers
266	for name, want := range map[string]string{
267		"authorization":       "Bearer subscription-token",
268		"chatgpt-account-id":  "account-1",
269		"openai-beta":         "responses=experimental",
270		"originator":          "codex_cli_rs",
271		"user-agent":          "codex_cli_rs/0.0.0",
272		"session-id":          "root-1",
273		"thread-id":           "thread-1",
274		"x-client-request-id": "thread-1",
275	} {
276		if got := gotHeaders.Get(name); got != want {
277			t.Fatalf("%s = %q, want %q", name, got, want)
278		}
279	}
280	if got := gotHeaders.Get("origin"); got != server.URL {
281		t.Fatalf("origin = %q, want %q", got, server.URL)
282	}
283	if got := (<-requests)["prompt_cache_key"]; got != "root-1" {
284		t.Fatalf("prompt_cache_key = %#v, want root cache ID", got)
285	}
286}
287
288func TestWithHTTPClientPreservesSubscriptionHeaders(t *testing.T) {
289	client := NewOpenAICodex("subscription-token", "account-1", "https://example.test/backend-api/codex/responses")
290	customHTTP := &http.Client{Transport: &http.Transport{}}
291	WithHTTPClient(client, customHTTP)
292
293	websocketClient := client.(*responsesWebSocketClient)
294	if websocketClient.http.http != customHTTP {
295		t.Fatal("subscription client was wrapped in the public OpenAI header-stripping transport")
296	}
297}
298
299func TestResponsesWebSocketWireBaselineChangeForcesFullRequest(t *testing.T) {
300	mappedModel := "gpt-5.6-sol-a"
301	client := newResponsesWebSocketClient(&codexClient{
302		providerName: "openai",
303		capabilities: responsesCapabilities{StablePromptCacheKey: true},
304		modelName:    func(string) string { return mappedModel },
305	}).(*responsesWebSocketClient)
306	first := Request{
307		Model:    "gpt-5.6-sol",
308		Context:  RequestContext{CacheSessionID: "cache-root-1", ThreadID: "thread-1", TurnID: "turn-1"},
309		Messages: []Message{{Role: RoleUser, Content: []Content{TextBlock{Text: "first"}}}},
310	}
311	baseline, err := client.canonicalResponsesBaseline(first, first.Context.CacheSessionID)
312	if err != nil {
313		t.Fatal(err)
314	}
315	session := &responsesWebSocketSession{}
316	session.recordCompleted(first, baseline, Message{Role: RoleAssistant, Content: []Content{TextBlock{Text: "ok"}}}, "resp_1")
317
318	mappedModel = "gpt-5.6-sol-b"
319	second := Request{
320		Model:   first.Model,
321		Context: RequestContext{CacheSessionID: "cache-root-1", ThreadID: "thread-1", TurnID: "turn-2"},
322		Messages: append(append([]Message(nil), first.Messages...),
323			Message{Role: RoleAssistant, Content: []Content{TextBlock{Text: "ok"}}},
324			Message{Role: RoleUser, Content: []Content{TextBlock{Text: "second"}}},
325		),
326	}
327	streamReq, previousResponseID := session.incrementalRequest(second, second.Context.CacheSessionID, client.canonicalResponsesBaseline)
328	if previousResponseID != "" {
329		t.Fatalf("previous response ID = %q, want full request after wire baseline change", previousResponseID)
330	}
331	if !reflect.DeepEqual(streamReq.Messages, second.Messages) {
332		t.Fatalf("stream messages = %#v, want complete transcript %#v", streamReq.Messages, second.Messages)
333	}
334}
335
336func TestResponsesWebSocketAssistantMismatchForcesCompleteTranscript(t *testing.T) {
337	client := newResponsesWebSocketClient(&codexClient{
338		providerName: "openai", capabilities: responsesCapabilities{StablePromptCacheKey: true},
339	}).(*responsesWebSocketClient)
340	first := Request{
341		Model: "gpt-5.6-sol", Context: RequestContext{CacheSessionID: "cache-root-1", ThreadID: "thread-1", TurnID: "turn-1"},
342		Messages: []Message{{Role: RoleUser, Content: []Content{TextBlock{Text: "first"}}}},
343	}
344	baseline, err := client.canonicalResponsesBaseline(first, first.Context.CacheSessionID)
345	if err != nil {
346		t.Fatal(err)
347	}
348	session := &responsesWebSocketSession{}
349	session.recordCompleted(first, baseline, Message{Role: RoleAssistant, Content: []Content{TextBlock{Text: "actual"}}}, "resp_1")
350	second := Request{
351		Model: first.Model, Context: RequestContext{CacheSessionID: "cache-root-1", ThreadID: "thread-1", TurnID: "turn-2"},
352		Messages: append(append([]Message(nil), first.Messages...),
353			Message{Role: RoleAssistant, Content: []Content{TextBlock{Text: "modified"}}},
354			Message{Role: RoleUser, Content: []Content{TextBlock{Text: "second"}}},
355		),
356	}
357	streamReq, previousResponseID := session.incrementalRequest(second, second.Context.CacheSessionID, client.canonicalResponsesBaseline)
358	if previousResponseID != "" {
359		t.Fatalf("previous response ID = %q, want complete request", previousResponseID)
360	}
361	if !reflect.DeepEqual(streamReq.Messages, second.Messages) {
362		t.Fatalf("stream messages = %#v, want complete transcript %#v", streamReq.Messages, second.Messages)
363	}
364}
365
366func TestResponsesWebSocketCancellationClosesOnlyActiveConnection(t *testing.T) {
367	closed := make(chan struct{})
368	var closeOnce sync.Once
369	server := httptest.NewServer(websocket.Handler(func(conn *websocket.Conn) {
370		defer closeOnce.Do(func() { close(closed) })
371		defer conn.Close()
372		var payload map[string]any
373		if err := websocket.JSON.Receive(conn, &payload); err != nil {
374			return
375		}
376		var ignored any
377		_ = websocket.JSON.Receive(conn, &ignored)
378	}))
379	defer server.Close()
380
381	client := newResponsesWebSocketClient(&codexClient{
382		token: "test-token", baseURL: server.URL + "/v1/responses", providerName: "openai",
383		capabilities: responsesCapabilities{StablePromptCacheKey: true},
384	})
385	ctx, cancel := context.WithCancel(context.Background())
386	stream, err := client.Stream(ctx, Request{Model: "gpt-5.6-sol", Context: RequestContext{CacheSessionID: "cache-root-1", ThreadID: "thread-1"}})
387	if err != nil {
388		t.Fatal(err)
389	}
390	cancel()
391	for range stream {
392	}
393	select {
394	case <-closed:
395	case <-time.After(time.Second):
396		t.Fatal("cancellation did not close active WebSocket")
397	}
398}
399
400func TestResponsesWebSocketIgnoresStaleTerminalFrameBeforeNextResponse(t *testing.T) {
401	server := httptest.NewServer(websocket.Handler(func(conn *websocket.Conn) {
402		defer conn.Close()
403		for turn := 1; turn <= 2; turn++ {
404			var payload map[string]any
405			if err := websocket.JSON.Receive(conn, &payload); err != nil {
406				return
407			}
408			if turn == 2 {
409				// A buffered terminal from the prior response must not complete
410				// this response before its response.created frame arrives.
411				_ = websocket.JSON.Send(conn, map[string]any{"type": "response.completed", "response": map[string]any{"id": "resp_1", "usage": map[string]any{}}})
412			}
413			id := "resp_" + string(rune('0'+turn))
414			text := "first"
415			if turn == 2 {
416				text = "second"
417			}
418			for _, event := range []map[string]any{
419				{"type": "response.created", "response": map[string]any{"id": id}},
420				{"type": "response.output_item.added", "output_index": 0, "item": map[string]any{"type": "message"}},
421				{"type": "response.output_text.delta", "output_index": 0, "delta": text},
422				{"type": "response.completed", "response": map[string]any{"id": id, "usage": map[string]any{}}},
423			} {
424				_ = websocket.JSON.Send(conn, event)
425			}
426		}
427	}))
428	defer server.Close()
429
430	client := newResponsesWebSocketClient(&codexClient{
431		token: "test-token", baseURL: server.URL + "/v1/responses", providerName: "openai",
432		capabilities: responsesCapabilities{StablePromptCacheKey: true}, http: &http.Client{},
433	})
434	first := Request{Model: "gpt-5.6-sol", Context: RequestContext{CacheSessionID: "cache-root-1", ThreadID: "thread-1"}, Messages: []Message{{Role: RoleUser, Content: []Content{TextBlock{Text: "first"}}}}}
435	assertCompletedResponseText(t, client, first, "first")
436	second := Request{Model: "gpt-5.6-sol", Context: RequestContext{CacheSessionID: "cache-root-1", ThreadID: "thread-1"}, Messages: append(append([]Message(nil), first.Messages...), Message{Role: RoleAssistant, Content: []Content{TextBlock{Text: "first"}}}, Message{Role: RoleUser, Content: []Content{TextBlock{Text: "second"}}})}
437	assertCompletedResponseText(t, client, second, "second")
438}
439
440func TestResponsesWebSocketCloseReleasesSessionSockets(t *testing.T) {
441	client := newResponsesWebSocketClient(&codexClient{}).(*responsesWebSocketClient)
442	client.session("resident-child")
443	if err := client.Close(); err != nil {
444		t.Fatalf("close responses websocket client: %v", err)
445	}
446	client.mu.Lock()
447	defer client.mu.Unlock()
448	if len(client.sessions) != 0 {
449		t.Fatalf("sessions retained after close: %d", len(client.sessions))
450	}
451}
452
453type cacheDiagnosticsLifecycle struct {
454	diagnostics []CacheDiagnostics
455}
456
457func (l *cacheDiagnosticsLifecycle) RequestAttempt(int, int)                {}
458func (l *cacheDiagnosticsLifecycle) RetryScheduled(int, int, time.Duration) {}
459
460func (l *cacheDiagnosticsLifecycle) CacheDiagnostics(diagnostics CacheDiagnostics) {
461	l.diagnostics = append(l.diagnostics, diagnostics)
462}
463
464func TestResponsesWebSocketUsesHTTPDiagnosticsForIncompatibleTransport(t *testing.T) {
465	lifecycle := &cacheDiagnosticsLifecycle{}
466	client := newResponsesWebSocketClient(&codexClient{
467		token: "test-token", baseURL: "https://example.test/v1/responses", providerName: "openai",
468		capabilities: responsesCapabilities{StablePromptCacheKey: true},
469		http: &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
470			return &http.Response{
471				StatusCode: http.StatusOK,
472				Header:     make(http.Header),
473				Body: io.NopCloser(strings.NewReader(`data: {"type":"response.completed","response":{"status":"completed","usage":{}}}
474
475`)),
476			}, nil
477		})},
478	})
479	req := Request{
480		Model: "gpt-5.6-sol", Context: RequestContext{CacheSessionID: "cache-root", ThreadID: "thread"}, Lifecycle: lifecycle,
481		Messages: []Message{{Role: RoleUser, Content: []Content{TextBlock{Text: "request"}}}},
482	}
483	stream, err := client.Stream(context.Background(), req)
484	if err != nil {
485		t.Fatal(err)
486	}
487	for range stream {
488	}
489	if len(lifecycle.diagnostics) != 1 || lifecycle.diagnostics[0].Transport != "http_sse" {
490		t.Fatalf("diagnostics = %#v, want only HTTP/SSE diagnostics", lifecycle.diagnostics)
491	}
492}
493
494func assertCompletedResponse(t *testing.T, client Client, req Request) {
495	assertCompletedResponseText(t, client, req, "ok")
496}
497
498func assertCompletedResponseText(t *testing.T, client Client, req Request, want string) {
499	t.Helper()
500	stream, err := client.Stream(context.Background(), req)
501	if err != nil {
502		t.Fatal(err)
503	}
504	var done EventDone
505	var got []Event
506	for event := range stream {
507		got = append(got, event)
508		if got, ok := event.(EventDone); ok {
509			done = got
510		}
511	}
512	if done.Err != nil || len(done.Message.Content) != 1 || done.Message.Content[0].(TextBlock).Text != want {
513		t.Fatalf("events = %#v; done = %#v", got, done)
514	}
515}