Skip to content

test(executor): fix two websocket test-server races behind CI flakes (MASTRA-771) #6

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ericlitman merged 1 commit into from
Oct 3, 2026
Merged

test(executor): fix two websocket test-server races behind CI flakes (MASTRA-771) #6

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 18 additions & 18 deletions internal/runtime/executor/codex_websockets_executor_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2667,15 +2667,11 @@ func TestCodexWebsockets_KeepalivePingDuringUpload_WithSession(t *testing.T) {
return nil
})

// Start server reader loop so server processes control frames.
readErrCh := make(chan error, 1)
// The reader answers control frames until the client request arrives.
requestRead := make(chan struct{})
go func() {
for {
if _, _, errRead := conn.ReadMessage(); errRead != nil {
readErrCh <- errRead
return
}
}
_, _, _ = conn.ReadMessage()
close(requestRead)
}()

// Wait until client has entered writeMessage and is actively holding writeMu.
Expand All @@ -2701,6 +2697,8 @@ func TestCodexWebsockets_KeepalivePingDuringUpload_WithSession(t *testing.T) {
return
}

// Reply only after the request arrives, so the reply and close cannot race the held write.
<-requestRead
// Now send terminal response.
respPayload := []byte(`{"type":"response.completed","response":{"id":"resp-1","status":"completed","output":[]}}`)
_ = conn.WriteMessage(websocket.TextMessage, respPayload)
Expand Down Expand Up @@ -2766,12 +2764,11 @@ func TestCodexWebsockets_KeepalivePingDuringUpload_Sessionless(t *testing.T) {
return nil
})

// The reader answers control frames until the client request arrives.
requestRead := make(chan struct{})
go func() {
for {
if _, _, errRead := conn.ReadMessage(); errRead != nil {
return
}
}
_, _, _ = conn.ReadMessage()
close(requestRead)
}()

// Wait until client has entered writeMessage on sessionless path.
Expand All @@ -2797,6 +2794,8 @@ func TestCodexWebsockets_KeepalivePingDuringUpload_Sessionless(t *testing.T) {
return
}

// Reply only after the request arrives, so the reply and close cannot race the held write.
<-requestRead
respPayload := []byte(`{"type":"response.completed","response":{"id":"resp-1","status":"completed","output":[]}}`)
_ = conn.WriteMessage(websocket.TextMessage, respPayload)
}))
Expand Down Expand Up @@ -2858,12 +2857,11 @@ func TestCodexWebsockets_KeepalivePingDuringUpload_NonstreamSessionless(t *testi
return nil
})

// The reader answers control frames until the client request arrives.
requestRead := make(chan struct{})
go func() {
for {
if _, _, errRead := conn.ReadMessage(); errRead != nil {
return
}
}
_, _, _ = conn.ReadMessage()
close(requestRead)
}()

// Wait until client has entered writeMessage on nonstream path.
Expand All @@ -2889,6 +2887,8 @@ func TestCodexWebsockets_KeepalivePingDuringUpload_NonstreamSessionless(t *testi
return
}

// Reply only after the request arrives, so the reply and close cannot race the held write.
<-requestRead
respPayload := []byte(`{"type":"response.completed","response":{"id":"resp-1","status":"completed","output":[]}}`)
_ = conn.WriteMessage(websocket.TextMessage, respPayload)
}))
Expand Down
3 changes: 2 additions & 1 deletion internal/runtime/executor/websocket_session_target_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -232,12 +232,13 @@ func TestWebsocketRetryBindFailureClearsActiveSessionState(t *testing.T) {
upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }}
var connections atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
// Count before the handshake completes so numbers follow dial order.
connection := connections.Add(1)
conn, errUpgrade := upgrader.Upgrade(w, r, nil)
if errUpgrade != nil {
t.Errorf("upgrade websocket: %v", errUpgrade)
return
}
connection := connections.Add(1)
defer func() { _ = conn.Close() }()
if connection == 1 {
_, _, _ = conn.ReadMessage()
Expand Down
22 changes: 12 additions & 10 deletions internal/runtime/executor/xai_websockets_executor_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2214,12 +2214,11 @@ func TestXAIWebsockets_KeepalivePingDuringUpload_WithSession(t *testing.T) {
return nil
})

// The reader answers control frames until the client request arrives.
requestRead := make(chan struct{})
go func() {
for {
if _, _, errReadLoop := conn.ReadMessage(); errReadLoop != nil {
return
}
}
_, _, _ = conn.ReadMessage()
close(requestRead)
}()

select {
Expand All @@ -2242,6 +2241,8 @@ func TestXAIWebsockets_KeepalivePingDuringUpload_WithSession(t *testing.T) {
return
}

// Reply only after the request arrives, so the reply and close cannot race the held write.
<-requestRead
respPayload := []byte(`{"type":"response.done","response":{"id":"resp-1","status":"completed","output":[]}}`)
_ = conn.WriteMessage(websocket.TextMessage, respPayload)
}))
Expand Down Expand Up @@ -2310,12 +2311,11 @@ func TestXAIWebsockets_KeepalivePingDuringUpload_Sessionless(t *testing.T) {
return nil
})

// The reader answers control frames until the client request arrives.
requestRead := make(chan struct{})
go func() {
for {
if _, _, errReadLoop := conn.ReadMessage(); errReadLoop != nil {
return
}
}
_, _, _ = conn.ReadMessage()
close(requestRead)
}()

select {
Expand All @@ -2338,6 +2338,8 @@ func TestXAIWebsockets_KeepalivePingDuringUpload_Sessionless(t *testing.T) {
return
}

// Reply only after the request arrives, so the reply and close cannot race the held write.
<-requestRead
respPayload := []byte(`{"type":"response.done","response":{"id":"resp-1","status":"completed","output":[]}}`)
_ = conn.WriteMessage(websocket.TextMessage, respPayload)
}))
Expand Down
Loading