From 564823e1a45f88a4c18f703a90465d77cbf7ceb4 Mon Sep 17 00:00:00 2001
From: Fedor Korotkov <fkorotkov@openai.com>
Date: Thu, 16 Jul 2026 17:10:15 -0400
Subject: [PATCH] Honor VM wait time in WebSocket requests

---
 .github/workflows/ci.yml |  2 +-
 pkg/client/client.go     | 31 +++++++++++++++++++++++++++++--
 pkg/client/vms_test.go   | 25 +++++++++++++++++++------
 3 files changed, 49 insertions(+), 9 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 933afe7..48f2410 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -60,7 +60,7 @@ jobs:
       - name: Pre-pull default Vetu image
         run: vetu pull ghcr.io/cirruslabs/ubuntu-amd64:24.04
       - name: Run tests
-        run: go test -v -count=1 ./...
+        run: go test -timeout=20m -v -count=1 ./...
 
   test-macos:
     name: Test (macOS)
diff --git a/pkg/client/client.go b/pkg/client/client.go
index 20420a7..842c937 100644
--- a/pkg/client/client.go
+++ b/pkg/client/client.go
@@ -13,6 +13,7 @@ import (
 	"net"
 	"net/http"
 	"net/url"
+	"strconv"
 	"time"
 
 	"github.com/cirruslabs/orchard/internal/config"
@@ -25,6 +26,8 @@ import (
 	"google.golang.org/grpc/metadata"
 )
 
+const defaultHTTPTimeout = 30 * time.Second
+
 type APIError struct {
 	StatusCode int
 }
@@ -107,7 +110,7 @@ func New(opts ...Option) (*Client, error) {
 		// the requests may hang indefinitely. See [1] for more details.
 		//
 		// [1]: https://github.com/cirruslabs/orchard/issues/152#issuecomment-1927091747
-		Timeout:   30 * time.Second,
+		Timeout:   defaultHTTPTimeout,
 		Transport: transport,
 	}
 
@@ -315,7 +318,7 @@ func (client *Client) wsRequestRaw(
 	endpointURL.RawQuery = values.Encode()
 
 	dialOptions := &websocket.DialOptions{
-		HTTPClient: client.httpClient,
+		HTTPClient: client.httpClientForWebSocket(params),
 		HTTPHeader: make(http.Header),
 	}
 
@@ -337,6 +340,30 @@ func (client *Client) wsRequestRaw(
 	return conn, nil
 }
 
+func (client *Client) httpClientForWebSocket(params map[string]string) *http.Client {
+	waitRaw, ok := params["wait"]
+	if !ok {
+		return client.httpClient
+	}
+
+	waitSeconds, err := strconv.ParseUint(waitRaw, 10, 16)
+	if err != nil {
+		return client.httpClient
+	}
+
+	waitTimeout := time.Duration(waitSeconds)*time.Second + defaultHTTPTimeout
+	if client.httpClient.Timeout == 0 || client.httpClient.Timeout >= waitTimeout {
+		return client.httpClient
+	}
+
+	return &http.Client{
+		CheckRedirect: client.httpClient.CheckRedirect,
+		Jar:           client.httpClient.Jar,
+		Timeout:       waitTimeout,
+		Transport:     client.httpClient.Transport,
+	}
+}
+
 func (client *Client) formatPath(path string) *url.URL {
 	endpointURL := &url.URL{
 		Scheme: client.baseURL.Scheme,
diff --git a/pkg/client/vms_test.go b/pkg/client/vms_test.go
index 710c30b..8fe2e79 100644
--- a/pkg/client/vms_test.go
+++ b/pkg/client/vms_test.go
@@ -4,21 +4,34 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"testing"
+	"time"
 
 	"github.com/coder/websocket"
 	"github.com/stretchr/testify/require"
 )
 
+func TestHTTPClientForWebSocketHonorsWait(t *testing.T) {
+	devClient, err := New(WithAddress("http://localhost"))
+	require.NoError(t, err)
+
+	httpClient := devClient.httpClientForWebSocket(map[string]string{"wait": "120"})
+
+	require.Equal(t, 150*time.Second, httpClient.Timeout)
+	require.Same(t, devClient.httpClient.Transport, httpClient.Transport)
+}
+
 func TestExecSessionBuildsReconnectableQuery(t *testing.T) {
 	var query map[string][]string
 
-	server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
-		query = request.URL.Query()
+	server := httptest.NewServer(
+		http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
+			query = request.URL.Query()
 
-		conn, err := websocket.Accept(writer, request, nil)
-		require.NoError(t, err)
-		defer conn.CloseNow()
-	}))
+			conn, err := websocket.Accept(writer, request, nil)
+			require.NoError(t, err)
+			defer conn.CloseNow()
+		}),
+	)
 	defer server.Close()
 
 	devClient, err := New(WithAddress(server.URL))