openai_responses_websocket_test.go RawHistoryFolder

1package provider
2
3import (
4	"context"
5	"encoding/json"
6	"net/http"
7	"net/http/httptest"
8	"strings"
9	"sync"
10	"testing"
11	"time"
12
13	"golang.org/x/net/websocket"
14)
15
16func TestResponsesWebSocketReconnectsMissingPreviousResponseWithFullContext(t *testing.T) {
17	var (
18		mu       sync.Mutex
19		requests []map[string]any
20		visits   int
21	)
22	ws := websocket.Handler(func(conn *websocket.Conn) {
23		defer conn.Close()
24		mu.Lock()
25		visits++
26		visit := visits
27		mu.Unlock()
28		if visit == 1 {
29			var first map[string]any
30			if err := websocket.JSON.Receive(conn, &first); err != nil {
31				return
32			}
33			mu.Lock()
34			requests = append(requests, first)
35			mu.Unlock()
36			for _, event := range []map[string]any{
37				{"type": "response.created", "response": map[string]any{"id": "resp_1"}},
38				{"type": "response.output_item.added", "output_index": 0, "item": map[string]any{"type": "message"}},
39				{"type": "response.output_text.delta", "output_index": 0, "delta": "ok"},
40				{"type": "response.completed", "response": map[string]any{"id": "resp_1", "usage": map[string]any{}}},
41			} {
42				_ = websocket.JSON.Send(conn, event)
43			}
44			var second map[string]any
45			if err := websocket.JSON.Receive(conn, &second); err != nil {
46				return
47			}
48			mu.Lock()
49			requests = append(requests, second)
50			mu.Unlock()
51			_ = websocket.JSON.Send(conn, map[string]any{
52				"type": "error", "status": 400,
53				"error": map[string]any{"code": "previous_response_not_found", "message": "evicted"},
54			})
55			return
56		}
57		var recovered map[string]any
58		if err := websocket.JSON.Receive(conn, &recovered); err != nil {
59			return
60		}
61		mu.Lock()
62		requests = append(requests, recovered)
63		mu.Unlock()
64		_ = websocket.JSON.Send(conn, map[string]any{"type": "response.created", "response": map[string]any{"id": "resp_2"}})
65		_ = websocket.JSON.Send(conn, map[string]any{"type": "response.output_item.added", "output_index": 0, "item": map[string]any{"type": "message"}})
66		_ = websocket.JSON.Send(conn, map[string]any{"type": "response.output_text.delta", "output_index": 0, "delta": "reconnected"})
67		_ = websocket.JSON.Send(conn, map[string]any{"type": "response.completed", "response": map[string]any{"id": "resp_2", "usage": map[string]any{}}})
68	})
69	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
70		if r.Header.Get("Upgrade") == "websocket" {
71			ws.ServeHTTP(w, r)
72			return
73		}
74		t.Errorf("unexpected HTTP fallback: %s %s", r.Method, r.URL)
75		w.WriteHeader(http.StatusInternalServerError)
76	}))
77	defer server.Close()
78
79	client := newResponsesWebSocketClient(&codexClient{
80		token: "test-token", baseURL: server.URL + "/v1/responses", providerName: "openai",
81		capabilities: responsesCapabilities{StablePromptCacheKey: true}, http: &http.Client{},
82	})
83	first := Request{Model: "gpt-5.6-sol", Context: RequestContext{SessionID: "session-1"}, Messages: []Message{{Role: RoleUser, Content: []Content{TextBlock{Text: "first"}}}}}
84	assertCompletedResponse(t, client, first)
85	lifecycle := &recordingRequestLifecycle{}
86	second := Request{Model: "gpt-5.6-sol", Context: RequestContext{SessionID: "session-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"}}})}
87	assertCompletedResponseText(t, client, second, "reconnected")
88	mu.Lock()
89	defer mu.Unlock()
90	if visits != 2 || len(requests) != 3 {
91		t.Fatalf("WebSocket visits/requests = %d/%d, want 2/3", visits, len(requests))
92	}
93	if len(lifecycle.ids) != 2 || lifecycle.ids[0] == lifecycle.ids[1] {
94		t.Fatalf("reconnect request IDs = %#v, want two distinct IDs", lifecycle.ids)
95	}
96	if requests[2]["previous_response_id"] != nil {
97		t.Fatalf("reconnect previous_response_id = %#v, want absent", requests[2]["previous_response_id"])
98	}
99	recovered, _ := json.Marshal(requests[2]["input"])
100	if !strings.Contains(string(recovered), "first") || !strings.Contains(string(recovered), "second") {
101		t.Fatalf("reconnect did not receive complete context: %s", recovered)
102	}
103}
104
105func TestResponsesWebSocketReusesSessionAndSendsIncrementalInput(t *testing.T) {
106	var (
107		mu       sync.Mutex
108		headers  []map[string]string
109		requests []map[string]any
110	)
111	server := httptest.NewServer(websocket.Handler(func(conn *websocket.Conn) {
112		defer conn.Close()
113		for i := 0; i < 2; i++ {
114			var payload map[string]any
115			if err := websocket.JSON.Receive(conn, &payload); err != nil {
116				t.Errorf("receive websocket request: %v", err)
117				return
118			}
119			mu.Lock()
120			requests = append(requests, payload)
121			headers = append(headers, map[string]string{
122				"authorization":      conn.Request().Header.Get("authorization"),
123				"chatgpt-account-id": conn.Request().Header.Get("chatgpt-account-id"),
124				"openai-beta":        conn.Request().Header.Get("openai-beta"),
125			})
126			mu.Unlock()
127			_ = websocket.JSON.Send(conn, map[string]any{
128				"type": "response.created",
129				"response": map[string]any{
130					"id": "resp_" + string(rune('1'+i)),
131				},
132			})
133			_ = websocket.JSON.Send(conn, map[string]any{
134				"type":         "response.output_item.added",
135				"output_index": 0,
136				"item":         map[string]any{"type": "message"},
137			})
138			_ = websocket.JSON.Send(conn, map[string]any{
139				"type":         "response.output_text.delta",
140				"output_index": 0,
141				"delta":        "ok",
142			})
143			_ = websocket.JSON.Send(conn, map[string]any{
144				"type": "response.completed",
145				"response": map[string]any{
146					"id":    "resp_" + string(rune('1'+i)),
147					"usage": map[string]any{"input_tokens": 1, "output_tokens": 1},
148				},
149			})
150		}
151	}))
152	defer server.Close()
153
154	inner := &codexClient{
155		token:        "test-token",
156		baseURL:      server.URL + "/v1/responses",
157		providerName: "openai",
158		capabilities: responsesCapabilities{StablePromptCacheKey: true},
159	}
160	client := newResponsesWebSocketClient(inner)
161
162	lifecycle := &recordingRequestLifecycle{}
163	first := Request{
164		Model:     "gpt-5.6-sol",
165		FastMode:  true,
166		Context:   RequestContext{SessionID: "session-1", TurnID: "turn-1"},
167		Lifecycle: lifecycle,
168		Messages:  []Message{{Role: RoleUser, Content: []Content{TextBlock{Text: "first"}}}},
169	}
170	assertCompletedResponse(t, client, first)
171	second := Request{
172		Model:   "gpt-5.6-sol",
173		Context: RequestContext{SessionID: "session-1", TurnID: "turn-2"},
174		Messages: append(append([]Message(nil), first.Messages...),
175			Message{Role: RoleAssistant, Content: []Content{TextBlock{Text: "ok"}}},
176			Message{Role: RoleUser, Content: []Content{TextBlock{Text: "second"}}},
177		),
178	}
179	assertCompletedResponse(t, client, second)
180	if len(lifecycle.ids) != 1 || lifecycle.ids[0] == "" {
181		t.Fatalf("WebSocket request IDs = %#v", lifecycle.ids)
182	}
183
184	mu.Lock()
185	defer mu.Unlock()
186	if len(requests) != 2 {
187		t.Fatalf("websocket requests = %d, want 2", len(requests))
188	}
189	if got := requests[0]["type"]; got != "response.create" {
190		t.Fatalf("first type = %#v", got)
191	}
192	if got := requests[0]["store"]; got != false {
193		t.Fatalf("store = %#v, want false", got)
194	}
195	if _, ok := requests[0]["stream"]; ok {
196		t.Fatalf("WebSocket payload unexpectedly carries stream: %#v", requests[0])
197	}
198	if got := requests[0]["prompt_cache_key"]; got != "session-1" {
199		t.Fatalf("prompt_cache_key = %#v", got)
200	}
201	if got := requests[0]["service_tier"]; got != fastModeServiceTier {
202		t.Fatalf("service_tier = %#v", got)
203	}
204	if got := requests[1]["previous_response_id"]; got != "resp_1" {
205		t.Fatalf("previous_response_id = %#v", got)
206	}
207	input, ok := requests[1]["input"].([]any)
208	if !ok || len(input) != 1 {
209		t.Fatalf("incremental input = %#v", requests[1]["input"])
210	}
211	encoded, _ := json.Marshal(input[0])
212	if !strings.Contains(string(encoded), "second") || strings.Contains(string(encoded), "ok") {
213		t.Fatalf("incremental input = %s", encoded)
214	}
215	for _, got := range headers {
216		if got["authorization"] != "Bearer test-token" {
217			t.Fatalf("authorization = %q", got["authorization"])
218		}
219		if got["chatgpt-account-id"] != "" || got["openai-beta"] != "" {
220			t.Fatalf("public WebSocket leaked Codex headers: %#v", got)
221		}
222	}
223}
224
225func TestResponsesWebSocketCancellationClosesOnlyActiveConnection(t *testing.T) {
226	closed := make(chan struct{})
227	var closeOnce sync.Once
228	server := httptest.NewServer(websocket.Handler(func(conn *websocket.Conn) {
229		defer closeOnce.Do(func() { close(closed) })
230		defer conn.Close()
231		var payload map[string]any
232		if err := websocket.JSON.Receive(conn, &payload); err != nil {
233			return
234		}
235		var ignored any
236		_ = websocket.JSON.Receive(conn, &ignored)
237	}))
238	defer server.Close()
239
240	client := newResponsesWebSocketClient(&codexClient{
241		token: "test-token", baseURL: server.URL + "/v1/responses", providerName: "openai",
242		capabilities: responsesCapabilities{StablePromptCacheKey: true},
243	})
244	ctx, cancel := context.WithCancel(context.Background())
245	stream, err := client.Stream(ctx, Request{Model: "gpt-5.6-sol", Context: RequestContext{SessionID: "session-1"}})
246	if err != nil {
247		t.Fatal(err)
248	}
249	cancel()
250	for range stream {
251	}
252	select {
253	case <-closed:
254	case <-time.After(time.Second):
255		t.Fatal("cancellation did not close active WebSocket")
256	}
257}
258
259func TestResponsesWebSocketIgnoresStaleTerminalFrameBeforeNextResponse(t *testing.T) {
260	server := httptest.NewServer(websocket.Handler(func(conn *websocket.Conn) {
261		defer conn.Close()
262		for turn := 1; turn <= 2; turn++ {
263			var payload map[string]any
264			if err := websocket.JSON.Receive(conn, &payload); err != nil {
265				return
266			}
267			if turn == 2 {
268				// A buffered terminal from the prior response must not complete
269				// this response before its response.created frame arrives.
270				_ = websocket.JSON.Send(conn, map[string]any{"type": "response.completed", "response": map[string]any{"id": "resp_1", "usage": map[string]any{}}})
271			}
272			id := "resp_" + string(rune('0'+turn))
273			text := "first"
274			if turn == 2 {
275				text = "second"
276			}
277			for _, event := range []map[string]any{
278				{"type": "response.created", "response": map[string]any{"id": id}},
279				{"type": "response.output_item.added", "output_index": 0, "item": map[string]any{"type": "message"}},
280				{"type": "response.output_text.delta", "output_index": 0, "delta": text},
281				{"type": "response.completed", "response": map[string]any{"id": id, "usage": map[string]any{}}},
282			} {
283				_ = websocket.JSON.Send(conn, event)
284			}
285		}
286	}))
287	defer server.Close()
288
289	client := newResponsesWebSocketClient(&codexClient{
290		token: "test-token", baseURL: server.URL + "/v1/responses", providerName: "openai",
291		capabilities: responsesCapabilities{StablePromptCacheKey: true}, http: &http.Client{},
292	})
293	first := Request{Model: "gpt-5.6-sol", Context: RequestContext{SessionID: "session-1"}, Messages: []Message{{Role: RoleUser, Content: []Content{TextBlock{Text: "first"}}}}}
294	assertCompletedResponseText(t, client, first, "first")
295	second := Request{Model: "gpt-5.6-sol", Context: RequestContext{SessionID: "session-1"}, Messages: append(append([]Message(nil), first.Messages...), Message{Role: RoleAssistant, Content: []Content{TextBlock{Text: "first"}}}, Message{Role: RoleUser, Content: []Content{TextBlock{Text: "second"}}})}
296	assertCompletedResponseText(t, client, second, "second")
297}
298
299func TestResponsesWebSocketCloseReleasesSessionSockets(t *testing.T) {
300	client := newResponsesWebSocketClient(&codexClient{}).(*responsesWebSocketClient)
301	client.session("resident-child")
302	if err := client.Close(); err != nil {
303		t.Fatalf("close responses websocket client: %v", err)
304	}
305	client.mu.Lock()
306	defer client.mu.Unlock()
307	if len(client.sessions) != 0 {
308		t.Fatalf("sessions retained after close: %d", len(client.sessions))
309	}
310}
311
312func assertCompletedResponse(t *testing.T, client Client, req Request) {
313	assertCompletedResponseText(t, client, req, "ok")
314}
315
316func assertCompletedResponseText(t *testing.T, client Client, req Request, want string) {
317	t.Helper()
318	stream, err := client.Stream(context.Background(), req)
319	if err != nil {
320		t.Fatal(err)
321	}
322	var done EventDone
323	var got []Event
324	for event := range stream {
325		got = append(got, event)
326		if got, ok := event.(EventDone); ok {
327			done = got
328		}
329	}
330	if done.Err != nil || len(done.Message.Content) != 1 || done.Message.Content[0].(TextBlock).Text != want {
331		t.Fatalf("events = %#v; done = %#v", got, done)
332	}
333}