From 70ff64476ce61409612590feb110146386b20328 Mon Sep 17 00:00:00 2001 From: Farhan Nawaz Date: Thu, 10 Sep 2026 17:13:48 +0530 Subject: [PATCH 1/8] fix: use websocket with invoke command --- cli/azd/extensions/azure.ai.rle/go.mod | 1 + cli/azd/extensions/azure.ai.rle/go.sum | 2 + .../azure.ai.rle/internal/cmd/invoke.go | 59 ++- .../azure.ai.rle/internal/cmd/invoke_test.go | 66 ++- .../azure.ai.rle/internal/project/runtime.go | 16 +- .../internal/project/websocket_runtime.go | 411 ++++++++++++++++++ .../project/websocket_runtime_test.go | 333 ++++++++++++++ 7 files changed, 870 insertions(+), 18 deletions(-) create mode 100644 cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime.go create mode 100644 cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime_test.go diff --git a/cli/azd/extensions/azure.ai.rle/go.mod b/cli/azd/extensions/azure.ai.rle/go.mod index efe393449f5..31998eee0a7 100644 --- a/cli/azd/extensions/azure.ai.rle/go.mod +++ b/cli/azd/extensions/azure.ai.rle/go.mod @@ -7,6 +7,7 @@ require ( github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.13.1 github.com/azure/azure-dev/cli/azd v1.25.0 github.com/fatih/color v1.18.0 + github.com/gorilla/websocket v1.5.3 github.com/spf13/cobra v1.10.1 ) diff --git a/cli/azd/extensions/azure.ai.rle/go.sum b/cli/azd/extensions/azure.ai.rle/go.sum index e2fb881cb18..9dcf6549391 100644 --- a/cli/azd/extensions/azure.ai.rle/go.sum +++ b/cli/azd/extensions/azure.ai.rle/go.sum @@ -121,6 +121,8 @@ github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/gorilla/css v1.0.1 h1:ntNaBIghp6JmvWnxbZKANoLyuXTPZ4cAMlo6RyhlbO8= github.com/gorilla/css v1.0.1/go.mod h1:BvnYkspnSzMmwRK+b8/xgNPLiIuNZr6vbZBTPQ2A3b0= +github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg= +github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= github.com/hexops/gotextdiff v1.0.3 h1:gitA9+qJrrTCsiCl7+kh75nPqQt1cx4ZkudSTLoUqJM= github.com/hexops/gotextdiff v1.0.3/go.mod h1:pSWU5MAI3yDq+fZBTazCSJysOMbxWL1BSow5/V2vxeg= github.com/hinshun/vt10x v0.0.0-20220119200601-820417d04eec h1:qv2VnGeEQHchGaZ/u7lxST/RaJw+cv273q79D81Xbog= diff --git a/cli/azd/extensions/azure.ai.rle/internal/cmd/invoke.go b/cli/azd/extensions/azure.ai.rle/internal/cmd/invoke.go index 989e21c9f38..1905d381653 100644 --- a/cli/azd/extensions/azure.ai.rle/internal/cmd/invoke.go +++ b/cli/azd/extensions/azure.ai.rle/internal/cmd/invoke.go @@ -7,6 +7,7 @@ import ( "context" "crypto/rand" "encoding/hex" + "encoding/json" "errors" "fmt" "io" @@ -147,10 +148,17 @@ func (a *remoteInvokeAction) Run() error { ); err != nil { return err } + runtimeSession := project.NewWebSocketRuntimeSession( + instanceUrl, + a.flags.timeout, + client.authorizationHeader, + ) + defer runtimeSession.Close() playgroundUrl, stopPlayground, err := remotePlaygroundUrlWithAuthorizationProvider( ctx, instanceUrl, client.authorizationHeader, + runtimeSession, ) if err != nil { return err @@ -159,13 +167,11 @@ func (a *remoteInvokeAction) Run() error { if err := ui.OpenBrowser(playgroundUrl); err != nil { _, _ = fmt.Fprintf(a.cmd.ErrOrStderr(), "Warning: failed to open playground UI: %v\n", err) } - return project.RunShellWithContextAndAuthorizationProvider( + return project.RunWebSocketShellWithSession( ctx, a.cmd.InOrStdin(), a.cmd.OutOrStdout(), - instanceUrl, - a.flags.timeout, - client.authorizationHeader, + runtimeSession, ) } @@ -500,6 +506,7 @@ func remotePlaygroundUrlWithAuthorizationProvider( ctx context.Context, sandboxUrl string, authorizationProvider project.AuthorizationProvider, + runtimeSessions ...*project.WebSocketRuntimeSession, ) (string, func(), error) { listener, err := net.Listen("tcp", "127.0.0.1:0") if err != nil { @@ -517,6 +524,7 @@ func remotePlaygroundUrlWithAuthorizationProvider( authorizationProvider, listener.Addr().String(), sessionToken, + runtimeSessions..., ), ReadHeaderTimeout: 5 * time.Second, } @@ -544,6 +552,7 @@ func remotePlaygroundHandler( authorizationProvider project.AuthorizationProvider, expectedHost string, sessionToken string, + runtimeSessions ...*project.WebSocketRuntimeSession, ) http.Handler { mux := http.NewServeMux() mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { @@ -558,7 +567,7 @@ func remotePlaygroundHandler( _, _ = io.WriteString(w, ui.RemotePlaygroundHTML) return } - proxyOpenEnvToSandbox(w, r, sandboxUrl, authorizationProvider) + proxyOpenEnvToSandbox(w, r, sandboxUrl, authorizationProvider, runtimeSessions...) }) return mux } @@ -631,6 +640,7 @@ func proxyOpenEnvToSandbox( r *http.Request, sandboxUrl string, authorizationProvider project.AuthorizationProvider, + runtimeSessions ...*project.WebSocketRuntimeSession, ) { operation := strings.Trim(r.URL.Path, "/") switch operation { @@ -648,6 +658,11 @@ func proxyOpenEnvToSandbox( http.NotFound(w, r) return } + if len(runtimeSessions) > 0 && runtimeSessions[0] != nil && + (operation == "reset" || operation == "step" || operation == "state") { + proxyStatefulOpenEnvOperation(w, r, operation, runtimeSessions[0]) + return + } targetUrl, err := project.RuntimeOperationURL(sandboxUrl, operation) if err != nil { @@ -687,6 +702,40 @@ func proxyOpenEnvToSandbox( _, _ = io.Copy(w, resp.Body) } +func proxyStatefulOpenEnvOperation( + w http.ResponseWriter, + r *http.Request, + operation string, + runtimeSession *project.WebSocketRuntimeSession, +) { + payload := "" + if operation == "reset" || operation == "step" { + body, err := io.ReadAll(http.MaxBytesReader(w, r.Body, 100*1024*1024)) + if err != nil { + http.Error(w, "invalid request body", http.StatusBadRequest) + return + } + payload = string(body) + } + if operation == "step" { + var request struct { + Action json.RawMessage `json:"action"` + } + if err := json.Unmarshal([]byte(payload), &request); err != nil || len(request.Action) == 0 { + http.Error(w, "step requires an action", http.StatusBadRequest) + return + } + payload = string(request.Action) + } + response, err := runtimeSession.Call(r.Context(), operation, payload) + if err != nil { + http.Error(w, err.Error(), http.StatusBadGateway) + return + } + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, response) +} + func withFoundryAPIVersion(runtimeUrl string) (string, error) { parsedUrl, err := url.Parse(runtimeUrl) if err != nil { diff --git a/cli/azd/extensions/azure.ai.rle/internal/cmd/invoke_test.go b/cli/azd/extensions/azure.ai.rle/internal/cmd/invoke_test.go index d6dd9748ba5..d83bcd00650 100644 --- a/cli/azd/extensions/azure.ai.rle/internal/cmd/invoke_test.go +++ b/cli/azd/extensions/azure.ai.rle/internal/cmd/invoke_test.go @@ -23,9 +23,11 @@ import ( "testing" "time" + "azure.ai.rle/internal/project" "azure.ai.rle/internal/ui" "github.com/azure/azure-dev/cli/azd/pkg/azdext" + "github.com/gorilla/websocket" ) const testFoundryProjectPath = "/api/projects/project-1" @@ -51,6 +53,27 @@ func TestInvokeRemoteCreatesInstanceAndRunsShell(t *testing.T) { switch r.URL.Path { case "/health": _, _ = w.Write([]byte(`{"status":"healthy"}`)) + case "/ws": + connection, err := (&websocket.Upgrader{}).Upgrade(w, r, nil) + if err != nil { + t.Errorf("upgrade WebSocket: %v", err) + return + } + defer connection.Close() + var request map[string]any + if err := connection.ReadJSON(&request); err != nil { + t.Errorf("read WebSocket request: %v", err) + return + } + if request["type"] != "state" { + t.Errorf("expected state request, got %#v", request) + } + if err := connection.WriteJSON(map[string]any{ + "type": "state", + "data": map[string]any{"state": "ready"}, + }); err != nil { + t.Errorf("write WebSocket response: %v", err) + } default: http.NotFound(w, r) } @@ -91,7 +114,7 @@ func TestInvokeRemoteCreatesInstanceAndRunsShell(t *testing.T) { useTestProjectEndpoint(t, controlPlane.URL) command := newInvokeCommand() - command.SetIn(strings.NewReader("health\nexit\n")) + command.SetIn(strings.NewReader("state\nexit\n")) var output bytes.Buffer command.SetOut(&output) command.SetErr(&output) @@ -104,8 +127,8 @@ func TestInvokeRemoteCreatesInstanceAndRunsShell(t *testing.T) { if strings.Contains(output.String(), envServer.URL) { t.Fatalf("expected instance data-plane URL to remain hidden, got %s", output.String()) } - if !strings.Contains(output.String(), `"status": "healthy"`) { - t.Fatalf("expected remote shell health output, got %s", output.String()) + if !strings.Contains(output.String(), `"state": "ready"`) { + t.Fatalf("expected remote shell state output, got %s", output.String()) } if !instanceDeleted || !groupDeleted { t.Fatal("expected remote invoke to delete the instance and group") @@ -925,22 +948,40 @@ func TestRemotePlaygroundProxyForwardsToSandbox(t *testing.T) { requestCount := 0 envServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { requestCount++ - w.Header().Set("Content-Type", "application/json") - switch r.URL.Path { - case "/web": - http.NotFound(w, r) - case "/state": - _, _ = w.Write([]byte(`{"step_count":3}`)) - default: + if r.URL.Path != "/ws" { http.NotFound(w, r) + return + } + connection, err := (&websocket.Upgrader{}).Upgrade(w, r, nil) + if err != nil { + t.Errorf("upgrade WebSocket: %v", err) + return + } + defer connection.Close() + var request map[string]any + if err := connection.ReadJSON(&request); err != nil { + t.Errorf("read WebSocket request: %v", err) + return + } + if request["type"] != "state" { + t.Errorf("expected state request, got %#v", request) + } + if err := connection.WriteJSON(map[string]any{ + "type": "state", + "data": map[string]any{"step_count": 3}, + }); err != nil { + t.Errorf("write WebSocket response: %v", err) } })) defer envServer.Close() + runtimeSession := project.NewWebSocketRuntimeSession(envServer.URL, 30, nil) + defer runtimeSession.Close() playgroundUrl, stop, err := remotePlaygroundUrlWithAuthorizationProvider( t.Context(), envServer.URL, nil, + runtimeSession, ) if err != nil { t.Fatal(err) @@ -986,8 +1027,9 @@ func TestRemotePlaygroundProxyForwardsToSandbox(t *testing.T) { if err != nil { t.Fatal(err) } - if string(body) != `{"step_count":3}` { - t.Fatalf("expected proxied state body, got %s", body) + var state map[string]any + if err := json.Unmarshal(body, &state); err != nil || state["step_count"] != float64(3) { + t.Fatalf("expected proxied state body, got %s (err: %v)", body, err) } if requestCount != 1 { t.Fatalf("expected one authorized backend request, got %d", requestCount) diff --git a/cli/azd/extensions/azure.ai.rle/internal/project/runtime.go b/cli/azd/extensions/azure.ai.rle/internal/project/runtime.go index d3017495189..75bf1cc72c3 100644 --- a/cli/azd/extensions/azure.ai.rle/internal/project/runtime.go +++ b/cli/azd/extensions/azure.ai.rle/internal/project/runtime.go @@ -26,6 +26,8 @@ type callOptions struct { type AuthorizationProvider func(context.Context) (string, error) +type runtimeCaller func(context.Context, string, string, *callOptions, AuthorizationProvider) (string, error) + func RunShellWithContext( ctx context.Context, input io.Reader, @@ -80,6 +82,18 @@ func runShell( baseUrl string, timeout int, authorizationProvider AuthorizationProvider, +) error { + return runShellWithCaller(ctx, input, output, baseUrl, timeout, authorizationProvider, call) +} + +func runShellWithCaller( + ctx context.Context, + input io.Reader, + output io.Writer, + baseUrl string, + timeout int, + authorizationProvider AuthorizationProvider, + caller runtimeCaller, ) error { fmt.Fprintln(output, "Environment runtime shell. Type help for commands, exit to quit.") scanner := bufio.NewScanner(input) @@ -125,7 +139,7 @@ func runShell( } } - response, err := call(ctx, baseUrl, operation, flags, authorizationProvider) + response, err := caller(ctx, baseUrl, operation, flags, authorizationProvider) if err != nil { fmt.Fprintf(output, "error: %v\n", err) continue diff --git a/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime.go b/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime.go new file mode 100644 index 00000000000..bb1f70292b7 --- /dev/null +++ b/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime.go @@ -0,0 +1,411 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package project + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/url" + "strings" + "sync" + "time" + + "github.com/azure/azure-dev/cli/azd/pkg/azdext" + "github.com/gorilla/websocket" +) + +const maxWebSocketMessageBytes = 100 * 1024 * 1024 + +type WebSocketRuntimeSession struct { + baseURL string + timeout int + authorizationProvider AuthorizationProvider + mu sync.Mutex + exchangeMu sync.Mutex + connection *websocket.Conn + terminalError error + closed bool +} + +func NewWebSocketRuntimeSession( + baseURL string, + timeout int, + authorizationProvider AuthorizationProvider, +) *WebSocketRuntimeSession { + return &WebSocketRuntimeSession{ + baseURL: baseURL, + timeout: timeout, + authorizationProvider: authorizationProvider, + } +} + +// RunWebSocketShellWithContextAndAuthorizationProvider uses one persistent WebSocket +// for stateful OpenEnv operations while retaining the safe HTTP endpoints. +func RunWebSocketShellWithContextAndAuthorizationProvider( + ctx context.Context, + input io.Reader, + output io.Writer, + baseURL string, + timeout int, + authorizationProvider AuthorizationProvider, +) error { + session := NewWebSocketRuntimeSession(baseURL, timeout, authorizationProvider) + defer session.Close() + return RunWebSocketShellWithSession(ctx, input, output, session) +} + +func RunWebSocketShellWithSession( + ctx context.Context, + input io.Reader, + output io.Writer, + session *WebSocketRuntimeSession, +) error { + done := make(chan error, 1) + go func() { + done <- runShellWithCaller( + ctx, + input, + output, + session.baseURL, + session.timeout, + session.authorizationProvider, + session.call, + ) + }() + select { + case err := <-done: + return err + case <-ctx.Done(): + session.Close() + fmt.Fprintln(output) + return nil + } +} + +func (c *WebSocketRuntimeSession) call( + ctx context.Context, + baseURL string, + operation string, + flags *callOptions, + authorizationProvider AuthorizationProvider, +) (string, error) { + switch operation { + case "reset", "step", "state": + return c.exchange(ctx, operation, flags) + default: + return call(ctx, baseURL, operation, flags, authorizationProvider) + } +} + +func (c *WebSocketRuntimeSession) Call( + ctx context.Context, + operation string, + payload string, +) (string, error) { + flags := &callOptions{timeout: c.timeout} + switch operation { + case "reset": + flags.body = payload + case "step": + flags.action = payload + } + return c.call(ctx, c.baseURL, operation, flags, c.authorizationProvider) +} + +func (c *WebSocketRuntimeSession) exchange( + ctx context.Context, + operation string, + flags *callOptions, +) (string, error) { + c.exchangeMu.Lock() + defer c.exchangeMu.Unlock() + if err := ctx.Err(); err != nil { + return "", err + } + if err := c.connect(ctx); err != nil { + return "", err + } + c.mu.Lock() + connection := c.connection + terminalError := c.terminalError + c.mu.Unlock() + if connection == nil { + if terminalError != nil { + return "", terminalError + } + return "", fmt.Errorf("OpenEnv WebSocket session closed") + } + + request, err := webSocketRequest(operation, flags) + if err != nil { + return "", err + } + + deadline, hasDeadline := operationDeadline(ctx, c.timeout) + if hasDeadline { + if err := connection.SetWriteDeadline(deadline); err != nil { + return "", c.failConnection(connection, fmt.Errorf("set OpenEnv WebSocket write deadline: %w", err)) + } + if err := connection.SetReadDeadline(deadline); err != nil { + return "", c.failConnection(connection, fmt.Errorf("set OpenEnv WebSocket read deadline: %w", err)) + } + } else { + _ = connection.SetWriteDeadline(time.Time{}) + _ = connection.SetReadDeadline(time.Time{}) + } + + if err := connection.WriteMessage(websocket.TextMessage, request); err != nil { + return "", c.failConnection( + connection, + fmt.Errorf("send OpenEnv WebSocket %s request: %w", operation, err), + ) + } + exchangeDone := make(chan struct{}) + cancellationHandled := make(chan struct{}) + go func() { + select { + case <-exchangeDone: + case <-ctx.Done(): + c.failConnection( + connection, + fmt.Errorf("OpenEnv WebSocket %s request canceled: %w", operation, ctx.Err()), + ) + } + close(cancellationHandled) + }() + messageType, response, err := connection.ReadMessage() + close(exchangeDone) + <-cancellationHandled + if err != nil { + return "", c.failConnection( + connection, + fmt.Errorf("receive OpenEnv WebSocket %s response: %w", operation, err), + ) + } + if messageType != websocket.TextMessage { + err := &azdext.LocalError{ + Message: "Environment runtime returned a non-text WebSocket response.", + Code: "rle_open_env_websocket_protocol_error", + Category: azdext.LocalErrorCategoryInternal, + } + return "", c.failConnection(connection, err) + } + result, terminal, err := parseWebSocketResponse(operation, response) + if terminal { + return "", c.failConnection(connection, err) + } + return result, err +} + +func (c *WebSocketRuntimeSession) connect(ctx context.Context) error { + c.mu.Lock() + defer c.mu.Unlock() + if c.closed { + return c.terminalError + } + if c.terminalError != nil { + return c.terminalError + } + if c.connection != nil { + return nil + } + + endpoint, err := RuntimeWebSocketURL(c.baseURL) + if err != nil { + return err + } + headers := http.Header{} + if c.authorizationProvider != nil { + authorization, err := c.authorizationProvider(ctx) + if err != nil { + return fmt.Errorf("authenticate to environment runtime: %w", err) + } + if authorization != "" { + headers.Set("Authorization", authorization) + } + } + + dialer := *websocket.DefaultDialer + dialer.HandshakeTimeout = 0 + if c.timeout > 0 { + dialer.HandshakeTimeout = time.Duration(c.timeout) * time.Second + } + connection, response, err := dialer.DialContext(ctx, endpoint, headers) + if err != nil { + detail := "" + if response != nil { + detail = readHealthErrorDetail(response.Body) + _ = response.Body.Close() + } + return &azdext.LocalError{ + Message: fmt.Sprintf( + "Environment runtime WebSocket connection failed%s: %v", + detail, + err, + ), + Code: "rle_open_env_websocket_connection_failed", + Category: azdext.LocalErrorCategoryUser, + Suggestion: "Check the remote RLE instance status and retry invoke.", + } + } + connection.SetReadLimit(maxWebSocketMessageBytes) + c.connection = connection + return nil +} + +func (c *WebSocketRuntimeSession) failConnection(connection *websocket.Conn, err error) error { + c.mu.Lock() + if c.connection == connection { + c.connection = nil + } + if c.terminalError == nil { + c.terminalError = &azdext.LocalError{ + Message: fmt.Sprintf("The OpenEnv WebSocket session is no longer usable: %v", err), + Code: "rle_open_env_websocket_session_failed", + Category: azdext.LocalErrorCategoryUser, + Suggestion: "Exit and run invoke again to start a new environment session.", + } + } + terminalError := c.terminalError + c.mu.Unlock() + _ = connection.Close() + return terminalError +} + +func (c *WebSocketRuntimeSession) Close() { + c.mu.Lock() + connection := c.connection + c.connection = nil + c.closed = true + if c.terminalError == nil { + c.terminalError = &azdext.LocalError{ + Message: "The OpenEnv WebSocket session is closed.", + Code: "rle_open_env_websocket_session_closed", + Category: azdext.LocalErrorCategoryUser, + Suggestion: "Run invoke again to start a new environment session.", + } + } + c.mu.Unlock() + if connection == nil { + return + } + deadline := time.Now().Add(2 * time.Second) + _ = connection.WriteControl( + websocket.CloseMessage, + websocket.FormatCloseMessage(websocket.CloseNormalClosure, "RLE invoke complete."), + deadline, + ) + _ = connection.Close() +} + +func parseWebSocketResponse(operation string, response []byte) (string, bool, error) { + var envelope struct { + Type string `json:"type"` + Data json.RawMessage `json:"data"` + } + if err := json.Unmarshal(response, &envelope); err != nil { + return "", true, &azdext.LocalError{ + Message: fmt.Sprintf("Environment runtime returned invalid WebSocket JSON: %v", err), + Code: "rle_open_env_websocket_protocol_error", + Category: azdext.LocalErrorCategoryInternal, + } + } + if envelope.Type == "error" { + return "", false, &azdext.LocalError{ + Message: fmt.Sprintf("Environment runtime rejected the %s request: %s", operation, prettyJson(envelope.Data)), + Code: "rle_open_env_websocket_request_failed", + Category: azdext.LocalErrorCategoryUser, + Suggestion: "Check the request payload and retry.", + } + } + expectedType := "observation" + if operation == "state" { + expectedType = "state" + } + if envelope.Type != expectedType || len(envelope.Data) == 0 { + return "", true, &azdext.LocalError{ + Message: fmt.Sprintf( + "Environment runtime returned WebSocket response type %q for %s; expected %q.", + envelope.Type, + operation, + expectedType, + ), + Code: "rle_open_env_websocket_protocol_error", + Category: azdext.LocalErrorCategoryInternal, + } + } + return prettyJson(envelope.Data), false, nil +} + +func RuntimeWebSocketURL(baseURL string) (string, error) { + endpoint, err := url.Parse(baseURL) + if err != nil { + return "", fmt.Errorf("parse environment runtime URL: %w", err) + } + switch strings.ToLower(endpoint.Scheme) { + case "http": + endpoint.Scheme = "ws" + case "https": + endpoint.Scheme = "wss" + default: + return "", fmt.Errorf("environment runtime URL must use HTTP or HTTPS") + } + endpoint.Path = strings.TrimRight(endpoint.Path, "/") + "/ws" + endpoint.RawPath = "" + return endpoint.String(), nil +} + +func webSocketRequest(operation string, flags *callOptions) ([]byte, error) { + request := map[string]any{"type": operation} + switch operation { + case "reset": + data, err := webSocketData(flags.body, "body", true) + if err != nil { + return nil, err + } + request["data"] = data + case "step": + data, err := webSocketData(flags.action, "action", false) + if err != nil { + return nil, err + } + request["data"] = data + case "state": + default: + return nil, fmt.Errorf("operation %q is not supported over OpenEnv WebSocket", operation) + } + return json.Marshal(request) +} + +func webSocketData(value string, flagName string, allowEmpty bool) (map[string]any, error) { + if strings.TrimSpace(value) == "" && allowEmpty { + return map[string]any{}, nil + } + data, err := validateJsonObject(value, flagName) + if err != nil { + return nil, err + } + var decoded map[string]any + if err := json.Unmarshal(data, &decoded); err != nil { + return nil, err + } + return decoded, nil +} + +func operationDeadline(ctx context.Context, timeoutSeconds int) (time.Time, bool) { + var deadline time.Time + hasDeadline := false + if timeoutSeconds > 0 { + deadline = time.Now().Add(time.Duration(timeoutSeconds) * time.Second) + hasDeadline = true + } + if contextDeadline, ok := ctx.Deadline(); ok && (!hasDeadline || contextDeadline.Before(deadline)) { + deadline = contextDeadline + hasDeadline = true + } + return deadline, hasDeadline +} diff --git a/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime_test.go b/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime_test.go new file mode 100644 index 00000000000..fb8e368a3c5 --- /dev/null +++ b/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime_test.go @@ -0,0 +1,333 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package project + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "slices" + "strings" + "testing" + + "github.com/gorilla/websocket" +) + +func TestRunWebSocketShellUsesPersistentSocketForStatefulOperations(t *testing.T) { + var requests []map[string]any + var safePaths []string + upgrades := 0 + upgrader := websocket.Upgrader{} + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Query().Get("api-version") != "test-version" { + t.Errorf("expected preserved API version, got %q", r.URL.Query().Get("api-version")) + } + if r.Header.Get("Authorization") != "******" { + t.Errorf("unexpected authorization header %q", r.Header.Get("Authorization")) + } + if r.URL.Path != "/ws" { + safePaths = append(safePaths, r.URL.Path) + _, _ = fmt.Fprintf(w, `{"path":%q}`, r.URL.Path) + return + } + + upgrades++ + connection, err := upgrader.Upgrade(w, r, nil) + if err != nil { + t.Errorf("upgrade WebSocket: %v", err) + return + } + defer connection.Close() + for { + var request map[string]any + if err := connection.ReadJSON(&request); err != nil { + if websocket.IsCloseError(err, websocket.CloseNormalClosure) { + return + } + t.Errorf("read WebSocket request: %v", err) + return + } + requests = append(requests, request) + responseType := "observation" + if request["type"] == "state" { + responseType = "state" + } + if err := connection.WriteJSON(map[string]any{ + "type": responseType, + "data": map[string]any{"requestType": request["type"]}, + }); err != nil { + t.Errorf("write WebSocket response: %v", err) + return + } + } + })) + defer server.Close() + + authorizationCalls := 0 + authorizationProvider := func(context.Context) (string, error) { + authorizationCalls++ + return "******", nil + } + input := strings.NewReader( + "reset {\"seed\":42}\nstep {\"message\":\"hello\"}\nstate\nhealth\nmetadata\nschema\nexit\n", + ) + var output bytes.Buffer + err := RunWebSocketShellWithContextAndAuthorizationProvider( + t.Context(), + input, + &output, + server.URL+"?api-version=test-version", + 30, + authorizationProvider, + ) + if err != nil { + t.Fatal(err) + } + + if upgrades != 1 { + t.Fatalf("expected one persistent WebSocket, got %d", upgrades) + } + if len(requests) != 3 { + t.Fatalf("expected three WebSocket requests, got %#v", requests) + } + assertWebSocketRequest(t, requests[0], "reset", map[string]any{"seed": float64(42)}) + assertWebSocketRequest(t, requests[1], "step", map[string]any{"message": "hello"}) + assertWebSocketRequest(t, requests[2], "state", nil) + if !slices.Equal(safePaths, []string{"/health", "/metadata", "/schema"}) { + t.Fatalf("unexpected safe HTTP operations: %v", safePaths) + } + if authorizationCalls != 4 { + t.Fatalf("expected one WebSocket and three HTTP authorization calls, got %d", authorizationCalls) + } + if !strings.Contains(output.String(), `"requestType": "step"`) { + t.Fatalf("expected formatted WebSocket response, got %s", output.String()) + } +} + +func TestRuntimeWebSocketURL(t *testing.T) { + tests := []struct { + baseURL string + want string + }{ + { + baseURL: "https://example.test/openenv?api-version=1", + want: "wss://example.test/openenv/ws?api-version=1", + }, + { + baseURL: "http://127.0.0.1:8080/openenv/", + want: "ws://127.0.0.1:8080/openenv/ws", + }, + } + for _, test := range tests { + got, err := RuntimeWebSocketURL(test.baseURL) + if err != nil { + t.Fatal(err) + } + if got != test.want { + t.Fatalf("RuntimeWebSocketURL(%q) = %q, want %q", test.baseURL, got, test.want) + } + } + if _, err := RuntimeWebSocketURL("ftp://example.test/openenv"); err == nil { + t.Fatal("expected unsupported scheme to fail") + } +} + +func TestWebSocketRequestUsesOpenEnvProtocol(t *testing.T) { + reset, err := webSocketRequest("reset", &callOptions{}) + if err != nil { + t.Fatal(err) + } + + step, err := webSocketRequest("step", &callOptions{action: `{"message":"hello"}`}) + if err != nil { + t.Fatal(err) + } + state, err := webSocketRequest("state", &callOptions{}) + if err != nil { + t.Fatal(err) + } + + assertJSONEqual(t, reset, `{"type":"reset","data":{}}`) + assertJSONEqual(t, step, `{"type":"step","data":{"message":"hello"}}`) + assertJSONEqual(t, state, `{"type":"state"}`) +} + +func TestParseWebSocketResponse(t *testing.T) { + result, terminal, err := parseWebSocketResponse( + "step", + []byte(`{"type":"observation","data":{"reward":1}}`), + ) + if err != nil || terminal || !strings.Contains(result, `"reward": 1`) { + t.Fatalf("unexpected successful response: result=%q terminal=%t err=%v", result, terminal, err) + } + + _, terminal, err = parseWebSocketResponse( + "step", + []byte(`{"type":"error","data":{"detail":"invalid action"}}`), + ) + if err == nil || terminal || !strings.Contains(err.Error(), "invalid action") { + t.Fatalf("unexpected OpenEnv error response: terminal=%t err=%v", terminal, err) + } + + _, terminal, err = parseWebSocketResponse( + "state", + []byte(`{"type":"observation","data":{}}`), + ) + if err == nil || !terminal { + t.Fatalf("expected mismatched response type to be terminal, got terminal=%t err=%v", terminal, err) + } +} + +func TestWebSocketSessionCloseInterruptsCallAndPreventsReconnect(t *testing.T) { + requestReceived := make(chan struct{}) + upgrades := 0 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + upgrades++ + connection, err := (&websocket.Upgrader{}).Upgrade(w, r, nil) + if err != nil { + t.Errorf("upgrade WebSocket: %v", err) + return + } + + defer connection.Close() + if _, _, err := connection.ReadMessage(); err != nil { + t.Errorf("read WebSocket request: %v", err) + return + } + close(requestReceived) + _, _, _ = connection.ReadMessage() + })) + defer server.Close() + + session := NewWebSocketRuntimeSession(server.URL, 30, nil) + callDone := make(chan error, 1) + go func() { + _, err := session.Call(t.Context(), "state", "") + callDone <- err + }() + <-requestReceived + session.Close() + if err := <-callDone; err == nil { + t.Fatal("expected closing the session to fail the active call") + } + if _, err := session.Call(t.Context(), "state", ""); err == nil { + t.Fatal("expected a closed session to reject subsequent calls") + } + if upgrades != 1 { + t.Fatalf("expected no reconnection after close, got %d connections", upgrades) + } +} + +func TestCanceledQueuedCallDoesNotReachWebSocket(t *testing.T) { + firstRequestReceived := make(chan struct{}) + releaseFirstRequest := make(chan struct{}) + requests := 0 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + connection, err := (&websocket.Upgrader{}).Upgrade(w, r, nil) + if err != nil { + t.Errorf("upgrade WebSocket: %v", err) + return + } + defer connection.Close() + for { + var request map[string]any + if err := connection.ReadJSON(&request); err != nil { + if websocket.IsCloseError(err, websocket.CloseNormalClosure) { + return + } + t.Errorf("read WebSocket request: %v", err) + return + } + requests++ + if requests == 1 { + close(firstRequestReceived) + <-releaseFirstRequest + } + if err := connection.WriteJSON(map[string]any{ + "type": "state", + "data": map[string]any{"request": requests}, + }); err != nil { + t.Errorf("write WebSocket response: %v", err) + return + } + } + })) + defer server.Close() + + session := NewWebSocketRuntimeSession(server.URL, 30, nil) + defer session.Close() + firstCallDone := make(chan error, 1) + go func() { + _, err := session.Call(t.Context(), "state", "") + firstCallDone <- err + }() + <-firstRequestReceived + + queuedContext, cancelQueued := context.WithCancel(t.Context()) + queuedCallDone := make(chan error, 1) + go func() { + _, err := session.Call(queuedContext, "state", "") + queuedCallDone <- err + }() + cancelQueued() + close(releaseFirstRequest) + + if err := <-firstCallDone; err != nil { + t.Fatalf("first call failed: %v", err) + } + if err := <-queuedCallDone; err == nil { + t.Fatal("expected canceled queued call to fail") + } + if requests != 1 { + t.Fatalf("expected canceled queued call not to reach WebSocket, got %d requests", requests) + } +} + +func assertWebSocketRequest( + t *testing.T, + request map[string]any, + requestType string, + data map[string]any, +) { + t.Helper() + if request["type"] != requestType { + t.Fatalf("request type = %#v, want %q", request["type"], requestType) + } + if data == nil { + if _, ok := request["data"]; ok { + t.Fatalf("did not expect data in %#v", request) + } + return + } + actual, ok := request["data"].(map[string]any) + if !ok || !mapsEqual(actual, data) { + t.Fatalf("request data = %#v, want %#v", request["data"], data) + } +} + +func assertJSONEqual(t *testing.T, actual []byte, expected string) { + t.Helper() + var actualValue any + var expectedValue any + if err := json.Unmarshal(actual, &actualValue); err != nil { + t.Fatal(err) + } + if err := json.Unmarshal([]byte(expected), &expectedValue); err != nil { + t.Fatal(err) + } + actualJSON, _ := json.Marshal(actualValue) + expectedJSON, _ := json.Marshal(expectedValue) + if !bytes.Equal(actualJSON, expectedJSON) { + t.Fatalf("JSON = %s, want %s", actual, expected) + } +} + +func mapsEqual(left map[string]any, right map[string]any) bool { + leftJSON, _ := json.Marshal(left) + rightJSON, _ := json.Marshal(right) + return bytes.Equal(leftJSON, rightJSON) +} From 52a6ad73ac4f0990e17f6b467575aba03cc2f835 Mon Sep 17 00:00:00 2001 From: Farhan Nawaz Date: Thu, 10 Sep 2026 19:28:27 +0530 Subject: [PATCH 2/8] fix: allow init for all openenv envs --- cli/azd/extensions/azure.ai.rle/README.md | 7 +++- .../extensions/azure.ai.rle/extension.yaml | 4 +- .../azure.ai.rle/internal/cmd/init.go | 16 +++---- .../azure.ai.rle/internal/cmd/root_test.go | 27 ++++++------ .../azure.ai.rle/internal/project/scaffold.go | 42 +++++++++++++------ .../internal/project/scaffold_test.go | 31 +++++++++++++- 6 files changed, 89 insertions(+), 38 deletions(-) diff --git a/cli/azd/extensions/azure.ai.rle/README.md b/cli/azd/extensions/azure.ai.rle/README.md index c614097bd85..87c93470feb 100644 --- a/cli/azd/extensions/azure.ai.rle/README.md +++ b/cli/azd/extensions/azure.ai.rle/README.md @@ -92,12 +92,15 @@ The default echo session downloads the Hugging Face `OpenEnv` repo, copies `envs The copied session does not keep `.git` metadata from the upstream repository. -Name the copied echo session: +Copy another environment from the OpenEnv `envs` catalog: ```powershell -azd ai rle init code_rl +azd ai rle init chess_env +cd .\chess_env ``` +The positional name selects `envs/` from OpenEnv and is also used for the local session directory and RLE environment name. + For an existing source folder, skip `init` and run commands directly from that folder. ### 2. Run locally diff --git a/cli/azd/extensions/azure.ai.rle/extension.yaml b/cli/azd/extensions/azure.ai.rle/extension.yaml index 93c51754c91..eb99f6e9e20 100644 --- a/cli/azd/extensions/azure.ai.rle/extension.yaml +++ b/cli/azd/extensions/azure.ai.rle/extension.yaml @@ -14,8 +14,8 @@ usage: azd ai rle [options] version: 0.4.0-preview examples: - name: init - description: Copy the OpenEnv echo sample into a local RLE environment. - usage: azd ai rle init + description: Copy an environment from the OpenEnv catalog; defaults to echo_env. + usage: azd ai rle init [environment-name] - name: publish description: Build, push, and create or update the RLE environment. usage: azd ai rle publish --version-bump major diff --git a/cli/azd/extensions/azure.ai.rle/internal/cmd/init.go b/cli/azd/extensions/azure.ai.rle/internal/cmd/init.go index 901d92c4d1a..81c45931f2a 100644 --- a/cli/azd/extensions/azure.ai.rle/internal/cmd/init.go +++ b/cli/azd/extensions/azure.ai.rle/internal/cmd/init.go @@ -22,24 +22,24 @@ type rleInitFlags struct { type initAction struct { cmd *cobra.Command flags *rleInitFlags - envNameOverride string + environmentName string } -var checkoutOpenEnvEchoSampleFunc = project.CheckoutOpenEnvEchoSample +var checkoutOpenEnvEnvironmentFunc = project.CheckoutOpenEnvEnvironment func newInitCommand() *cobra.Command { flags := &rleInitFlags{} cmd := &cobra.Command{ - Use: "init", + Use: "init [environment-name]", Short: "Initialize a local RLE environment", Args: cobra.MaximumNArgs(1), RunE: func(cmd *cobra.Command, args []string) error { - envNameOverride := "" + environmentName := "" if len(args) == 1 { - envNameOverride = args[0] + environmentName = args[0] } - return (&initAction{cmd: cmd, flags: flags, envNameOverride: envNameOverride}).Run() + return (&initAction{cmd: cmd, flags: flags, environmentName: environmentName}).Run() }, } @@ -62,7 +62,7 @@ func newInitCommand() *cobra.Command { } func (a *initAction) Run() error { - envName := firstNonEmpty(a.envNameOverride, "echo_env") + envName := firstNonEmpty(a.environmentName, "echo_env") var err error envName, err = project.ValidateEnvironmentName(envName) if err != nil { @@ -74,7 +74,7 @@ func (a *initAction) Run() error { } } - sessionDir, err := checkoutOpenEnvEchoSampleFunc(envName, ".", a.flags.force) + sessionDir, err := checkoutOpenEnvEnvironmentFunc(envName, ".", a.flags.force) if err != nil { return err } diff --git a/cli/azd/extensions/azure.ai.rle/internal/cmd/root_test.go b/cli/azd/extensions/azure.ai.rle/internal/cmd/root_test.go index d0797f8d801..3cdf4ceab00 100644 --- a/cli/azd/extensions/azure.ai.rle/internal/cmd/root_test.go +++ b/cli/azd/extensions/azure.ai.rle/internal/cmd/root_test.go @@ -242,7 +242,7 @@ func TestLifecycleCommandsRejectPositionalArguments(t *testing.T) { func TestInitCopiesOpenEnvEchoSampleByDefault(t *testing.T) { tempDir := t.TempDir() t.Chdir(tempDir) - stubOpenEnvEchoCheckout(t) + stubOpenEnvCheckout(t, "echo_env") command := newInitCommand() var output bytes.Buffer @@ -280,13 +280,13 @@ func TestInitCopiesOpenEnvEchoSampleByDefault(t *testing.T) { } } -func TestInitUsesPositionalNameForDefaultSample(t *testing.T) { +func TestInitUsesPositionalNameToSelectOpenEnvEnvironment(t *testing.T) { tempDir := t.TempDir() t.Chdir(tempDir) - stubOpenEnvEchoCheckout(t) + stubOpenEnvCheckout(t, "chess_env") command := newInitCommand() - command.SetArgs([]string{"code_rl"}) + command.SetArgs([]string{"chess_env"}) var output bytes.Buffer command.SetOut(&output) command.SetErr(&output) @@ -294,7 +294,7 @@ func TestInitUsesPositionalNameForDefaultSample(t *testing.T) { t.Fatal(err) } - sessionDir := filepath.Join(tempDir, "code_rl") + sessionDir := filepath.Join(tempDir, "chess_env") // The test reads state from its own temporary session directory. stateBytes, err := os.ReadFile(filepath.Join(sessionDir, rleStateFile)) //nolint:gosec if err != nil { @@ -304,8 +304,8 @@ func TestInitUsesPositionalNameForDefaultSample(t *testing.T) { if err := json.Unmarshal(stateBytes, &state); err != nil { t.Fatal(err) } - if state.EnvironmentName != "code_rl" { - t.Fatalf("expected code_rl environment name, got %q", state.EnvironmentName) + if state.EnvironmentName != "chess_env" { + t.Fatalf("expected chess_env environment name, got %q", state.EnvironmentName) } } @@ -368,10 +368,13 @@ func TestInitNextStepsUseShellAppropriateSyntax(t *testing.T) { } } -func stubOpenEnvEchoCheckout(t *testing.T) { +func stubOpenEnvCheckout(t *testing.T, expectedName string) { t.Helper() - old := checkoutOpenEnvEchoSampleFunc - checkoutOpenEnvEchoSampleFunc = func(name string, dest string, force bool) (string, error) { + old := checkoutOpenEnvEnvironmentFunc + checkoutOpenEnvEnvironmentFunc = func(name string, dest string, force bool) (string, error) { + if name != expectedName { + t.Fatalf("expected OpenEnv environment %q, got %q", expectedName, name) + } sessionDir := filepath.Join(dest, name) if force { if err := os.RemoveAll(sessionDir); err != nil { @@ -388,12 +391,12 @@ func stubOpenEnvEchoCheckout(t *testing.T) { if err := os.WriteFile(filepath.Join(serverDir, "Dockerfile"), []byte("FROM scratch\n"), 0600); err != nil { return "", err } - if err := os.WriteFile(filepath.Join(sessionDir, "openenv.yaml"), []byte("name: echo_env\n"), 0600); err != nil { + if err := os.WriteFile(filepath.Join(sessionDir, "openenv.yaml"), []byte("name: "+name+"\n"), 0600); err != nil { return "", err } return sessionDir, nil } t.Cleanup(func() { - checkoutOpenEnvEchoSampleFunc = old + checkoutOpenEnvEnvironmentFunc = old }) } diff --git a/cli/azd/extensions/azure.ai.rle/internal/project/scaffold.go b/cli/azd/extensions/azure.ai.rle/internal/project/scaffold.go index c0a1a3b3a67..71d51959a52 100644 --- a/cli/azd/extensions/azure.ai.rle/internal/project/scaffold.go +++ b/cli/azd/extensions/azure.ai.rle/internal/project/scaffold.go @@ -13,9 +13,8 @@ import ( ) const ( - openEnvRepoUrl = "https://github.com/huggingface/OpenEnv.git" - openEnvRepoRef = "main" - openEnvEchoSamplePath = "envs/echo_env" + openEnvRepoUrl = "https://github.com/huggingface/OpenEnv.git" + openEnvRepoRef = "main" ) func createRleSessionDir(name string, dest string, force bool) (string, error) { @@ -41,15 +40,12 @@ func createRleSessionDir(name string, dest string, force bool) (string, error) { return sessionDir, nil } -func CheckoutOpenEnvEchoSample(name string, dest string, force bool) (string, error) { +func CheckoutOpenEnvEnvironment(name string, dest string, force bool) (string, error) { name, err := ValidateEnvironmentName(name) if err != nil { return "", err } - sessionDir, err := createRleSessionDir(name, dest, force) - if err != nil { - return "", err - } + sourcePath := openEnvEnvironmentPath(name) tempDir, err := os.MkdirTemp("", "azd-rle-open-env-*") if err != nil { return "", err @@ -69,17 +65,39 @@ func CheckoutOpenEnvEchoSample(name string, dest string, force bool) (string, er ); err != nil { return "", err } - if err := runGitCheckout("-C", tempDir, "sparse-checkout", "set", openEnvEchoSamplePath); err != nil { + if err := runGitCheckout("-C", tempDir, "sparse-checkout", "set", sourcePath); err != nil { return "", err } - sourceDir := filepath.Join(tempDir, filepath.FromSlash(openEnvEchoSamplePath)) + sourceDir := filepath.Join(tempDir, filepath.FromSlash(sourcePath)) + return copyOpenEnvEnvironment(sourceDir, name, dest, force) +} + +func copyOpenEnvEnvironment(sourceDir string, name string, dest string, force bool) (string, error) { + if _, err := os.Stat(sourceDir); os.IsNotExist(err) { + return "", &azdext.LocalError{ + Message: fmt.Sprintf("OpenEnv environment %q was not found.", name), + Code: "rle_open_env_environment_not_found", + Category: azdext.LocalErrorCategoryUser, + Suggestion: fmt.Sprintf("Choose an environment from %s/tree/%s/envs.", strings.TrimSuffix(openEnvRepoUrl, ".git"), openEnvRepoRef), + } + } else if err != nil { + return "", err + } + sessionDir, err := createRleSessionDir(name, dest, force) + if err != nil { + return "", err + } if err := copyDirectory(sourceDir, sessionDir); err != nil { return "", err } return sessionDir, nil } +func openEnvEnvironmentPath(name string) string { + return "envs/" + name +} + func runGitCheckout(args ...string) error { if _, err := exec.LookPath("git"); err != nil { return &azdext.LocalError{ @@ -89,12 +107,12 @@ func runGitCheckout(args ...string) error { Suggestion: "Install Git, then retry azd ai rle init.", } } - process := exec.Command("git", args...) //nolint:gosec // args are fixed by init's OpenEnv sample checkout flow. + process := exec.Command("git", args...) //nolint:gosec // Arguments are passed directly; the user value is validated as an environment name. process.Env = os.Environ() output, err := process.CombinedOutput() if err != nil { return &azdext.LocalError{ - Message: fmt.Sprintf("Failed to checkout OpenEnv echo sample: %v", err), + Message: fmt.Sprintf("Failed to checkout OpenEnv environment: %v", err), Code: "rle_open_env_checkout_failed", Category: azdext.LocalErrorCategoryUser, Suggestion: strings.TrimSpace(string(output)), diff --git a/cli/azd/extensions/azure.ai.rle/internal/project/scaffold_test.go b/cli/azd/extensions/azure.ai.rle/internal/project/scaffold_test.go index 23d701ee1ca..15eaee024bf 100644 --- a/cli/azd/extensions/azure.ai.rle/internal/project/scaffold_test.go +++ b/cli/azd/extensions/azure.ai.rle/internal/project/scaffold_test.go @@ -71,17 +71,44 @@ func TestCopyDirectoryRejectsFileSource(t *testing.T) { } } -func TestCheckoutOpenEnvEchoSampleRejectsInvalidNameBeforeChangingDestination(t *testing.T) { +func TestCheckoutOpenEnvEnvironmentRejectsInvalidNameBeforeChangingDestination(t *testing.T) { destDir := t.TempDir() sentinel := filepath.Join(destDir, "sentinel.txt") if err := os.WriteFile(sentinel, []byte("keep"), 0600); err != nil { t.Fatal(err) } - if _, err := CheckoutOpenEnvEchoSample("../bad", destDir, true); err == nil { + if _, err := CheckoutOpenEnvEnvironment("../bad", destDir, true); err == nil { t.Fatal("expected invalid environment name to be rejected") } + if _, err := os.Stat(sentinel); err != nil { t.Fatalf("expected destination to be unchanged: %v", err) } } + +func TestOpenEnvEnvironmentPath(t *testing.T) { + if got := openEnvEnvironmentPath("chess_env"); got != "envs/chess_env" { + t.Fatalf("expected selected OpenEnv environment path, got %q", got) + } +} + +func TestCopyOpenEnvEnvironmentValidatesSourceBeforeReplacingDestination(t *testing.T) { + destDir := t.TempDir() + sessionDir := filepath.Join(destDir, "missing_env") + if err := os.MkdirAll(sessionDir, 0750); err != nil { + t.Fatal(err) + } + sentinel := filepath.Join(sessionDir, "keep.txt") + if err := os.WriteFile(sentinel, []byte("keep"), 0600); err != nil { + t.Fatal(err) + } + + _, err := copyOpenEnvEnvironment(filepath.Join(t.TempDir(), "missing"), "missing_env", destDir, true) + if err == nil { + t.Fatal("expected missing OpenEnv environment to fail") + } + if _, statErr := os.Stat(sentinel); statErr != nil { + t.Fatalf("expected destination to remain unchanged after catalog lookup failure: %v", statErr) + } +} From 602f5e6a4cdf499a9a6e7dc3d84bdb189f7ff4fa Mon Sep 17 00:00:00 2001 From: Farhan Nawaz Date: Thu, 10 Sep 2026 20:22:24 +0530 Subject: [PATCH 3/8] fix(rle): address websocket review feedback Preserve positional init naming with an explicit OpenEnv source flag. Add bounded WebSocket keepalives and drain browser responses without allowing client disconnects to corrupt the shared session. Authored-by: GitHub Copilot CLI v1.0.68 Model: GPT-5.4 (gpt-5.4) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- cli/azd/extensions/azure.ai.rle/README.md | 13 ++- .../extensions/azure.ai.rle/extension.yaml | 4 +- .../azure.ai.rle/internal/cmd/init.go | 27 +++-- .../azure.ai.rle/internal/cmd/invoke.go | 4 +- .../azure.ai.rle/internal/cmd/invoke_test.go | 64 +++++++++++ .../azure.ai.rle/internal/cmd/root_test.go | 49 +++++++-- .../azure.ai.rle/internal/project/scaffold.go | 30 ++++-- .../internal/project/scaffold_test.go | 10 +- .../internal/project/websocket_runtime.go | 100 +++++++++++++++--- .../project/websocket_runtime_test.go | 89 ++++++++++++++-- 10 files changed, 334 insertions(+), 56 deletions(-) diff --git a/cli/azd/extensions/azure.ai.rle/README.md b/cli/azd/extensions/azure.ai.rle/README.md index 87c93470feb..619633c632a 100644 --- a/cli/azd/extensions/azure.ai.rle/README.md +++ b/cli/azd/extensions/azure.ai.rle/README.md @@ -95,11 +95,20 @@ The copied session does not keep `.git` metadata from the upstream repository. Copy another environment from the OpenEnv `envs` catalog: ```powershell -azd ai rle init chess_env +azd ai rle init --source chess_env cd .\chess_env ``` -The positional name selects `envs/` from OpenEnv and is also used for the local session directory and RLE environment name. +The `--source` value selects `envs/` from OpenEnv. When the positional environment name is omitted, +the source name is also used for the local session directory and RLE environment name. + +Use a different local and RLE environment name while copying a catalog environment: + +```powershell +azd ai rle init my_chess_env --source chess_env +``` + +For compatibility, a positional name without `--source` still copies `echo_env` into a session with that name. For an existing source folder, skip `init` and run commands directly from that folder. diff --git a/cli/azd/extensions/azure.ai.rle/extension.yaml b/cli/azd/extensions/azure.ai.rle/extension.yaml index eb99f6e9e20..f8d859be73c 100644 --- a/cli/azd/extensions/azure.ai.rle/extension.yaml +++ b/cli/azd/extensions/azure.ai.rle/extension.yaml @@ -14,8 +14,8 @@ usage: azd ai rle [options] version: 0.4.0-preview examples: - name: init - description: Copy an environment from the OpenEnv catalog; defaults to echo_env. - usage: azd ai rle init [environment-name] + description: Initialize an environment from the OpenEnv catalog; defaults to echo_env. + usage: azd ai rle init [environment-name] [--source ] - name: publish description: Build, push, and create or update the RLE environment. usage: azd ai rle publish --version-bump major diff --git a/cli/azd/extensions/azure.ai.rle/internal/cmd/init.go b/cli/azd/extensions/azure.ai.rle/internal/cmd/init.go index 81c45931f2a..be26c683fbd 100644 --- a/cli/azd/extensions/azure.ai.rle/internal/cmd/init.go +++ b/cli/azd/extensions/azure.ai.rle/internal/cmd/init.go @@ -16,13 +16,14 @@ import ( ) type rleInitFlags struct { - force bool + force bool + source string } type initAction struct { cmd *cobra.Command flags *rleInitFlags - environmentName string + envNameOverride string } var checkoutOpenEnvEnvironmentFunc = project.CheckoutOpenEnvEnvironment @@ -35,11 +36,11 @@ func newInitCommand() *cobra.Command { Short: "Initialize a local RLE environment", Args: cobra.MaximumNArgs(1), RunE: func(cmd *cobra.Command, args []string) error { - environmentName := "" + envNameOverride := "" if len(args) == 1 { - environmentName = args[0] + envNameOverride = args[0] } - return (&initAction{cmd: cmd, flags: flags, environmentName: environmentName}).Run() + return (&initAction{cmd: cmd, flags: flags, envNameOverride: envNameOverride}).Run() }, } @@ -50,6 +51,7 @@ func newInitCommand() *cobra.Command { help.WriteString(" rle init [environment-name] [flags]\n") help.WriteString("Flags:\n") help.WriteString(" --force Overwrite generated files in an existing non-empty session directory\n") + help.WriteString(" --source OpenEnv catalog environment to copy (default \"echo_env\")\n") help.WriteString(" -h, --help help for init\n") if cmd.InheritedFlags().HasAvailableFlags() { help.WriteString("Global Flags:\n") @@ -58,11 +60,13 @@ func newInitCommand() *cobra.Command { _, _ = fmt.Fprint(cmd.OutOrStdout(), help.String()) }) cmd.Flags().BoolVar(&flags.force, "force", false, "Overwrite generated files in an existing non-empty session directory") + cmd.Flags().StringVar(&flags.source, "source", "", "OpenEnv catalog environment to copy") return cmd } func (a *initAction) Run() error { - envName := firstNonEmpty(a.environmentName, "echo_env") + sourceName := firstNonEmpty(a.flags.source, "echo_env") + envName := firstNonEmpty(a.envNameOverride, a.flags.source, "echo_env") var err error envName, err = project.ValidateEnvironmentName(envName) if err != nil { @@ -73,8 +77,17 @@ func (a *initAction) Run() error { Suggestion: "Use snake_case starting with a letter, for example code_rl.", } } + sourceName, err = project.ValidateEnvironmentName(sourceName) + if err != nil { + return &azdext.LocalError{ + Message: err.Error(), + Code: "rle_invalid_open_env_source", + Category: azdext.LocalErrorCategoryUser, + Suggestion: "Choose an OpenEnv catalog name in snake_case, for example chess_env.", + } + } - sessionDir, err := checkoutOpenEnvEnvironmentFunc(envName, ".", a.flags.force) + sessionDir, err := checkoutOpenEnvEnvironmentFunc(sourceName, envName, ".", a.flags.force) if err != nil { return err } diff --git a/cli/azd/extensions/azure.ai.rle/internal/cmd/invoke.go b/cli/azd/extensions/azure.ai.rle/internal/cmd/invoke.go index 1905d381653..3c9fb7dd78a 100644 --- a/cli/azd/extensions/azure.ai.rle/internal/cmd/invoke.go +++ b/cli/azd/extensions/azure.ai.rle/internal/cmd/invoke.go @@ -727,13 +727,13 @@ func proxyStatefulOpenEnvOperation( } payload = string(request.Action) } - response, err := runtimeSession.Call(r.Context(), operation, payload) + response, err := runtimeSession.CallAndDrain(r.Context(), operation, payload) if err != nil { http.Error(w, err.Error(), http.StatusBadGateway) return } w.Header().Set("Content-Type", "application/json") - _, _ = io.WriteString(w, response) + _, _ = io.WriteString(w, response) //nolint:gosec // The response is served as JSON, not executable HTML. } func withFoundryAPIVersion(runtimeUrl string) (string, error) { diff --git a/cli/azd/extensions/azure.ai.rle/internal/cmd/invoke_test.go b/cli/azd/extensions/azure.ai.rle/internal/cmd/invoke_test.go index d83bcd00650..0273177d9df 100644 --- a/cli/azd/extensions/azure.ai.rle/internal/cmd/invoke_test.go +++ b/cli/azd/extensions/azure.ai.rle/internal/cmd/invoke_test.go @@ -20,6 +20,7 @@ import ( "slices" "strconv" "strings" + "sync/atomic" "testing" "time" @@ -1036,6 +1037,69 @@ func TestRemotePlaygroundProxyForwardsToSandbox(t *testing.T) { } } +func TestRemotePlaygroundCancellationDoesNotFailSharedSession(t *testing.T) { + firstRequestReceived := make(chan struct{}) + releaseFirstResponse := make(chan struct{}) + var requestCount atomic.Int32 + envServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + connection, err := (&websocket.Upgrader{}).Upgrade(w, r, nil) + if err != nil { + t.Errorf("upgrade WebSocket: %v", err) + return + } + defer connection.Close() + for { + var request map[string]any + if err := connection.ReadJSON(&request); err != nil { + if websocket.IsCloseError(err, websocket.CloseNormalClosure) { + return + } + t.Errorf("read WebSocket request: %v", err) + return + } + currentRequest := requestCount.Add(1) + if currentRequest == 1 { + close(firstRequestReceived) + <-releaseFirstResponse + } + if err := connection.WriteJSON(map[string]any{ + "type": "state", + "data": map[string]any{"request": currentRequest}, + }); err != nil { + t.Errorf("write WebSocket response: %v", err) + return + } + } + })) + defer envServer.Close() + + runtimeSession := project.NewWebSocketRuntimeSession(envServer.URL, 30, nil) + defer runtimeSession.Close() + requestContext, cancelRequest := context.WithCancel(t.Context()) + request := httptest.NewRequest(http.MethodGet, "/state", nil).WithContext(requestContext) + recorder := httptest.NewRecorder() + proxyDone := make(chan struct{}) + go func() { + proxyStatefulOpenEnvOperation(recorder, request, "state", runtimeSession) + close(proxyDone) + }() + + <-firstRequestReceived + cancelRequest() + close(releaseFirstResponse) + <-proxyDone + + if recorder.Code != http.StatusOK { + t.Fatalf("expected canceled browser request to drain successfully, got %d", recorder.Code) + } + if _, err := runtimeSession.Call(t.Context(), "state", ""); err != nil { + t.Fatalf("expected shared session to remain usable: %v", err) + } + if requestCount.Load() != 2 { + t.Fatalf("expected two state requests on the shared session, got %d", requestCount.Load()) + } +} + func TestRemotePlaygroundProxyRefreshesAuthorizationForEachRequest(t *testing.T) { var authorizations []string envServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { diff --git a/cli/azd/extensions/azure.ai.rle/internal/cmd/root_test.go b/cli/azd/extensions/azure.ai.rle/internal/cmd/root_test.go index 3cdf4ceab00..ae8634efdb5 100644 --- a/cli/azd/extensions/azure.ai.rle/internal/cmd/root_test.go +++ b/cli/azd/extensions/azure.ai.rle/internal/cmd/root_test.go @@ -6,6 +6,7 @@ package cmd import ( "bytes" "encoding/json" + "io" "os" "path/filepath" "strings" @@ -120,6 +121,9 @@ func TestLifecycleFlagsAlignWithHostedAgentConventions(t *testing.T) { if flag := initCommand.Flags().Lookup("name"); flag != nil { t.Fatal("expected init not to expose --name") } + if flag := initCommand.Flags().Lookup("source"); flag == nil { + t.Fatal("expected init to expose --source") + } runCommand, _, err := rootCmd.Find([]string{"run"}) if err != nil { @@ -280,13 +284,13 @@ func TestInitCopiesOpenEnvEchoSampleByDefault(t *testing.T) { } } -func TestInitUsesPositionalNameToSelectOpenEnvEnvironment(t *testing.T) { +func TestInitPreservesPositionalEnvironmentName(t *testing.T) { tempDir := t.TempDir() t.Chdir(tempDir) - stubOpenEnvCheckout(t, "chess_env") + stubOpenEnvCheckout(t, "echo_env", "code_rl") command := newInitCommand() - command.SetArgs([]string{"chess_env"}) + command.SetArgs([]string{"code_rl"}) var output bytes.Buffer command.SetOut(&output) command.SetErr(&output) @@ -294,7 +298,7 @@ func TestInitUsesPositionalNameToSelectOpenEnvEnvironment(t *testing.T) { t.Fatal(err) } - sessionDir := filepath.Join(tempDir, "chess_env") + sessionDir := filepath.Join(tempDir, "code_rl") // The test reads state from its own temporary session directory. stateBytes, err := os.ReadFile(filepath.Join(sessionDir, rleStateFile)) //nolint:gosec if err != nil { @@ -304,6 +308,32 @@ func TestInitUsesPositionalNameToSelectOpenEnvEnvironment(t *testing.T) { if err := json.Unmarshal(stateBytes, &state); err != nil { t.Fatal(err) } + if state.EnvironmentName != "code_rl" { + t.Fatalf("expected code_rl environment name, got %q", state.EnvironmentName) + } +} + +func TestInitUsesOpenEnvSourceAsDefaultEnvironmentName(t *testing.T) { + tempDir := t.TempDir() + t.Chdir(tempDir) + stubOpenEnvCheckout(t, "chess_env", "chess_env") + + command := newInitCommand() + command.SetArgs([]string{"--source", "chess_env"}) + command.SetOut(io.Discard) + command.SetErr(io.Discard) + if err := command.Execute(); err != nil { + t.Fatal(err) + } + + stateBytes, err := os.ReadFile(filepath.Join(tempDir, "chess_env", rleStateFile)) //nolint:gosec + if err != nil { + t.Fatal(err) + } + var state rleState + if err := json.Unmarshal(stateBytes, &state); err != nil { + t.Fatal(err) + } if state.EnvironmentName != "chess_env" { t.Fatalf("expected chess_env environment name, got %q", state.EnvironmentName) } @@ -368,12 +398,15 @@ func TestInitNextStepsUseShellAppropriateSyntax(t *testing.T) { } } -func stubOpenEnvCheckout(t *testing.T, expectedName string) { +func stubOpenEnvCheckout(t *testing.T, expectedSource string, expectedName ...string) { t.Helper() old := checkoutOpenEnvEnvironmentFunc - checkoutOpenEnvEnvironmentFunc = func(name string, dest string, force bool) (string, error) { - if name != expectedName { - t.Fatalf("expected OpenEnv environment %q, got %q", expectedName, name) + checkoutOpenEnvEnvironmentFunc = func(sourceName string, name string, dest string, force bool) (string, error) { + if sourceName != expectedSource { + t.Fatalf("expected OpenEnv source %q, got %q", expectedSource, sourceName) + } + if len(expectedName) == 1 && name != expectedName[0] { + t.Fatalf("expected environment name %q, got %q", expectedName[0], name) } sessionDir := filepath.Join(dest, name) if force { diff --git a/cli/azd/extensions/azure.ai.rle/internal/project/scaffold.go b/cli/azd/extensions/azure.ai.rle/internal/project/scaffold.go index 71d51959a52..47987fb1db0 100644 --- a/cli/azd/extensions/azure.ai.rle/internal/project/scaffold.go +++ b/cli/azd/extensions/azure.ai.rle/internal/project/scaffold.go @@ -40,12 +40,16 @@ func createRleSessionDir(name string, dest string, force bool) (string, error) { return sessionDir, nil } -func CheckoutOpenEnvEnvironment(name string, dest string, force bool) (string, error) { - name, err := ValidateEnvironmentName(name) +func CheckoutOpenEnvEnvironment(sourceName string, name string, dest string, force bool) (string, error) { + sourceName, err := ValidateEnvironmentName(sourceName) if err != nil { return "", err } - sourcePath := openEnvEnvironmentPath(name) + name, err = ValidateEnvironmentName(name) + if err != nil { + return "", err + } + sourcePath := openEnvEnvironmentPath(sourceName) tempDir, err := os.MkdirTemp("", "azd-rle-open-env-*") if err != nil { return "", err @@ -70,16 +74,21 @@ func CheckoutOpenEnvEnvironment(name string, dest string, force bool) (string, e } sourceDir := filepath.Join(tempDir, filepath.FromSlash(sourcePath)) - return copyOpenEnvEnvironment(sourceDir, name, dest, force) + return copyOpenEnvEnvironment(sourceDir, sourceName, name, dest, force) } -func copyOpenEnvEnvironment(sourceDir string, name string, dest string, force bool) (string, error) { +func copyOpenEnvEnvironment(sourceDir string, sourceName string, name string, dest string, force bool) (string, error) { if _, err := os.Stat(sourceDir); os.IsNotExist(err) { + catalogURL := strings.TrimSuffix(openEnvRepoUrl, ".git") return "", &azdext.LocalError{ - Message: fmt.Sprintf("OpenEnv environment %q was not found.", name), - Code: "rle_open_env_environment_not_found", - Category: azdext.LocalErrorCategoryUser, - Suggestion: fmt.Sprintf("Choose an environment from %s/tree/%s/envs.", strings.TrimSuffix(openEnvRepoUrl, ".git"), openEnvRepoRef), + Message: fmt.Sprintf("OpenEnv environment %q was not found.", sourceName), + Code: "rle_open_env_environment_not_found", + Category: azdext.LocalErrorCategoryUser, + Suggestion: fmt.Sprintf( + "Choose an environment from %s/tree/%s/envs.", + catalogURL, + openEnvRepoRef, + ), } } else if err != nil { return "", err @@ -107,7 +116,8 @@ func runGitCheckout(args ...string) error { Suggestion: "Install Git, then retry azd ai rle init.", } } - process := exec.Command("git", args...) //nolint:gosec // Arguments are passed directly; the user value is validated as an environment name. + // The user-provided sparse path is validated as an environment name before reaching this command. + process := exec.Command("git", args...) //nolint:gosec process.Env = os.Environ() output, err := process.CombinedOutput() if err != nil { diff --git a/cli/azd/extensions/azure.ai.rle/internal/project/scaffold_test.go b/cli/azd/extensions/azure.ai.rle/internal/project/scaffold_test.go index 15eaee024bf..76a6842971f 100644 --- a/cli/azd/extensions/azure.ai.rle/internal/project/scaffold_test.go +++ b/cli/azd/extensions/azure.ai.rle/internal/project/scaffold_test.go @@ -78,7 +78,7 @@ func TestCheckoutOpenEnvEnvironmentRejectsInvalidNameBeforeChangingDestination(t t.Fatal(err) } - if _, err := CheckoutOpenEnvEnvironment("../bad", destDir, true); err == nil { + if _, err := CheckoutOpenEnvEnvironment("../bad", "target", destDir, true); err == nil { t.Fatal("expected invalid environment name to be rejected") } @@ -104,7 +104,13 @@ func TestCopyOpenEnvEnvironmentValidatesSourceBeforeReplacingDestination(t *test t.Fatal(err) } - _, err := copyOpenEnvEnvironment(filepath.Join(t.TempDir(), "missing"), "missing_env", destDir, true) + _, err := copyOpenEnvEnvironment( + filepath.Join(t.TempDir(), "missing"), + "missing_env", + "target_env", + destDir, + true, + ) if err == nil { t.Fatal("expected missing OpenEnv environment to fail") } diff --git a/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime.go b/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime.go index bb1f70292b7..5d5d15f80f3 100644 --- a/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime.go +++ b/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime.go @@ -18,7 +18,12 @@ import ( "github.com/gorilla/websocket" ) -const maxWebSocketMessageBytes = 100 * 1024 * 1024 +const ( + maxWebSocketMessageBytes = 100 * 1024 * 1024 + webSocketPingInterval = 20 * time.Second + webSocketPingTimeout = 20 * time.Second + webSocketDrainTimeout = 60 * time.Second +) type WebSocketRuntimeSession struct { baseURL string @@ -27,6 +32,9 @@ type WebSocketRuntimeSession struct { mu sync.Mutex exchangeMu sync.Mutex connection *websocket.Conn + connectionDone chan struct{} + keepAliveInterval time.Duration + drainTimeout time.Duration terminalError error closed bool } @@ -40,6 +48,8 @@ func NewWebSocketRuntimeSession( baseURL: baseURL, timeout: timeout, authorizationProvider: authorizationProvider, + keepAliveInterval: webSocketPingInterval, + drainTimeout: webSocketDrainTimeout, } } @@ -95,7 +105,7 @@ func (c *WebSocketRuntimeSession) call( ) (string, error) { switch operation { case "reset", "step", "state": - return c.exchange(ctx, operation, flags) + return c.exchange(ctx, operation, flags, true) default: return call(ctx, baseURL, operation, flags, authorizationProvider) } @@ -116,10 +126,28 @@ func (c *WebSocketRuntimeSession) Call( return c.call(ctx, c.baseURL, operation, flags, c.authorizationProvider) } +// CallAndDrain preserves cancellation until a request is sent, then drains its response +// so a disconnected HTTP client cannot desynchronize the shared WebSocket session. +func (c *WebSocketRuntimeSession) CallAndDrain( + ctx context.Context, + operation string, + payload string, +) (string, error) { + flags := &callOptions{timeout: c.timeout} + switch operation { + case "reset": + flags.body = payload + case "step": + flags.action = payload + } + return c.exchange(ctx, operation, flags, false) +} + func (c *WebSocketRuntimeSession) exchange( ctx context.Context, operation string, flags *callOptions, + cancelAfterSend bool, ) (string, error) { c.exchangeMu.Lock() defer c.exchangeMu.Unlock() @@ -146,6 +174,10 @@ func (c *WebSocketRuntimeSession) exchange( } deadline, hasDeadline := operationDeadline(ctx, c.timeout) + if !cancelAfterSend && !hasDeadline { + deadline = time.Now().Add(c.drainTimeout) + hasDeadline = true + } if hasDeadline { if err := connection.SetWriteDeadline(deadline); err != nil { return "", c.failConnection(connection, fmt.Errorf("set OpenEnv WebSocket write deadline: %w", err)) @@ -164,22 +196,28 @@ func (c *WebSocketRuntimeSession) exchange( fmt.Errorf("send OpenEnv WebSocket %s request: %w", operation, err), ) } - exchangeDone := make(chan struct{}) - cancellationHandled := make(chan struct{}) - go func() { - select { - case <-exchangeDone: - case <-ctx.Done(): - c.failConnection( - connection, - fmt.Errorf("OpenEnv WebSocket %s request canceled: %w", operation, ctx.Err()), - ) - } - close(cancellationHandled) - }() + var exchangeDone chan struct{} + var cancellationHandled chan struct{} + if cancelAfterSend { + exchangeDone = make(chan struct{}) + cancellationHandled = make(chan struct{}) + go func() { + select { + case <-exchangeDone: + case <-ctx.Done(): + _ = c.failConnection( + connection, + fmt.Errorf("OpenEnv WebSocket %s request canceled: %w", operation, ctx.Err()), + ) + } + close(cancellationHandled) + }() + } messageType, response, err := connection.ReadMessage() - close(exchangeDone) - <-cancellationHandled + if cancelAfterSend { + close(exchangeDone) + <-cancellationHandled + } if err != nil { return "", c.failConnection( connection, @@ -254,13 +292,33 @@ func (c *WebSocketRuntimeSession) connect(ctx context.Context) error { } connection.SetReadLimit(maxWebSocketMessageBytes) c.connection = connection + c.connectionDone = make(chan struct{}) + go c.keepAlive(connection, c.connectionDone) return nil } +func (c *WebSocketRuntimeSession) keepAlive(connection *websocket.Conn, done <-chan struct{}) { + ticker := time.NewTicker(c.keepAliveInterval) + defer ticker.Stop() + for { + select { + case <-done: + return + case <-ticker.C: + deadline := time.Now().Add(webSocketPingTimeout) + if err := connection.WriteControl(websocket.PingMessage, nil, deadline); err != nil { + _ = c.failConnection(connection, fmt.Errorf("send OpenEnv WebSocket keepalive: %w", err)) + return + } + } + } +} + func (c *WebSocketRuntimeSession) failConnection(connection *websocket.Conn, err error) error { c.mu.Lock() if c.connection == connection { c.connection = nil + c.stopKeepAliveLocked() } if c.terminalError == nil { c.terminalError = &azdext.LocalError{ @@ -280,6 +338,7 @@ func (c *WebSocketRuntimeSession) Close() { c.mu.Lock() connection := c.connection c.connection = nil + c.stopKeepAliveLocked() c.closed = true if c.terminalError == nil { c.terminalError = &azdext.LocalError{ @@ -302,6 +361,13 @@ func (c *WebSocketRuntimeSession) Close() { _ = connection.Close() } +func (c *WebSocketRuntimeSession) stopKeepAliveLocked() { + if c.connectionDone != nil { + close(c.connectionDone) + c.connectionDone = nil + } +} + func parseWebSocketResponse(operation string, response []byte) (string, bool, error) { var envelope struct { Type string `json:"type"` diff --git a/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime_test.go b/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime_test.go index fb8e368a3c5..cd870bba061 100644 --- a/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime_test.go +++ b/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime_test.go @@ -13,6 +13,7 @@ import ( "slices" "strings" "testing" + "time" "github.com/gorilla/websocket" ) @@ -31,7 +32,7 @@ func TestRunWebSocketShellUsesPersistentSocketForStatefulOperations(t *testing.T } if r.URL.Path != "/ws" { safePaths = append(safePaths, r.URL.Path) - _, _ = fmt.Fprintf(w, `{"path":%q}`, r.URL.Path) + _, _ = fmt.Fprintf(w, `{"path":%q}`, r.URL.Path) //nolint:gosec // Test response uses JSON encoding. return } @@ -45,10 +46,6 @@ func TestRunWebSocketShellUsesPersistentSocketForStatefulOperations(t *testing.T for { var request map[string]any if err := connection.ReadJSON(&request); err != nil { - if websocket.IsCloseError(err, websocket.CloseNormalClosure) { - return - } - t.Errorf("read WebSocket request: %v", err) return } requests = append(requests, request) @@ -136,6 +133,58 @@ func TestRuntimeWebSocketURL(t *testing.T) { } } +func TestWebSocketSessionSendsKeepalivePings(t *testing.T) { + pingReceived := make(chan struct{}, 1) + serverDone := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + defer close(serverDone) + connection, err := (&websocket.Upgrader{}).Upgrade(w, r, nil) + if err != nil { + t.Errorf("upgrade WebSocket: %v", err) + return + } + defer connection.Close() + connection.SetPingHandler(func(data string) error { + select { + case pingReceived <- struct{}{}: + default: + } + return connection.WriteControl( + websocket.PongMessage, + []byte(data), + time.Now().Add(time.Second), + ) + }) + for { + var request map[string]any + if err := connection.ReadJSON(&request); err != nil { + return + } + if err := connection.WriteJSON(map[string]any{ + "type": "state", + "data": map[string]any{}, + }); err != nil { + t.Errorf("write WebSocket response: %v", err) + return + } + } + })) + defer server.Close() + + session := NewWebSocketRuntimeSession(server.URL, 30, nil) + session.keepAliveInterval = 10 * time.Millisecond + if _, err := session.Call(t.Context(), "state", ""); err != nil { + t.Fatal(err) + } + select { + case <-pingReceived: + case <-time.After(time.Second): + t.Fatal("expected WebSocket keepalive ping") + } + session.Close() + <-serverDone +} + func TestWebSocketRequestUsesOpenEnvProtocol(t *testing.T) { reset, err := webSocketRequest("reset", &callOptions{}) if err != nil { @@ -270,7 +319,7 @@ func TestCanceledQueuedCallDoesNotReachWebSocket(t *testing.T) { queuedContext, cancelQueued := context.WithCancel(t.Context()) queuedCallDone := make(chan error, 1) go func() { - _, err := session.Call(queuedContext, "state", "") + _, err := session.CallAndDrain(queuedContext, "state", "") queuedCallDone <- err }() cancelQueued() @@ -287,6 +336,34 @@ func TestCanceledQueuedCallDoesNotReachWebSocket(t *testing.T) { } } +func TestCallAndDrainUsesBoundedReadWhenTimeoutDisabled(t *testing.T) { + requestReceived := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + connection, err := (&websocket.Upgrader{}).Upgrade(w, r, nil) + if err != nil { + t.Errorf("upgrade WebSocket: %v", err) + return + } + defer connection.Close() + if _, _, err := connection.ReadMessage(); err != nil { + t.Errorf("read WebSocket request: %v", err) + return + } + close(requestReceived) + _, _, _ = connection.ReadMessage() + })) + defer server.Close() + + session := NewWebSocketRuntimeSession(server.URL, 0, nil) + session.drainTimeout = 10 * time.Millisecond + defer session.Close() + _, err := session.CallAndDrain(t.Context(), "state", "") + if err == nil { + t.Fatal("expected bounded drain to fail when the runtime does not respond") + } + <-requestReceived +} + func assertWebSocketRequest( t *testing.T, request map[string]any, From f8c712fd44fb782138f8f044ac95af7d02cd53d5 Mon Sep 17 00:00:00 2001 From: Farhan Nawaz Date: Thu, 10 Sep 2026 20:28:22 +0530 Subject: [PATCH 4/8] fix(rle): satisfy extension spell check Use repository-recognized wording in the WebSocket drain comment. Authored-by: GitHub Copilot CLI v1.0.68 Model: GPT-5.4 (gpt-5.4) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../azure.ai.rle/internal/project/websocket_runtime.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime.go b/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime.go index 5d5d15f80f3..59b955f2ab7 100644 --- a/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime.go +++ b/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime.go @@ -127,7 +127,7 @@ func (c *WebSocketRuntimeSession) Call( } // CallAndDrain preserves cancellation until a request is sent, then drains its response -// so a disconnected HTTP client cannot desynchronize the shared WebSocket session. +// so a disconnected HTTP client cannot disrupt the shared WebSocket protocol sequence. func (c *WebSocketRuntimeSession) CallAndDrain( ctx context.Context, operation string, From c83edda87ff89ade9b94662a3d1d01f28de83ac8 Mon Sep 17 00:00:00 2001 From: Farhan Nawaz Date: Thu, 10 Sep 2026 22:40:34 +0530 Subject: [PATCH 5/8] fix(rle): harden catalog init and websocket errors Use the positional init argument as the OpenEnv catalog selector and suggest close matches without replacing existing destinations. Treat connection-level OpenEnv errors as terminal and align client timeouts and message limits with the service contract. Authored-by: GitHub Copilot CLI v1.0.68 Model: GPT-5.4 (gpt-5.4) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- cli/azd/extensions/azure.ai.rle/README.md | 14 +- .../extensions/azure.ai.rle/extension.yaml | 4 +- .../azure.ai.rle/internal/cmd/init.go | 27 +--- .../azure.ai.rle/internal/cmd/invoke.go | 2 +- .../azure.ai.rle/internal/cmd/root_test.go | 49 ++----- .../azure.ai.rle/internal/project/scaffold.go | 120 ++++++++++++++---- .../internal/project/scaffold_test.go | 32 +++-- .../internal/project/websocket_runtime.go | 54 ++++++-- .../project/websocket_runtime_test.go | 60 +++++++++ 9 files changed, 247 insertions(+), 115 deletions(-) diff --git a/cli/azd/extensions/azure.ai.rle/README.md b/cli/azd/extensions/azure.ai.rle/README.md index 619633c632a..53ed86a3fa1 100644 --- a/cli/azd/extensions/azure.ai.rle/README.md +++ b/cli/azd/extensions/azure.ai.rle/README.md @@ -95,20 +95,12 @@ The copied session does not keep `.git` metadata from the upstream repository. Copy another environment from the OpenEnv `envs` catalog: ```powershell -azd ai rle init --source chess_env +azd ai rle init chess_env cd .\chess_env ``` -The `--source` value selects `envs/` from OpenEnv. When the positional environment name is omitted, -the source name is also used for the local session directory and RLE environment name. - -Use a different local and RLE environment name while copying a catalog environment: - -```powershell -azd ai rle init my_chess_env --source chess_env -``` - -For compatibility, a positional name without `--source` still copies `echo_env` into a session with that name. +The positional name selects `envs/` from OpenEnv and is also used for the local session directory +and RLE environment name. If the name is not present in the catalog, `init` suggests the closest available name. For an existing source folder, skip `init` and run commands directly from that folder. diff --git a/cli/azd/extensions/azure.ai.rle/extension.yaml b/cli/azd/extensions/azure.ai.rle/extension.yaml index f8d859be73c..eb99f6e9e20 100644 --- a/cli/azd/extensions/azure.ai.rle/extension.yaml +++ b/cli/azd/extensions/azure.ai.rle/extension.yaml @@ -14,8 +14,8 @@ usage: azd ai rle [options] version: 0.4.0-preview examples: - name: init - description: Initialize an environment from the OpenEnv catalog; defaults to echo_env. - usage: azd ai rle init [environment-name] [--source ] + description: Copy an environment from the OpenEnv catalog; defaults to echo_env. + usage: azd ai rle init [environment-name] - name: publish description: Build, push, and create or update the RLE environment. usage: azd ai rle publish --version-bump major diff --git a/cli/azd/extensions/azure.ai.rle/internal/cmd/init.go b/cli/azd/extensions/azure.ai.rle/internal/cmd/init.go index be26c683fbd..81c45931f2a 100644 --- a/cli/azd/extensions/azure.ai.rle/internal/cmd/init.go +++ b/cli/azd/extensions/azure.ai.rle/internal/cmd/init.go @@ -16,14 +16,13 @@ import ( ) type rleInitFlags struct { - force bool - source string + force bool } type initAction struct { cmd *cobra.Command flags *rleInitFlags - envNameOverride string + environmentName string } var checkoutOpenEnvEnvironmentFunc = project.CheckoutOpenEnvEnvironment @@ -36,11 +35,11 @@ func newInitCommand() *cobra.Command { Short: "Initialize a local RLE environment", Args: cobra.MaximumNArgs(1), RunE: func(cmd *cobra.Command, args []string) error { - envNameOverride := "" + environmentName := "" if len(args) == 1 { - envNameOverride = args[0] + environmentName = args[0] } - return (&initAction{cmd: cmd, flags: flags, envNameOverride: envNameOverride}).Run() + return (&initAction{cmd: cmd, flags: flags, environmentName: environmentName}).Run() }, } @@ -51,7 +50,6 @@ func newInitCommand() *cobra.Command { help.WriteString(" rle init [environment-name] [flags]\n") help.WriteString("Flags:\n") help.WriteString(" --force Overwrite generated files in an existing non-empty session directory\n") - help.WriteString(" --source OpenEnv catalog environment to copy (default \"echo_env\")\n") help.WriteString(" -h, --help help for init\n") if cmd.InheritedFlags().HasAvailableFlags() { help.WriteString("Global Flags:\n") @@ -60,13 +58,11 @@ func newInitCommand() *cobra.Command { _, _ = fmt.Fprint(cmd.OutOrStdout(), help.String()) }) cmd.Flags().BoolVar(&flags.force, "force", false, "Overwrite generated files in an existing non-empty session directory") - cmd.Flags().StringVar(&flags.source, "source", "", "OpenEnv catalog environment to copy") return cmd } func (a *initAction) Run() error { - sourceName := firstNonEmpty(a.flags.source, "echo_env") - envName := firstNonEmpty(a.envNameOverride, a.flags.source, "echo_env") + envName := firstNonEmpty(a.environmentName, "echo_env") var err error envName, err = project.ValidateEnvironmentName(envName) if err != nil { @@ -77,17 +73,8 @@ func (a *initAction) Run() error { Suggestion: "Use snake_case starting with a letter, for example code_rl.", } } - sourceName, err = project.ValidateEnvironmentName(sourceName) - if err != nil { - return &azdext.LocalError{ - Message: err.Error(), - Code: "rle_invalid_open_env_source", - Category: azdext.LocalErrorCategoryUser, - Suggestion: "Choose an OpenEnv catalog name in snake_case, for example chess_env.", - } - } - sessionDir, err := checkoutOpenEnvEnvironmentFunc(sourceName, envName, ".", a.flags.force) + sessionDir, err := checkoutOpenEnvEnvironmentFunc(envName, ".", a.flags.force) if err != nil { return err } diff --git a/cli/azd/extensions/azure.ai.rle/internal/cmd/invoke.go b/cli/azd/extensions/azure.ai.rle/internal/cmd/invoke.go index 3c9fb7dd78a..89778ad39fc 100644 --- a/cli/azd/extensions/azure.ai.rle/internal/cmd/invoke.go +++ b/cli/azd/extensions/azure.ai.rle/internal/cmd/invoke.go @@ -41,7 +41,7 @@ var validateSandboxURL = validateRemoteSandboxURL func newInvokeCommand() *cobra.Command { flags := &remoteInvokeFlags{ - timeout: 30, + timeout: 60, } cmd := &cobra.Command{ diff --git a/cli/azd/extensions/azure.ai.rle/internal/cmd/root_test.go b/cli/azd/extensions/azure.ai.rle/internal/cmd/root_test.go index ae8634efdb5..3cdf4ceab00 100644 --- a/cli/azd/extensions/azure.ai.rle/internal/cmd/root_test.go +++ b/cli/azd/extensions/azure.ai.rle/internal/cmd/root_test.go @@ -6,7 +6,6 @@ package cmd import ( "bytes" "encoding/json" - "io" "os" "path/filepath" "strings" @@ -121,9 +120,6 @@ func TestLifecycleFlagsAlignWithHostedAgentConventions(t *testing.T) { if flag := initCommand.Flags().Lookup("name"); flag != nil { t.Fatal("expected init not to expose --name") } - if flag := initCommand.Flags().Lookup("source"); flag == nil { - t.Fatal("expected init to expose --source") - } runCommand, _, err := rootCmd.Find([]string{"run"}) if err != nil { @@ -284,13 +280,13 @@ func TestInitCopiesOpenEnvEchoSampleByDefault(t *testing.T) { } } -func TestInitPreservesPositionalEnvironmentName(t *testing.T) { +func TestInitUsesPositionalNameToSelectOpenEnvEnvironment(t *testing.T) { tempDir := t.TempDir() t.Chdir(tempDir) - stubOpenEnvCheckout(t, "echo_env", "code_rl") + stubOpenEnvCheckout(t, "chess_env") command := newInitCommand() - command.SetArgs([]string{"code_rl"}) + command.SetArgs([]string{"chess_env"}) var output bytes.Buffer command.SetOut(&output) command.SetErr(&output) @@ -298,7 +294,7 @@ func TestInitPreservesPositionalEnvironmentName(t *testing.T) { t.Fatal(err) } - sessionDir := filepath.Join(tempDir, "code_rl") + sessionDir := filepath.Join(tempDir, "chess_env") // The test reads state from its own temporary session directory. stateBytes, err := os.ReadFile(filepath.Join(sessionDir, rleStateFile)) //nolint:gosec if err != nil { @@ -308,32 +304,6 @@ func TestInitPreservesPositionalEnvironmentName(t *testing.T) { if err := json.Unmarshal(stateBytes, &state); err != nil { t.Fatal(err) } - if state.EnvironmentName != "code_rl" { - t.Fatalf("expected code_rl environment name, got %q", state.EnvironmentName) - } -} - -func TestInitUsesOpenEnvSourceAsDefaultEnvironmentName(t *testing.T) { - tempDir := t.TempDir() - t.Chdir(tempDir) - stubOpenEnvCheckout(t, "chess_env", "chess_env") - - command := newInitCommand() - command.SetArgs([]string{"--source", "chess_env"}) - command.SetOut(io.Discard) - command.SetErr(io.Discard) - if err := command.Execute(); err != nil { - t.Fatal(err) - } - - stateBytes, err := os.ReadFile(filepath.Join(tempDir, "chess_env", rleStateFile)) //nolint:gosec - if err != nil { - t.Fatal(err) - } - var state rleState - if err := json.Unmarshal(stateBytes, &state); err != nil { - t.Fatal(err) - } if state.EnvironmentName != "chess_env" { t.Fatalf("expected chess_env environment name, got %q", state.EnvironmentName) } @@ -398,15 +368,12 @@ func TestInitNextStepsUseShellAppropriateSyntax(t *testing.T) { } } -func stubOpenEnvCheckout(t *testing.T, expectedSource string, expectedName ...string) { +func stubOpenEnvCheckout(t *testing.T, expectedName string) { t.Helper() old := checkoutOpenEnvEnvironmentFunc - checkoutOpenEnvEnvironmentFunc = func(sourceName string, name string, dest string, force bool) (string, error) { - if sourceName != expectedSource { - t.Fatalf("expected OpenEnv source %q, got %q", expectedSource, sourceName) - } - if len(expectedName) == 1 && name != expectedName[0] { - t.Fatalf("expected environment name %q, got %q", expectedName[0], name) + checkoutOpenEnvEnvironmentFunc = func(name string, dest string, force bool) (string, error) { + if name != expectedName { + t.Fatalf("expected OpenEnv environment %q, got %q", expectedName, name) } sessionDir := filepath.Join(dest, name) if force { diff --git a/cli/azd/extensions/azure.ai.rle/internal/project/scaffold.go b/cli/azd/extensions/azure.ai.rle/internal/project/scaffold.go index 47987fb1db0..17d7838a10b 100644 --- a/cli/azd/extensions/azure.ai.rle/internal/project/scaffold.go +++ b/cli/azd/extensions/azure.ai.rle/internal/project/scaffold.go @@ -40,16 +40,12 @@ func createRleSessionDir(name string, dest string, force bool) (string, error) { return sessionDir, nil } -func CheckoutOpenEnvEnvironment(sourceName string, name string, dest string, force bool) (string, error) { - sourceName, err := ValidateEnvironmentName(sourceName) +func CheckoutOpenEnvEnvironment(name string, dest string, force bool) (string, error) { + name, err := ValidateEnvironmentName(name) if err != nil { return "", err } - name, err = ValidateEnvironmentName(name) - if err != nil { - return "", err - } - sourcePath := openEnvEnvironmentPath(sourceName) + sourcePath := openEnvEnvironmentPath(name) tempDir, err := os.MkdirTemp("", "azd-rle-open-env-*") if err != nil { return "", err @@ -69,27 +65,24 @@ func CheckoutOpenEnvEnvironment(sourceName string, name string, dest string, for ); err != nil { return "", err } + environmentNames, err := listOpenEnvEnvironments(tempDir) + if err != nil { + return "", err + } + if !containsString(environmentNames, name) { + return "", openEnvEnvironmentNotFoundError(name, environmentNames) + } if err := runGitCheckout("-C", tempDir, "sparse-checkout", "set", sourcePath); err != nil { return "", err } sourceDir := filepath.Join(tempDir, filepath.FromSlash(sourcePath)) - return copyOpenEnvEnvironment(sourceDir, sourceName, name, dest, force) + return copyOpenEnvEnvironment(sourceDir, name, dest, force) } -func copyOpenEnvEnvironment(sourceDir string, sourceName string, name string, dest string, force bool) (string, error) { +func copyOpenEnvEnvironment(sourceDir string, name string, dest string, force bool) (string, error) { if _, err := os.Stat(sourceDir); os.IsNotExist(err) { - catalogURL := strings.TrimSuffix(openEnvRepoUrl, ".git") - return "", &azdext.LocalError{ - Message: fmt.Sprintf("OpenEnv environment %q was not found.", sourceName), - Code: "rle_open_env_environment_not_found", - Category: azdext.LocalErrorCategoryUser, - Suggestion: fmt.Sprintf( - "Choose an environment from %s/tree/%s/envs.", - catalogURL, - openEnvRepoRef, - ), - } + return "", openEnvEnvironmentNotFoundError(name, nil) } else if err != nil { return "", err } @@ -107,9 +100,90 @@ func openEnvEnvironmentPath(name string) string { return "envs/" + name } +func listOpenEnvEnvironments(repoDir string) ([]string, error) { + output, err := runGitCommand("-C", repoDir, "ls-tree", "--name-only", openEnvRepoRef+":envs") + if err != nil { + return nil, err + } + return strings.Fields(string(output)), nil +} + +func openEnvEnvironmentNotFoundError(name string, environmentNames []string) error { + catalogURL := strings.TrimSuffix(openEnvRepoUrl, ".git") + suggestion := fmt.Sprintf( + "Choose an environment from %s/tree/%s/envs.", + catalogURL, + openEnvRepoRef, + ) + if closest := closestEnvironmentName(name, environmentNames); closest != "" { + suggestion = fmt.Sprintf("Did you mean %q? %s", closest, suggestion) + } + return &azdext.LocalError{ + Message: fmt.Sprintf("OpenEnv environment %q was not found.", name), + Code: "rle_open_env_environment_not_found", + Category: azdext.LocalErrorCategoryUser, + Suggestion: suggestion, + } +} + +func closestEnvironmentName(name string, environmentNames []string) string { + closest := "" + closestDistance := len(name) + 1 + for _, candidate := range environmentNames { + distance := editDistance(name, candidate) + if distance < closestDistance { + closest = candidate + closestDistance = distance + } + } + maxDistance := max(2, len(name)/3) + if closestDistance > maxDistance { + return "" + } + return closest +} + +func editDistance(left string, right string) int { + previous := make([]int, len(right)+1) + for index := range previous { + previous[index] = index + } + for leftIndex, leftRune := range left { + current := make([]int, len(right)+1) + current[0] = leftIndex + 1 + for rightIndex, rightRune := range right { + substitutionCost := 0 + if leftRune != rightRune { + substitutionCost = 1 + } + current[rightIndex+1] = min( + current[rightIndex]+1, + previous[rightIndex+1]+1, + previous[rightIndex]+substitutionCost, + ) + } + previous = current + } + return previous[len(right)] +} + +func containsString(values []string, expected string) bool { + for _, value := range values { + if value == expected { + return true + } + } + return false +} + func runGitCheckout(args ...string) error { + _, err := runGitCommand(args...) + return err +} + +func runGitCommand(args ...string) ([]byte, error) { if _, err := exec.LookPath("git"); err != nil { - return &azdext.LocalError{ + return nil, &azdext.LocalError{ Message: "Could not find \"git\" on PATH.", Code: "rle_git_not_found", Category: azdext.LocalErrorCategoryUser, @@ -121,14 +195,14 @@ func runGitCheckout(args ...string) error { process.Env = os.Environ() output, err := process.CombinedOutput() if err != nil { - return &azdext.LocalError{ + return nil, &azdext.LocalError{ Message: fmt.Sprintf("Failed to checkout OpenEnv environment: %v", err), Code: "rle_open_env_checkout_failed", Category: azdext.LocalErrorCategoryUser, Suggestion: strings.TrimSpace(string(output)), } } - return nil + return output, nil } func copyDirectory(sourceDir string, destDir string) error { diff --git a/cli/azd/extensions/azure.ai.rle/internal/project/scaffold_test.go b/cli/azd/extensions/azure.ai.rle/internal/project/scaffold_test.go index 76a6842971f..0915164193e 100644 --- a/cli/azd/extensions/azure.ai.rle/internal/project/scaffold_test.go +++ b/cli/azd/extensions/azure.ai.rle/internal/project/scaffold_test.go @@ -4,10 +4,14 @@ package project import ( + "errors" "os" "path/filepath" "runtime" + "strings" "testing" + + "github.com/azure/azure-dev/cli/azd/pkg/azdext" ) func TestCopyDirectorySkipsGitMetadata(t *testing.T) { @@ -78,7 +82,7 @@ func TestCheckoutOpenEnvEnvironmentRejectsInvalidNameBeforeChangingDestination(t t.Fatal(err) } - if _, err := CheckoutOpenEnvEnvironment("../bad", "target", destDir, true); err == nil { + if _, err := CheckoutOpenEnvEnvironment("../bad", destDir, true); err == nil { t.Fatal("expected invalid environment name to be rejected") } @@ -104,13 +108,7 @@ func TestCopyOpenEnvEnvironmentValidatesSourceBeforeReplacingDestination(t *test t.Fatal(err) } - _, err := copyOpenEnvEnvironment( - filepath.Join(t.TempDir(), "missing"), - "missing_env", - "target_env", - destDir, - true, - ) + _, err := copyOpenEnvEnvironment(filepath.Join(t.TempDir(), "missing"), "missing_env", destDir, true) if err == nil { t.Fatal("expected missing OpenEnv environment to fail") } @@ -118,3 +116,21 @@ func TestCopyOpenEnvEnvironmentValidatesSourceBeforeReplacingDestination(t *test t.Fatalf("expected destination to remain unchanged after catalog lookup failure: %v", statErr) } } + +func TestOpenEnvEnvironmentNotFoundSuggestsClosestCatalogName(t *testing.T) { + err := openEnvEnvironmentNotFoundError( + "ches_env", + []string{"atari_env", "chess_env", "echo_env"}, + ) + var localError *azdext.LocalError + if !errors.As(err, &localError) || + !strings.Contains(localError.Suggestion, `Did you mean "chess_env"?`) { + t.Fatalf("expected closest catalog suggestion, got %v", err) + } +} + +func TestClosestEnvironmentNameRejectsDistantMatch(t *testing.T) { + if got := closestEnvironmentName("unknown_env", []string{"chess_env", "echo_env"}); got != "" { + t.Fatalf("expected no distant suggestion, got %q", got) + } +} diff --git a/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime.go b/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime.go index 59b955f2ab7..09697381c04 100644 --- a/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime.go +++ b/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime.go @@ -19,10 +19,11 @@ import ( ) const ( - maxWebSocketMessageBytes = 100 * 1024 * 1024 - webSocketPingInterval = 20 * time.Second - webSocketPingTimeout = 20 * time.Second - webSocketDrainTimeout = 60 * time.Second + maxWebSocketMessageBytes = 8 * 1024 * 1024 + webSocketHandshakeTimeout = 30 * time.Second + webSocketPingInterval = 20 * time.Second + webSocketPingTimeout = 20 * time.Second + webSocketDrainTimeout = 60 * time.Second ) type WebSocketRuntimeSession struct { @@ -172,6 +173,19 @@ func (c *WebSocketRuntimeSession) exchange( if err != nil { return "", err } + if len(request) > maxWebSocketMessageBytes { + return "", &azdext.LocalError{ + Message: fmt.Sprintf( + "OpenEnv WebSocket %s request is %d bytes; the RLE service limit is %d bytes.", + operation, + len(request), + maxWebSocketMessageBytes, + ), + Code: "rle_open_env_websocket_request_too_large", + Category: azdext.LocalErrorCategoryUser, + Suggestion: "Reduce the request payload and retry.", + } + } deadline, hasDeadline := operationDeadline(ctx, c.timeout) if !cancelAfterSend && !hasDeadline { @@ -268,8 +282,8 @@ func (c *WebSocketRuntimeSession) connect(ctx context.Context) error { } dialer := *websocket.DefaultDialer - dialer.HandshakeTimeout = 0 - if c.timeout > 0 { + dialer.HandshakeTimeout = webSocketHandshakeTimeout + if c.timeout > 0 && time.Duration(c.timeout)*time.Second < dialer.HandshakeTimeout { dialer.HandshakeTimeout = time.Duration(c.timeout) * time.Second } connection, response, err := dialer.DialContext(ctx, endpoint, headers) @@ -381,11 +395,24 @@ func parseWebSocketResponse(operation string, response []byte) (string, bool, er } } if envelope.Type == "error" { - return "", false, &azdext.LocalError{ - Message: fmt.Sprintf("Environment runtime rejected the %s request: %s", operation, prettyJson(envelope.Data)), + var errorData struct { + Code string `json:"code"` + } + _ = json.Unmarshal(envelope.Data, &errorData) + terminal := isTerminalOpenEnvError(errorData.Code) + suggestion := "Check the request payload and retry." + if terminal { + suggestion = "Exit and run invoke again to start a new environment session." + } + return "", terminal, &azdext.LocalError{ + Message: fmt.Sprintf( + "Environment runtime rejected the %s request: %s", + operation, + prettyJson(envelope.Data), + ), Code: "rle_open_env_websocket_request_failed", Category: azdext.LocalErrorCategoryUser, - Suggestion: "Check the request payload and retry.", + Suggestion: suggestion, } } expectedType := "observation" @@ -407,6 +434,15 @@ func parseWebSocketResponse(operation string, response []byte) (string, bool, er return prettyJson(envelope.Data), false, nil } +func isTerminalOpenEnvError(code string) bool { + switch code { + case "CAPACITY_REACHED", "FACTORY_ERROR", "SESSION_ERROR": + return true + default: + return false + } +} + func RuntimeWebSocketURL(baseURL string) (string, error) { endpoint, err := url.Parse(baseURL) if err != nil { diff --git a/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime_test.go b/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime_test.go index cd870bba061..cdc2d705f56 100644 --- a/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime_test.go +++ b/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime_test.go @@ -7,6 +7,7 @@ import ( "bytes" "context" "encoding/json" + "errors" "fmt" "net/http" "net/http/httptest" @@ -15,6 +16,7 @@ import ( "testing" "time" + "github.com/azure/azure-dev/cli/azd/pkg/azdext" "github.com/gorilla/websocket" ) @@ -222,6 +224,28 @@ func TestParseWebSocketResponse(t *testing.T) { t.Fatalf("unexpected OpenEnv error response: terminal=%t err=%v", terminal, err) } + for _, code := range []string{"CAPACITY_REACHED", "FACTORY_ERROR", "SESSION_ERROR"} { + _, terminal, err = parseWebSocketResponse( + "step", + []byte(fmt.Sprintf(`{"type":"error","data":{"code":%q,"detail":"failed"}}`, code)), + ) + var localError *azdext.LocalError + if err == nil || !terminal || !errors.As(err, &localError) || + !strings.Contains(localError.Suggestion, "run invoke again") { + t.Fatalf("expected %s to be terminal with reinvoke guidance: terminal=%t err=%v", code, terminal, err) + } + } + + _, terminal, err = parseWebSocketResponse( + "step", + []byte(`{"type":"error","data":{"code":"VALIDATION_ERROR","detail":"invalid action"}}`), + ) + var localError *azdext.LocalError + if err == nil || terminal || !errors.As(err, &localError) || + !strings.Contains(localError.Suggestion, "payload and retry") { + t.Fatalf("expected validation error to be recoverable: terminal=%t err=%v", terminal, err) + } + _, terminal, err = parseWebSocketResponse( "state", []byte(`{"type":"observation","data":{}}`), @@ -231,6 +255,42 @@ func TestParseWebSocketResponse(t *testing.T) { } } +func TestWebSocketSessionErrorPreventsReconnect(t *testing.T) { + upgrades := 0 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + upgrades++ + connection, err := (&websocket.Upgrader{}).Upgrade(w, r, nil) + if err != nil { + t.Error(err) + return + } + defer connection.Close() + if _, _, err := connection.ReadMessage(); err != nil { + t.Error(err) + return + } + if err := connection.WriteMessage( + websocket.TextMessage, + []byte(`{"type":"error","data":{"code":"SESSION_ERROR","detail":"session failed"}}`), + ); err != nil { + t.Error(err) + } + })) + defer server.Close() + + session := NewWebSocketRuntimeSession(server.URL, 30, nil) + defer session.Close() + if _, err := session.Call(t.Context(), "state", ""); err == nil { + t.Fatal("expected the session error to fail the first call") + } + if _, err := session.Call(t.Context(), "state", ""); err == nil { + t.Fatal("expected the terminal session error to fail the second call") + } + if upgrades != 1 { + t.Fatalf("expected no reconnection after session error, got %d connections", upgrades) + } +} + func TestWebSocketSessionCloseInterruptsCallAndPreventsReconnect(t *testing.T) { requestReceived := make(chan struct{}) upgrades := 0 From 5f9c60f4115f0477c9471daff233bc1106bc4ca0 Mon Sep 17 00:00:00 2001 From: Farhan Nawaz Date: Thu, 10 Sep 2026 22:59:43 +0530 Subject: [PATCH 6/8] fix: review comments --- .../azure.ai.rle/internal/project/scaffold.go | 8 ++------ .../azure.ai.rle/internal/project/scaffold_test.go | 5 ++--- .../internal/project/websocket_runtime_test.go | 10 +++++----- 3 files changed, 9 insertions(+), 14 deletions(-) diff --git a/cli/azd/extensions/azure.ai.rle/internal/project/scaffold.go b/cli/azd/extensions/azure.ai.rle/internal/project/scaffold.go index 17d7838a10b..163ff49648b 100644 --- a/cli/azd/extensions/azure.ai.rle/internal/project/scaffold.go +++ b/cli/azd/extensions/azure.ai.rle/internal/project/scaffold.go @@ -7,6 +7,7 @@ import ( "os" "os/exec" "path/filepath" + "slices" "strings" "github.com/azure/azure-dev/cli/azd/pkg/azdext" @@ -168,12 +169,7 @@ func editDistance(left string, right string) int { } func containsString(values []string, expected string) bool { - for _, value := range values { - if value == expected { - return true - } - } - return false + return slices.Contains(values, expected) } func runGitCheckout(args ...string) error { diff --git a/cli/azd/extensions/azure.ai.rle/internal/project/scaffold_test.go b/cli/azd/extensions/azure.ai.rle/internal/project/scaffold_test.go index 0915164193e..d95e0939578 100644 --- a/cli/azd/extensions/azure.ai.rle/internal/project/scaffold_test.go +++ b/cli/azd/extensions/azure.ai.rle/internal/project/scaffold_test.go @@ -122,9 +122,8 @@ func TestOpenEnvEnvironmentNotFoundSuggestsClosestCatalogName(t *testing.T) { "ches_env", []string{"atari_env", "chess_env", "echo_env"}, ) - var localError *azdext.LocalError - if !errors.As(err, &localError) || - !strings.Contains(localError.Suggestion, `Did you mean "chess_env"?`) { + localError, ok := errors.AsType[*azdext.LocalError](err) + if !ok || !strings.Contains(localError.Suggestion, `Did you mean "chess_env"?`) { t.Fatalf("expected closest catalog suggestion, got %v", err) } } diff --git a/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime_test.go b/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime_test.go index cdc2d705f56..0c0768048bf 100644 --- a/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime_test.go +++ b/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime_test.go @@ -227,10 +227,10 @@ func TestParseWebSocketResponse(t *testing.T) { for _, code := range []string{"CAPACITY_REACHED", "FACTORY_ERROR", "SESSION_ERROR"} { _, terminal, err = parseWebSocketResponse( "step", - []byte(fmt.Sprintf(`{"type":"error","data":{"code":%q,"detail":"failed"}}`, code)), + fmt.Appendf(nil, `{"type":"error","data":{"code":%q,"detail":"failed"}}`, code), ) - var localError *azdext.LocalError - if err == nil || !terminal || !errors.As(err, &localError) || + localError, ok := errors.AsType[*azdext.LocalError](err) + if err == nil || !terminal || !ok || !strings.Contains(localError.Suggestion, "run invoke again") { t.Fatalf("expected %s to be terminal with reinvoke guidance: terminal=%t err=%v", code, terminal, err) } @@ -240,8 +240,8 @@ func TestParseWebSocketResponse(t *testing.T) { "step", []byte(`{"type":"error","data":{"code":"VALIDATION_ERROR","detail":"invalid action"}}`), ) - var localError *azdext.LocalError - if err == nil || terminal || !errors.As(err, &localError) || + localError, ok := errors.AsType[*azdext.LocalError](err) + if err == nil || terminal || !ok || !strings.Contains(localError.Suggestion, "payload and retry") { t.Fatalf("expected validation error to be recoverable: terminal=%t err=%v", terminal, err) } From 8f9de8987bee58e6e59463a563100546541e2387 Mon Sep 17 00:00:00 2001 From: Farhan Nawaz Date: Fri, 11 Sep 2026 17:41:49 +0530 Subject: [PATCH 7/8] fix: init command changes --- cli/azd/extensions/azure.ai.rle/CHANGELOG.md | 3 +- cli/azd/extensions/azure.ai.rle/README.md | 24 +-- .../extensions/azure.ai.rle/extension.yaml | 6 +- .../azure.ai.rle/internal/cmd/init.go | 95 +++++++-- .../azure.ai.rle/internal/cmd/root_test.go | 141 +++++++------ .../azure.ai.rle/internal/project/scaffold.go | 191 ++++++++---------- .../internal/project/scaffold_test.go | 116 +++++++---- cli/azd/extensions/azure.ai.rle/version.txt | 2 +- 8 files changed, 333 insertions(+), 245 deletions(-) diff --git a/cli/azd/extensions/azure.ai.rle/CHANGELOG.md b/cli/azd/extensions/azure.ai.rle/CHANGELOG.md index d99903acb97..f50a9b523a9 100644 --- a/cli/azd/extensions/azure.ai.rle/CHANGELOG.md +++ b/cli/azd/extensions/azure.ai.rle/CHANGELOG.md @@ -1,6 +1,6 @@ # Release History -## 0.4.0-preview (Unreleased) +## 0.4.1-preview (Unreleased) - Align environment discovery and remote invocation with the refreshed RLE service routes and cursor-based response contracts. - Use `/rl_environments` consistently for environment and instance lifecycle APIs. @@ -10,6 +10,7 @@ - Delete the temporary instance and group on exit with Ctrl+C-independent cleanup and concise terminal status. - Persist the environment name as `environmentName` while continuing to read legacy `name` state files. - Authenticate and API-version OpenEnv gateway requests on the configured Foundry project origin, wait for runtime health before reporting readiness, and route the browser playground through an authenticated local proxy. +- Initialize a required local folder by interactively selecting and sparsely downloading an environment from the RLE samples repository. ## 0.3.0-preview diff --git a/cli/azd/extensions/azure.ai.rle/README.md b/cli/azd/extensions/azure.ai.rle/README.md index 53ed86a3fa1..a8dd7b88a5e 100644 --- a/cli/azd/extensions/azure.ai.rle/README.md +++ b/cli/azd/extensions/azure.ai.rle/README.md @@ -10,7 +10,7 @@ Install: - Azure CLI (`az`): https://learn.microsoft.com/cli/azure/install-azure-cli - Docker Desktop: https://www.docker.com/products/docker-desktop/ - Go, if building from source: https://go.dev/doc/install -- Git, if building from source: https://git-scm.com/downloads +- Git, required by `azd ai rle init` to download samples: https://git-scm.com/downloads Verify: @@ -81,26 +81,18 @@ az acr login --name ### 1. Initialize an environment session -Default echo session: +Choose the local folder name: ```powershell -azd ai rle init -cd .\echo_env +azd ai rle init my_environment ``` -The default echo session downloads the Hugging Face `OpenEnv` repo, copies `envs/echo_env` into the session folder, and writes `.azd-rle.json` with the local `environmentName`. Existing state files that use the legacy `name` property remain supported. +`init` reads the available environments from +[rle-samples](https://github.com/sujit-kamireddy/rle-samples) and prompts you to select one. +Only the selected sample is downloaded and copied into `.\my_environment`, and `.azd-rle.json` +stores `my_environment` as the RLE environment name. -The copied session does not keep `.git` metadata from the upstream repository. - -Copy another environment from the OpenEnv `envs` catalog: - -```powershell -azd ai rle init chess_env -cd .\chess_env -``` - -The positional name selects `envs/` from OpenEnv and is also used for the local session directory -and RLE environment name. If the name is not present in the catalog, `init` suggests the closest available name. +The copied session does not keep `.git` metadata from the sample repository. For an existing source folder, skip `init` and run commands directly from that folder. diff --git a/cli/azd/extensions/azure.ai.rle/extension.yaml b/cli/azd/extensions/azure.ai.rle/extension.yaml index eb99f6e9e20..c4224103f6f 100644 --- a/cli/azd/extensions/azure.ai.rle/extension.yaml +++ b/cli/azd/extensions/azure.ai.rle/extension.yaml @@ -11,11 +11,11 @@ tags: - ai - rle usage: azd ai rle [options] -version: 0.4.0-preview +version: 0.4.1-preview examples: - name: init - description: Copy an environment from the OpenEnv catalog; defaults to echo_env. - usage: azd ai rle init [environment-name] + description: Select an RLE sample and copy it into a new local folder. + usage: azd ai rle init - name: publish description: Build, push, and create or update the RLE environment. usage: azd ai rle publish --version-bump major diff --git a/cli/azd/extensions/azure.ai.rle/internal/cmd/init.go b/cli/azd/extensions/azure.ai.rle/internal/cmd/init.go index 81c45931f2a..fb8d2c8e9ce 100644 --- a/cli/azd/extensions/azure.ai.rle/internal/cmd/init.go +++ b/cli/azd/extensions/azure.ai.rle/internal/cmd/init.go @@ -4,6 +4,7 @@ package cmd import ( + "context" "fmt" "os" "runtime" @@ -20,34 +21,40 @@ type rleInitFlags struct { } type initAction struct { - cmd *cobra.Command - flags *rleInitFlags - environmentName string + cmd *cobra.Command + flags *rleInitFlags + folderName string } -var checkoutOpenEnvEnvironmentFunc = project.CheckoutOpenEnvEnvironment +type rleSampleCatalog interface { + SampleNames() []string + Copy(sampleName string, folderName string, dest string, force bool) (string, error) + Close() error +} + +var loadRleSampleCatalogFunc = func() (rleSampleCatalog, error) { + return project.LoadRleSampleCatalog() +} + +var selectRleSampleFunc = selectRleSample func newInitCommand() *cobra.Command { flags := &rleInitFlags{} cmd := &cobra.Command{ - Use: "init [environment-name]", - Short: "Initialize a local RLE environment", - Args: cobra.MaximumNArgs(1), + Use: "init ", + Short: "Initialize a local RLE environment from a sample", + Args: cobra.ExactArgs(1), RunE: func(cmd *cobra.Command, args []string) error { - environmentName := "" - if len(args) == 1 { - environmentName = args[0] - } - return (&initAction{cmd: cmd, flags: flags, environmentName: environmentName}).Run() + return (&initAction{cmd: cmd, flags: flags, folderName: args[0]}).Run() }, } cmd.SetHelpFunc(func(cmd *cobra.Command, args []string) { var help strings.Builder - help.WriteString("Initialize a local RLE environment\n") + help.WriteString("Initialize a local RLE environment from a sample\n") help.WriteString("Usage:\n") - help.WriteString(" rle init [environment-name] [flags]\n") + help.WriteString(" rle init [flags]\n") help.WriteString("Flags:\n") help.WriteString(" --force Overwrite generated files in an existing non-empty session directory\n") help.WriteString(" -h, --help help for init\n") @@ -62,9 +69,7 @@ func newInitCommand() *cobra.Command { } func (a *initAction) Run() error { - envName := firstNonEmpty(a.environmentName, "echo_env") - var err error - envName, err = project.ValidateEnvironmentName(envName) + folderName, err := project.ValidateEnvironmentName(a.folderName) if err != nil { return &azdext.LocalError{ Message: err.Error(), @@ -74,20 +79,70 @@ func (a *initAction) Run() error { } } - sessionDir, err := checkoutOpenEnvEnvironmentFunc(envName, ".", a.flags.force) + catalog, err := loadRleSampleCatalogFunc() + if err != nil { + return err + } + defer func() { + _ = catalog.Close() + }() + sampleName, err := selectRleSampleFunc(a.cmd.Context(), catalog.SampleNames()) + if err != nil { + return err + } + sessionDir, err := catalog.Copy(sampleName, folderName, ".", a.flags.force) if err != nil { return err } - if err := saveRleStateIn(sessionDir, defaultRleState(envName)); err != nil { + if err := saveRleStateIn(sessionDir, defaultRleState(folderName)); err != nil { return err } displayDir := "." + string(os.PathSeparator) + sessionDir + if _, err := fmt.Fprintf(a.cmd.OutOrStdout(), "Copied RLE sample %q.\n", sampleName); err != nil { + return err + } _, err = fmt.Fprint(a.cmd.OutOrStdout(), initNextSteps(displayDir, runtime.GOOS, os.Getenv("SHELL"))) return err } +func selectRleSample(ctx context.Context, sampleNames []string) (string, error) { + if len(sampleNames) == 0 { + return "", &azdext.LocalError{ + Message: "No RLE samples are available.", + Code: "rle_samples_empty", + Category: azdext.LocalErrorCategoryUser, + Suggestion: "Add a sample to the RLE samples repository, then retry.", + } + } + choices := make([]*azdext.SelectChoice, len(sampleNames)) + for index, sampleName := range sampleNames { + choices[index] = &azdext.SelectChoice{Label: sampleName, Value: sampleName} + } + azdClient, err := azdext.NewAzdClient() + if err != nil { + return "", fmt.Errorf("create azd client for sample selection: %w", err) + } + defer azdClient.Close() + response, err := azdClient.Prompt().Select(azdext.WithAccessToken(ctx), &azdext.SelectRequest{ + Options: &azdext.SelectOptions{ + Message: "Select an RLE sample", + Choices: choices, + DisplayNumbers: new(true), + EnableFiltering: new(true), + }, + }) + if err != nil { + return "", fmt.Errorf("select RLE sample: %w", err) + } + selectedIndex := int(response.GetValue()) + if selectedIndex < 0 || selectedIndex >= len(sampleNames) { + return "", fmt.Errorf("invalid RLE sample selection index: %d", selectedIndex) + } + return sampleNames[selectedIndex], nil +} + func initNextSteps(displayDir string, goos string, shell string) string { projectEndpoint := `https://.services.ai.azure.com/api/projects/` registryEndpoint := `.azurecr.io` @@ -111,7 +166,7 @@ func initNextSteps(displayDir string, goos string, shell string) string { } return fmt.Sprintf( - "Created OpenEnv-style environment at: %s\n"+ + "Created RLE environment at: %s\n"+ "\nRun locally:\n"+ " cd \"%s\"\n"+ " azd ai rle run\n"+ diff --git a/cli/azd/extensions/azure.ai.rle/internal/cmd/root_test.go b/cli/azd/extensions/azure.ai.rle/internal/cmd/root_test.go index 3cdf4ceab00..15d6b1f1a86 100644 --- a/cli/azd/extensions/azure.ai.rle/internal/cmd/root_test.go +++ b/cli/azd/extensions/azure.ai.rle/internal/cmd/root_test.go @@ -5,9 +5,11 @@ package cmd import ( "bytes" + "context" "encoding/json" "os" "path/filepath" + "slices" "strings" "testing" ) @@ -232,19 +234,23 @@ func TestLifecycleCommandsRejectPositionalArguments(t *testing.T) { t.Fatalf("expected init command to be registered: %v", err) } if err := initCommand.Args(initCommand, []string{"custom_env"}); err != nil { - t.Fatalf("expected init to accept one positional environment name: %v", err) + t.Fatalf("expected init to accept one positional folder name: %v", err) + } + if err := initCommand.Args(initCommand, nil); err == nil { + t.Fatal("expected init to require a folder name") } if err := initCommand.Args(initCommand, []string{"one", "two"}); err == nil { t.Fatal("expected init to reject multiple positional arguments") } } -func TestInitCopiesOpenEnvEchoSampleByDefault(t *testing.T) { +func TestInitSelectsSampleAndCopiesItToNamedFolder(t *testing.T) { tempDir := t.TempDir() t.Chdir(tempDir) - stubOpenEnvCheckout(t, "echo_env") + stubRleSampleCatalog(t, []string{"echo", "wordle"}, "wordle", "training_env") command := newInitCommand() + command.SetArgs([]string{"training_env"}) var output bytes.Buffer command.SetOut(&output) command.SetErr(&output) @@ -252,7 +258,7 @@ func TestInitCopiesOpenEnvEchoSampleByDefault(t *testing.T) { t.Fatal(err) } - sessionDir := filepath.Join(tempDir, "echo_env") + sessionDir := filepath.Join(tempDir, "training_env") // The test reads state from its own temporary session directory. stateBytes, err := os.ReadFile(filepath.Join(sessionDir, rleStateFile)) //nolint:gosec if err != nil { @@ -262,11 +268,11 @@ func TestInitCopiesOpenEnvEchoSampleByDefault(t *testing.T) { if err := json.Unmarshal(stateBytes, &state); err != nil { t.Fatal(err) } - if state.EnvironmentName != "echo_env" { - t.Fatalf("expected echo_env environment name, got %q", state.EnvironmentName) + if state.EnvironmentName != "training_env" { + t.Fatalf("expected training_env environment name, got %q", state.EnvironmentName) } if _, err := os.Stat(filepath.Join(sessionDir, "server", "Dockerfile")); err != nil { - t.Fatalf("expected copied OpenEnv server Dockerfile: %v", err) + t.Fatalf("expected copied RLE sample server Dockerfile: %v", err) } if _, err := os.Stat(filepath.Join(sessionDir, ".git")); !os.IsNotExist(err) { t.Fatalf("expected copied sample not to include .git metadata, got err=%v", err) @@ -274,38 +280,12 @@ func TestInitCopiesOpenEnvEchoSampleByDefault(t *testing.T) { if strings.Contains(output.String(), sessionDir) { t.Fatalf("expected init output not to use absolute cd path, got %s", output.String()) } - expectedCd := `cd "` + "." + string(os.PathSeparator) + "echo_env" + `"` + expectedCd := `cd "` + "." + string(os.PathSeparator) + "training_env" + `"` if !strings.Contains(output.String(), expectedCd) { t.Fatalf("expected init output to quote relative cd path, got %s", output.String()) } -} - -func TestInitUsesPositionalNameToSelectOpenEnvEnvironment(t *testing.T) { - tempDir := t.TempDir() - t.Chdir(tempDir) - stubOpenEnvCheckout(t, "chess_env") - - command := newInitCommand() - command.SetArgs([]string{"chess_env"}) - var output bytes.Buffer - command.SetOut(&output) - command.SetErr(&output) - if err := command.Execute(); err != nil { - t.Fatal(err) - } - - sessionDir := filepath.Join(tempDir, "chess_env") - // The test reads state from its own temporary session directory. - stateBytes, err := os.ReadFile(filepath.Join(sessionDir, rleStateFile)) //nolint:gosec - if err != nil { - t.Fatal(err) - } - var state rleState - if err := json.Unmarshal(stateBytes, &state); err != nil { - t.Fatal(err) - } - if state.EnvironmentName != "chess_env" { - t.Fatalf("expected chess_env environment name, got %q", state.EnvironmentName) + if !strings.Contains(output.String(), `Copied RLE sample "wordle".`) { + t.Fatalf("expected selected sample in output, got %s", output.String()) } } @@ -368,35 +348,74 @@ func TestInitNextStepsUseShellAppropriateSyntax(t *testing.T) { } } -func stubOpenEnvCheckout(t *testing.T, expectedName string) { - t.Helper() - old := checkoutOpenEnvEnvironmentFunc - checkoutOpenEnvEnvironmentFunc = func(name string, dest string, force bool) (string, error) { - if name != expectedName { - t.Fatalf("expected OpenEnv environment %q, got %q", expectedName, name) - } - sessionDir := filepath.Join(dest, name) - if force { - if err := os.RemoveAll(sessionDir); err != nil { - return "", err - } - } - if err := os.MkdirAll(sessionDir, 0750); err != nil { - return "", err - } - serverDir := filepath.Join(sessionDir, "server") - if err := os.MkdirAll(serverDir, 0750); err != nil { - return "", err - } - if err := os.WriteFile(filepath.Join(serverDir, "Dockerfile"), []byte("FROM scratch\n"), 0600); err != nil { +type testRleSampleCatalog struct { + t *testing.T + sampleNames []string + expectedSampleName string + expectedFolderName string +} + +func (c *testRleSampleCatalog) SampleNames() []string { + return c.sampleNames +} + +func (c *testRleSampleCatalog) Copy( + sampleName string, + folderName string, + dest string, + force bool, +) (string, error) { + c.t.Helper() + if sampleName != c.expectedSampleName { + c.t.Fatalf("expected RLE sample %q, got %q", c.expectedSampleName, sampleName) + } + if folderName != c.expectedFolderName { + c.t.Fatalf("expected folder name %q, got %q", c.expectedFolderName, folderName) + } + sessionDir := filepath.Join(dest, folderName) + if force { + if err := os.RemoveAll(sessionDir); err != nil { return "", err } - if err := os.WriteFile(filepath.Join(sessionDir, "openenv.yaml"), []byte("name: "+name+"\n"), 0600); err != nil { - return "", err + } + if err := os.MkdirAll(filepath.Join(sessionDir, "server"), 0750); err != nil { + return "", err + } + if err := os.WriteFile(filepath.Join(sessionDir, "server", "Dockerfile"), []byte("FROM scratch\n"), 0600); err != nil { + return "", err + } + return sessionDir, nil +} + +func (c *testRleSampleCatalog) Close() error { + return nil +} + +func stubRleSampleCatalog( + t *testing.T, + sampleNames []string, + selectedSampleName string, + expectedFolderName string, +) { + t.Helper() + oldLoad := loadRleSampleCatalogFunc + oldSelect := selectRleSampleFunc + loadRleSampleCatalogFunc = func() (rleSampleCatalog, error) { + return &testRleSampleCatalog{ + t: t, + sampleNames: sampleNames, + expectedSampleName: selectedSampleName, + expectedFolderName: expectedFolderName, + }, nil + } + selectRleSampleFunc = func(_ context.Context, actualSampleNames []string) (string, error) { + if !slices.Equal(actualSampleNames, sampleNames) { + t.Fatalf("expected sample names %v, got %v", sampleNames, actualSampleNames) } - return sessionDir, nil + return selectedSampleName, nil } t.Cleanup(func() { - checkoutOpenEnvEnvironmentFunc = old + loadRleSampleCatalogFunc = oldLoad + selectRleSampleFunc = oldSelect }) } diff --git a/cli/azd/extensions/azure.ai.rle/internal/project/scaffold.go b/cli/azd/extensions/azure.ai.rle/internal/project/scaffold.go index 163ff49648b..6cd5a6ee91a 100644 --- a/cli/azd/extensions/azure.ai.rle/internal/project/scaffold.go +++ b/cli/azd/extensions/azure.ai.rle/internal/project/scaffold.go @@ -14,10 +14,15 @@ import ( ) const ( - openEnvRepoUrl = "https://github.com/huggingface/OpenEnv.git" - openEnvRepoRef = "main" + rleSamplesRepoURL = "https://github.com/sujit-kamireddy/rle-samples.git" + rleSamplesRepoRef = "main" ) +type RleSampleCatalog struct { + repoDir string + sampleNames []string +} + func createRleSessionDir(name string, dest string, force bool) (string, error) { sessionDir := filepath.Join(dest, name) if entries, err := os.ReadDir(sessionDir); err == nil && len(entries) > 0 && !force { @@ -41,140 +46,115 @@ func createRleSessionDir(name string, dest string, force bool) (string, error) { return sessionDir, nil } -func CheckoutOpenEnvEnvironment(name string, dest string, force bool) (string, error) { - name, err := ValidateEnvironmentName(name) - if err != nil { - return "", err - } - sourcePath := openEnvEnvironmentPath(name) - tempDir, err := os.MkdirTemp("", "azd-rle-open-env-*") +func LoadRleSampleCatalog() (*RleSampleCatalog, error) { + return loadRleSampleCatalog(rleSamplesRepoURL, rleSamplesRepoRef) +} + +func loadRleSampleCatalog(repoURL string, repoRef string) (*RleSampleCatalog, error) { + tempDir, err := os.MkdirTemp("", "azd-rle-samples-*") if err != nil { - return "", err + return nil, err } - defer func() { - _ = os.RemoveAll(tempDir) - }() - - if err := runGitCheckout( + if _, err := runGitCommand( "clone", "--depth", "1", "--filter=blob:none", "--sparse", - "--branch", openEnvRepoRef, - openEnvRepoUrl, + "--branch", repoRef, + "--single-branch", + repoURL, tempDir, ); err != nil { - return "", err + _ = os.RemoveAll(tempDir) + return nil, err } - environmentNames, err := listOpenEnvEnvironments(tempDir) + + sampleNames, err := listRleSamples(tempDir, repoRef) if err != nil { - return "", err - } - if !containsString(environmentNames, name) { - return "", openEnvEnvironmentNotFoundError(name, environmentNames) + _ = os.RemoveAll(tempDir) + return nil, err } - if err := runGitCheckout("-C", tempDir, "sparse-checkout", "set", sourcePath); err != nil { - return "", err + if len(sampleNames) == 0 { + _ = os.RemoveAll(tempDir) + return nil, &azdext.LocalError{ + Message: "The RLE samples repository does not contain any sample environments.", + Code: "rle_samples_empty", + Category: azdext.LocalErrorCategoryUser, + Suggestion: fmt.Sprintf( + "Add sample directories to %s, then retry.", + strings.TrimSuffix(rleSamplesRepoURL, ".git"), + ), + } } + return &RleSampleCatalog{ + repoDir: tempDir, + sampleNames: sampleNames, + }, nil +} - sourceDir := filepath.Join(tempDir, filepath.FromSlash(sourcePath)) - return copyOpenEnvEnvironment(sourceDir, name, dest, force) +func (c *RleSampleCatalog) SampleNames() []string { + return slices.Clone(c.sampleNames) } -func copyOpenEnvEnvironment(sourceDir string, name string, dest string, force bool) (string, error) { - if _, err := os.Stat(sourceDir); os.IsNotExist(err) { - return "", openEnvEnvironmentNotFoundError(name, nil) - } else if err != nil { - return "", err - } - sessionDir, err := createRleSessionDir(name, dest, force) +func (c *RleSampleCatalog) Copy(sampleName string, folderName string, dest string, force bool) (string, error) { + folderName, err := ValidateEnvironmentName(folderName) if err != nil { return "", err } - if err := copyDirectory(sourceDir, sessionDir); err != nil { + sourcePath := filepath.ToSlash(filepath.Join("envs", sampleName)) + if _, err := runGitCommand("-C", c.repoDir, "sparse-checkout", "set", sourcePath); err != nil { return "", err } - return sessionDir, nil + sourceDir := filepath.Join(c.repoDir, filepath.FromSlash(sourcePath)) + return copyRleSample(sourceDir, folderName, dest, force) } -func openEnvEnvironmentPath(name string) string { - return "envs/" + name +func (c *RleSampleCatalog) Close() error { + return os.RemoveAll(c.repoDir) } -func listOpenEnvEnvironments(repoDir string) ([]string, error) { - output, err := runGitCommand("-C", repoDir, "ls-tree", "--name-only", openEnvRepoRef+":envs") +func listRleSamples(repoDir string, repoRef string) ([]string, error) { + output, err := runGitCommand( + "-C", + repoDir, + "ls-tree", + "-d", + "--name-only", + repoRef+":envs", + ) if err != nil { return nil, err } - return strings.Fields(string(output)), nil -} - -func openEnvEnvironmentNotFoundError(name string, environmentNames []string) error { - catalogURL := strings.TrimSuffix(openEnvRepoUrl, ".git") - suggestion := fmt.Sprintf( - "Choose an environment from %s/tree/%s/envs.", - catalogURL, - openEnvRepoRef, - ) - if closest := closestEnvironmentName(name, environmentNames); closest != "" { - suggestion = fmt.Sprintf("Did you mean %q? %s", closest, suggestion) - } - return &azdext.LocalError{ - Message: fmt.Sprintf("OpenEnv environment %q was not found.", name), - Code: "rle_open_env_environment_not_found", - Category: azdext.LocalErrorCategoryUser, - Suggestion: suggestion, - } + sampleNames := strings.Fields(string(output)) + sampleNames = slices.DeleteFunc(sampleNames, func(name string) bool { + return strings.HasPrefix(name, ".") + }) + slices.Sort(sampleNames) + return sampleNames, nil } -func closestEnvironmentName(name string, environmentNames []string) string { - closest := "" - closestDistance := len(name) + 1 - for _, candidate := range environmentNames { - distance := editDistance(name, candidate) - if distance < closestDistance { - closest = candidate - closestDistance = distance +func copyRleSample(sourceDir string, folderName string, dest string, force bool) (string, error) { + sourceInfo, err := os.Stat(sourceDir) + if os.IsNotExist(err) { + return "", &azdext.LocalError{ + Message: fmt.Sprintf("RLE sample source %q was not found.", sourceDir), + Code: "rle_sample_source_not_found", + Category: azdext.LocalErrorCategoryInternal, + Suggestion: "Run azd ai rle init again to refresh the sample list.", } + } else if err != nil { + return "", err + } else if !sourceInfo.IsDir() { + return "", fmt.Errorf("RLE sample source %q is not a directory", sourceDir) } - maxDistance := max(2, len(name)/3) - if closestDistance > maxDistance { - return "" + sessionDir, err := createRleSessionDir(folderName, dest, force) + if err != nil { + return "", err } - return closest -} - -func editDistance(left string, right string) int { - previous := make([]int, len(right)+1) - for index := range previous { - previous[index] = index - } - for leftIndex, leftRune := range left { - current := make([]int, len(right)+1) - current[0] = leftIndex + 1 - for rightIndex, rightRune := range right { - substitutionCost := 0 - if leftRune != rightRune { - substitutionCost = 1 - } - current[rightIndex+1] = min( - current[rightIndex]+1, - previous[rightIndex+1]+1, - previous[rightIndex]+substitutionCost, - ) - } - previous = current + if err := copyDirectory(sourceDir, sessionDir); err != nil { + return "", err } - return previous[len(right)] -} - -func containsString(values []string, expected string) bool { - return slices.Contains(values, expected) -} - -func runGitCheckout(args ...string) error { - _, err := runGitCommand(args...) - return err + return sessionDir, nil } func runGitCommand(args ...string) ([]byte, error) { @@ -186,14 +166,13 @@ func runGitCommand(args ...string) ([]byte, error) { Suggestion: "Install Git, then retry azd ai rle init.", } } - // The user-provided sparse path is validated as an environment name before reaching this command. process := exec.Command("git", args...) //nolint:gosec process.Env = os.Environ() output, err := process.CombinedOutput() if err != nil { return nil, &azdext.LocalError{ - Message: fmt.Sprintf("Failed to checkout OpenEnv environment: %v", err), - Code: "rle_open_env_checkout_failed", + Message: fmt.Sprintf("Failed to download RLE samples: %v", err), + Code: "rle_samples_download_failed", Category: azdext.LocalErrorCategoryUser, Suggestion: strings.TrimSpace(string(output)), } diff --git a/cli/azd/extensions/azure.ai.rle/internal/project/scaffold_test.go b/cli/azd/extensions/azure.ai.rle/internal/project/scaffold_test.go index d95e0939578..da63f32549c 100644 --- a/cli/azd/extensions/azure.ai.rle/internal/project/scaffold_test.go +++ b/cli/azd/extensions/azure.ai.rle/internal/project/scaffold_test.go @@ -4,14 +4,12 @@ package project import ( - "errors" "os" + "os/exec" "path/filepath" "runtime" - "strings" + "slices" "testing" - - "github.com/azure/azure-dev/cli/azd/pkg/azdext" ) func TestCopyDirectorySkipsGitMetadata(t *testing.T) { @@ -75,31 +73,92 @@ func TestCopyDirectoryRejectsFileSource(t *testing.T) { } } -func TestCheckoutOpenEnvEnvironmentRejectsInvalidNameBeforeChangingDestination(t *testing.T) { - destDir := t.TempDir() - sentinel := filepath.Join(destDir, "sentinel.txt") - if err := os.WriteFile(sentinel, []byte("keep"), 0600); err != nil { +func TestRleSampleCatalogUsesSparseCheckout(t *testing.T) { + if _, err := exec.LookPath("git"); err != nil { + t.Skip("git is not available") + } + + sourceRepo := t.TempDir() + runTestGit(t, sourceRepo, "init", "--initial-branch=main") + for _, sampleName := range []string{"code_rl", "math_rl"} { + sampleDir := filepath.Join(sourceRepo, "envs", sampleName) + if err := os.MkdirAll(sampleDir, 0750); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(sampleDir, "sample.txt"), []byte(sampleName), 0600); err != nil { + t.Fatal(err) + } + } + if err := os.WriteFile(filepath.Join(sourceRepo, "README.md"), []byte("samples"), 0600); err != nil { t.Fatal(err) } + runTestGit(t, sourceRepo, "add", ".") + runTestGit( + t, + sourceRepo, + "-c", "user.name=RLE Tests", + "-c", "user.email=rle-tests@example.com", + "commit", "-m", "Add samples", + ) - if _, err := CheckoutOpenEnvEnvironment("../bad", destDir, true); err == nil { - t.Fatal("expected invalid environment name to be rejected") + catalog, err := loadRleSampleCatalog(sourceRepo, "main") + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if err := catalog.Close(); err != nil { + t.Errorf("close sample catalog: %v", err) + } + }) + if !slices.Equal(catalog.SampleNames(), []string{"code_rl", "math_rl"}) { + t.Fatalf("expected sorted sample names, got %v", catalog.SampleNames()) + } + if _, err := os.Stat(filepath.Join(catalog.repoDir, "envs")); !os.IsNotExist(err) { + t.Fatalf("expected sample contents not to be checked out before selection, got err=%v", err) } - if _, err := os.Stat(sentinel); err != nil { - t.Fatalf("expected destination to be unchanged: %v", err) + sessionDir, err := catalog.Copy("math_rl", "training_env", t.TempDir(), false) + if err != nil { + t.Fatal(err) + } + if _, err := os.Stat(filepath.Join(sessionDir, "sample.txt")); err != nil { + t.Fatalf("expected selected sample to be copied: %v", err) + } + if _, err := os.Stat(filepath.Join(catalog.repoDir, "envs", "code_rl")); !os.IsNotExist(err) { + t.Fatalf("expected unselected sample not to be checked out, got err=%v", err) } } -func TestOpenEnvEnvironmentPath(t *testing.T) { - if got := openEnvEnvironmentPath("chess_env"); got != "envs/chess_env" { - t.Fatalf("expected selected OpenEnv environment path, got %q", got) +func TestCopyRleSampleRenamesDestination(t *testing.T) { + sourceDir := t.TempDir() + if err := os.WriteFile(filepath.Join(sourceDir, "sample.txt"), []byte("content"), 0600); err != nil { + t.Fatal(err) + } + destDir := t.TempDir() + sessionDir, err := copyRleSample(sourceDir, "my_environment", destDir, false) + if err != nil { + t.Fatal(err) + } + if sessionDir != filepath.Join(destDir, "my_environment") { + t.Fatalf("expected renamed destination, got %q", sessionDir) + } + if _, err := os.Stat(filepath.Join(sessionDir, "sample.txt")); err != nil { + t.Fatalf("expected sample file in renamed destination: %v", err) + } +} + +func runTestGit(t *testing.T, dir string, args ...string) { + t.Helper() + command := exec.Command("git", args...) //nolint:gosec + command.Dir = dir + if output, err := command.CombinedOutput(); err != nil { + t.Fatalf("git %v failed: %v\n%s", args, err, output) } } -func TestCopyOpenEnvEnvironmentValidatesSourceBeforeReplacingDestination(t *testing.T) { +func TestCopyRleSampleValidatesSourceBeforeReplacingDestination(t *testing.T) { destDir := t.TempDir() - sessionDir := filepath.Join(destDir, "missing_env") + sessionDir := filepath.Join(destDir, "my_environment") if err := os.MkdirAll(sessionDir, 0750); err != nil { t.Fatal(err) } @@ -108,28 +167,11 @@ func TestCopyOpenEnvEnvironmentValidatesSourceBeforeReplacingDestination(t *test t.Fatal(err) } - _, err := copyOpenEnvEnvironment(filepath.Join(t.TempDir(), "missing"), "missing_env", destDir, true) + _, err := copyRleSample(filepath.Join(t.TempDir(), "missing"), "my_environment", destDir, true) if err == nil { - t.Fatal("expected missing OpenEnv environment to fail") + t.Fatal("expected missing RLE sample to fail") } if _, statErr := os.Stat(sentinel); statErr != nil { - t.Fatalf("expected destination to remain unchanged after catalog lookup failure: %v", statErr) - } -} - -func TestOpenEnvEnvironmentNotFoundSuggestsClosestCatalogName(t *testing.T) { - err := openEnvEnvironmentNotFoundError( - "ches_env", - []string{"atari_env", "chess_env", "echo_env"}, - ) - localError, ok := errors.AsType[*azdext.LocalError](err) - if !ok || !strings.Contains(localError.Suggestion, `Did you mean "chess_env"?`) { - t.Fatalf("expected closest catalog suggestion, got %v", err) - } -} - -func TestClosestEnvironmentNameRejectsDistantMatch(t *testing.T) { - if got := closestEnvironmentName("unknown_env", []string{"chess_env", "echo_env"}); got != "" { - t.Fatalf("expected no distant suggestion, got %q", got) + t.Fatalf("expected destination to remain unchanged after sample lookup failure: %v", statErr) } } diff --git a/cli/azd/extensions/azure.ai.rle/version.txt b/cli/azd/extensions/azure.ai.rle/version.txt index 6bcee0a14ab..34dab3faf78 100644 --- a/cli/azd/extensions/azure.ai.rle/version.txt +++ b/cli/azd/extensions/azure.ai.rle/version.txt @@ -1 +1 @@ -0.4.0-preview \ No newline at end of file +0.4.1-preview \ No newline at end of file From 3e8938f7e060feceb5ad510059bd91349a7e3733 Mon Sep 17 00:00:00 2001 From: Farhan Nawaz Date: Fri, 11 Sep 2026 20:29:15 +0530 Subject: [PATCH 8/8] fix: websocket handshake retries and init command updates --- cli/azd/extensions/azure.ai.rle/README.md | 20 +- .../extensions/azure.ai.rle/extension.yaml | 2 +- .../azure.ai.rle/internal/cmd/init.go | 78 +++++- .../azure.ai.rle/internal/cmd/root.go | 2 +- .../azure.ai.rle/internal/cmd/root_test.go | 77 +++++- .../internal/project/websocket_runtime.go | 115 ++++++--- .../project/websocket_runtime_test.go | 241 +++++++++++++++++- 7 files changed, 470 insertions(+), 65 deletions(-) diff --git a/cli/azd/extensions/azure.ai.rle/README.md b/cli/azd/extensions/azure.ai.rle/README.md index a8dd7b88a5e..05faf29ecde 100644 --- a/cli/azd/extensions/azure.ai.rle/README.md +++ b/cli/azd/extensions/azure.ai.rle/README.md @@ -81,16 +81,28 @@ az acr login --name ### 1. Initialize an environment session -Choose the local folder name: +Select a sample and use its name for the local folder and RLE environment: ```powershell -azd ai rle init my_environment +azd ai rle init ``` `init` reads the available environments from [rle-samples](https://github.com/sujit-kamireddy/rle-samples) and prompts you to select one. -Only the selected sample is downloaded and copied into `.\my_environment`, and `.azd-rle.json` -stores `my_environment` as the RLE environment name. +Only the selected sample is downloaded. For example, selecting `code_rl` copies it into `.\code_rl` +and stores `code_rl` as the RLE environment name in `.azd-rle.json`. + +To use a different local folder and RLE environment name, provide it before selecting a sample: + +```powershell +azd ai rle init my_environment +``` + +When prompts are disabled, the required positional name selects the sample and is also used for the folder: + +```powershell +azd ai rle init code_rl --no-prompt +``` The copied session does not keep `.git` metadata from the sample repository. diff --git a/cli/azd/extensions/azure.ai.rle/extension.yaml b/cli/azd/extensions/azure.ai.rle/extension.yaml index c4224103f6f..c19bea8bde0 100644 --- a/cli/azd/extensions/azure.ai.rle/extension.yaml +++ b/cli/azd/extensions/azure.ai.rle/extension.yaml @@ -15,7 +15,7 @@ version: 0.4.1-preview examples: - name: init description: Select an RLE sample and copy it into a new local folder. - usage: azd ai rle init + usage: azd ai rle init [folder-name] - name: publish description: Build, push, and create or update the RLE environment. usage: azd ai rle publish --version-bump major diff --git a/cli/azd/extensions/azure.ai.rle/internal/cmd/init.go b/cli/azd/extensions/azure.ai.rle/internal/cmd/init.go index fb8d2c8e9ce..61719371e9b 100644 --- a/cli/azd/extensions/azure.ai.rle/internal/cmd/init.go +++ b/cli/azd/extensions/azure.ai.rle/internal/cmd/init.go @@ -8,6 +8,7 @@ import ( "fmt" "os" "runtime" + "slices" "strings" "azure.ai.rle/internal/project" @@ -24,6 +25,7 @@ type initAction struct { cmd *cobra.Command flags *rleInitFlags folderName string + noPrompt bool } type rleSampleCatalog interface { @@ -38,15 +40,24 @@ var loadRleSampleCatalogFunc = func() (rleSampleCatalog, error) { var selectRleSampleFunc = selectRleSample -func newInitCommand() *cobra.Command { +func newInitCommand(noPrompt *bool) *cobra.Command { flags := &rleInitFlags{} cmd := &cobra.Command{ - Use: "init ", + Use: "init [folder-name]", Short: "Initialize a local RLE environment from a sample", - Args: cobra.ExactArgs(1), + Args: cobra.MaximumNArgs(1), RunE: func(cmd *cobra.Command, args []string) error { - return (&initAction{cmd: cmd, flags: flags, folderName: args[0]}).Run() + folderName := "" + if len(args) == 1 { + folderName = args[0] + } + return (&initAction{ + cmd: cmd, + flags: flags, + folderName: folderName, + noPrompt: noPrompt != nil && *noPrompt, + }).Run() }, } @@ -54,7 +65,7 @@ func newInitCommand() *cobra.Command { var help strings.Builder help.WriteString("Initialize a local RLE environment from a sample\n") help.WriteString("Usage:\n") - help.WriteString(" rle init [flags]\n") + help.WriteString(" rle init [folder-name] [flags]\n") help.WriteString("Flags:\n") help.WriteString(" --force Overwrite generated files in an existing non-empty session directory\n") help.WriteString(" -h, --help help for init\n") @@ -69,13 +80,20 @@ func newInitCommand() *cobra.Command { } func (a *initAction) Run() error { - folderName, err := project.ValidateEnvironmentName(a.folderName) - if err != nil { + if a.noPrompt && a.folderName == "" { return &azdext.LocalError{ - Message: err.Error(), - Code: "rle_invalid_environment_name", + Message: "A sample name is required when prompts are disabled.", + Code: "rle_sample_name_required", Category: azdext.LocalErrorCategoryUser, - Suggestion: "Use snake_case starting with a letter, for example code_rl.", + Suggestion: "Run azd ai rle init --no-prompt.", + } + } + folderName := a.folderName + if folderName != "" { + var err error + folderName, err = validateRleFolderName(folderName) + if err != nil { + return err } } @@ -86,10 +104,20 @@ func (a *initAction) Run() error { defer func() { _ = catalog.Close() }() - sampleName, err := selectRleSampleFunc(a.cmd.Context(), catalog.SampleNames()) + requestedSample := "" + if a.noPrompt { + requestedSample = folderName + } + sampleName, err := resolveRleSample(a.cmd.Context(), requestedSample, catalog.SampleNames()) if err != nil { return err } + if folderName == "" { + folderName, err = validateRleFolderName(sampleName) + if err != nil { + return err + } + } sessionDir, err := catalog.Copy(sampleName, folderName, ".", a.flags.force) if err != nil { return err @@ -107,6 +135,34 @@ func (a *initAction) Run() error { return err } +func validateRleFolderName(folderName string) (string, error) { + validated, err := project.ValidateEnvironmentName(folderName) + if err == nil { + return validated, nil + } + return "", &azdext.LocalError{ + Message: err.Error(), + Code: "rle_invalid_environment_name", + Category: azdext.LocalErrorCategoryUser, + Suggestion: "Use snake_case starting with a letter, for example code_rl.", + } +} + +func resolveRleSample(ctx context.Context, requestedSample string, sampleNames []string) (string, error) { + if requestedSample == "" { + return selectRleSampleFunc(ctx, sampleNames) + } + if slices.Contains(sampleNames, requestedSample) { + return requestedSample, nil + } + return "", &azdext.LocalError{ + Message: fmt.Sprintf("RLE sample %q was not found.", requestedSample), + Code: "rle_sample_not_found", + Category: azdext.LocalErrorCategoryUser, + Suggestion: fmt.Sprintf("Choose one of the available samples: %s.", strings.Join(sampleNames, ", ")), + } +} + func selectRleSample(ctx context.Context, sampleNames []string) (string, error) { if len(sampleNames) == 0 { return "", &azdext.LocalError{ diff --git a/cli/azd/extensions/azure.ai.rle/internal/cmd/root.go b/cli/azd/extensions/azure.ai.rle/internal/cmd/root.go index a650ca91f66..b17411e7bb7 100644 --- a/cli/azd/extensions/azure.ai.rle/internal/cmd/root.go +++ b/cli/azd/extensions/azure.ai.rle/internal/cmd/root.go @@ -38,7 +38,7 @@ func NewRootCommand() *cobra.Command { userCommands := []*cobra.Command{ newListCommand(&extCtx.OutputFormat), newShowCommand(&extCtx.OutputFormat), - newInitCommand(), + newInitCommand(&extCtx.NoPrompt), newInvokeCommand(), newPublishCommand(), newRunCommand(), diff --git a/cli/azd/extensions/azure.ai.rle/internal/cmd/root_test.go b/cli/azd/extensions/azure.ai.rle/internal/cmd/root_test.go index 15d6b1f1a86..6d2420a86dd 100644 --- a/cli/azd/extensions/azure.ai.rle/internal/cmd/root_test.go +++ b/cli/azd/extensions/azure.ai.rle/internal/cmd/root_test.go @@ -7,11 +7,14 @@ import ( "bytes" "context" "encoding/json" + "errors" "os" "path/filepath" "slices" "strings" "testing" + + "github.com/azure/azure-dev/cli/azd/pkg/azdext" ) func TestNewRootCommandIncludesExpectedCommands(t *testing.T) { @@ -236,8 +239,8 @@ func TestLifecycleCommandsRejectPositionalArguments(t *testing.T) { if err := initCommand.Args(initCommand, []string{"custom_env"}); err != nil { t.Fatalf("expected init to accept one positional folder name: %v", err) } - if err := initCommand.Args(initCommand, nil); err == nil { - t.Fatal("expected init to require a folder name") + if err := initCommand.Args(initCommand, nil); err != nil { + t.Fatalf("expected interactive init to allow an omitted folder name: %v", err) } if err := initCommand.Args(initCommand, []string{"one", "two"}); err == nil { t.Fatal("expected init to reject multiple positional arguments") @@ -249,7 +252,8 @@ func TestInitSelectsSampleAndCopiesItToNamedFolder(t *testing.T) { t.Chdir(tempDir) stubRleSampleCatalog(t, []string{"echo", "wordle"}, "wordle", "training_env") - command := newInitCommand() + noPrompt := false + command := newInitCommand(&noPrompt) command.SetArgs([]string{"training_env"}) var output bytes.Buffer command.SetOut(&output) @@ -289,6 +293,73 @@ func TestInitSelectsSampleAndCopiesItToNamedFolder(t *testing.T) { } } +func TestInitWithoutFolderUsesSelectedSampleName(t *testing.T) { + tempDir := t.TempDir() + t.Chdir(tempDir) + stubRleSampleCatalog(t, []string{"echo", "wordle"}, "wordle", "wordle") + + noPrompt := false + command := newInitCommand(&noPrompt) + command.SetArgs(nil) + if err := command.Execute(); err != nil { + t.Fatal(err) + } + if _, err := os.Stat(filepath.Join(tempDir, "wordle", rleStateFile)); err != nil { + t.Fatalf("expected selected sample folder and state: %v", err) + } +} + +func TestInitNoPromptUsesPositionalNameAsSampleAndFolder(t *testing.T) { + tempDir := t.TempDir() + t.Chdir(tempDir) + t.Setenv(rleEnableEnvVar, "true") + + oldLoad := loadRleSampleCatalogFunc + oldSelect := selectRleSampleFunc + loadRleSampleCatalogFunc = func() (rleSampleCatalog, error) { + return &testRleSampleCatalog{ + t: t, + sampleNames: []string{"echo", "wordle"}, + expectedSampleName: "wordle", + expectedFolderName: "wordle", + }, nil + } + selectRleSampleFunc = func(context.Context, []string) (string, error) { + t.Fatal("expected no-prompt init to bypass the prompt") + return "", nil + } + t.Cleanup(func() { + loadRleSampleCatalogFunc = oldLoad + selectRleSampleFunc = oldSelect + }) + + command := NewRootCommand() + command.SetArgs([]string{"init", "wordle", "--no-prompt"}) + if err := command.Execute(); err != nil { + t.Fatal(err) + } +} + +func TestInitNoPromptRequiresPositionalSampleName(t *testing.T) { + t.Setenv(rleEnableEnvVar, "true") + command := NewRootCommand() + command.SetArgs([]string{"init", "--no-prompt"}) + err := command.Execute() + localError, ok := errors.AsType[*azdext.LocalError](err) + if !ok || localError.Code != "rle_sample_name_required" { + t.Fatalf("expected missing no-prompt sample error, got %v", err) + } +} + +func TestResolveRleSampleRejectsUnknownNonInteractiveSample(t *testing.T) { + _, err := resolveRleSample(t.Context(), "missing", []string{"echo", "wordle"}) + localError, ok := errors.AsType[*azdext.LocalError](err) + if !ok || localError.Code != "rle_sample_not_found" || + !strings.Contains(localError.Suggestion, "echo, wordle") { + t.Fatalf("expected available sample guidance, got %v", err) + } +} + func TestInitNextStepsUseShellAppropriateSyntax(t *testing.T) { tests := []struct { name string diff --git a/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime.go b/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime.go index 09697381c04..f69972bd384 100644 --- a/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime.go +++ b/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime.go @@ -6,8 +6,10 @@ package project import ( "context" "encoding/json" + "errors" "fmt" "io" + "net" "net/http" "net/url" "strings" @@ -26,6 +28,8 @@ const ( webSocketDrainTimeout = 60 * time.Second ) +var defaultWebSocketHandshakeRetryDelays = []time.Duration{time.Second, 2 * time.Second} + type WebSocketRuntimeSession struct { baseURL string timeout int @@ -35,6 +39,7 @@ type WebSocketRuntimeSession struct { connection *websocket.Conn connectionDone chan struct{} keepAliveInterval time.Duration + handshakeRetryDelays []time.Duration drainTimeout time.Duration terminalError error closed bool @@ -50,6 +55,7 @@ func NewWebSocketRuntimeSession( timeout: timeout, authorizationProvider: authorizationProvider, keepAliveInterval: webSocketPingInterval, + handshakeRetryDelays: defaultWebSocketHandshakeRetryDelays, drainTimeout: webSocketDrainTimeout, } } @@ -152,10 +158,17 @@ func (c *WebSocketRuntimeSession) exchange( ) (string, error) { c.exchangeMu.Lock() defer c.exchangeMu.Unlock() - if err := ctx.Err(); err != nil { + deadline, hasDeadline := operationDeadline(ctx, c.timeout) + operationCtx := ctx + cancel := func() {} + if hasDeadline { + operationCtx, cancel = context.WithDeadline(ctx, deadline) + } + defer cancel() + if err := operationCtx.Err(); err != nil { return "", err } - if err := c.connect(ctx); err != nil { + if err := c.connect(operationCtx); err != nil { return "", err } c.mu.Lock() @@ -187,11 +200,6 @@ func (c *WebSocketRuntimeSession) exchange( } } - deadline, hasDeadline := operationDeadline(ctx, c.timeout) - if !cancelAfterSend && !hasDeadline { - deadline = time.Now().Add(c.drainTimeout) - hasDeadline = true - } if hasDeadline { if err := connection.SetWriteDeadline(deadline); err != nil { return "", c.failConnection(connection, fmt.Errorf("set OpenEnv WebSocket write deadline: %w", err)) @@ -210,28 +218,30 @@ func (c *WebSocketRuntimeSession) exchange( fmt.Errorf("send OpenEnv WebSocket %s request: %w", operation, err), ) } - var exchangeDone chan struct{} - var cancellationHandled chan struct{} - if cancelAfterSend { - exchangeDone = make(chan struct{}) - cancellationHandled = make(chan struct{}) - go func() { - select { - case <-exchangeDone: - case <-ctx.Done(): + exchangeDone := make(chan struct{}) + cancellationHandled := make(chan struct{}) + go func() { + select { + case <-exchangeDone: + case <-ctx.Done(): + if cancelAfterSend { _ = c.failConnection( connection, fmt.Errorf("OpenEnv WebSocket %s request canceled: %w", operation, ctx.Err()), ) + } else { + drainDeadline := time.Now().Add(c.drainTimeout) + if hasDeadline && deadline.Before(drainDeadline) { + drainDeadline = deadline + } + _ = connection.SetReadDeadline(drainDeadline) } - close(cancellationHandled) - }() - } + } + close(cancellationHandled) + }() messageType, response, err := connection.ReadMessage() - if cancelAfterSend { - close(exchangeDone) - <-cancellationHandled - } + close(exchangeDone) + <-cancellationHandled if err != nil { return "", c.failConnection( connection, @@ -266,6 +276,13 @@ func (c *WebSocketRuntimeSession) connect(ctx context.Context) error { return nil } + connectCtx := ctx + cancel := func() {} + if deadline, hasDeadline := operationDeadline(ctx, c.timeout); hasDeadline { + connectCtx, cancel = context.WithDeadline(ctx, deadline) + } + defer cancel() + endpoint, err := RuntimeWebSocketURL(c.baseURL) if err != nil { return err @@ -286,14 +303,24 @@ func (c *WebSocketRuntimeSession) connect(ctx context.Context) error { if c.timeout > 0 && time.Duration(c.timeout)*time.Second < dialer.HandshakeTimeout { dialer.HandshakeTimeout = time.Duration(c.timeout) * time.Second } - connection, response, err := dialer.DialContext(ctx, endpoint, headers) - if err != nil { + + for attempt := 0; ; attempt++ { + connection, response, err := dialer.DialContext(connectCtx, endpoint, headers) + if err == nil { + connection.SetReadLimit(maxWebSocketMessageBytes) + c.connection = connection + c.connectionDone = make(chan struct{}) + go c.keepAlive(connection, c.connectionDone) + return nil + } detail := "" + retryable := isRetryableWebSocketHandshakeError(err) if response != nil { + retryable = isRetryableWebSocketHandshakeStatus(response.StatusCode) detail = readHealthErrorDetail(response.Body) _ = response.Body.Close() } - return &azdext.LocalError{ + connectionError := &azdext.LocalError{ Message: fmt.Sprintf( "Environment runtime WebSocket connection failed%s: %v", detail, @@ -303,12 +330,38 @@ func (c *WebSocketRuntimeSession) connect(ctx context.Context) error { Category: azdext.LocalErrorCategoryUser, Suggestion: "Check the remote RLE instance status and retry invoke.", } + if !retryable || attempt >= len(c.handshakeRetryDelays) { + return connectionError + } + timer := time.NewTimer(c.handshakeRetryDelays[attempt]) + select { + case <-timer.C: + case <-connectCtx.Done(): + if !timer.Stop() { + <-timer.C + } + return connectCtx.Err() + } + } +} + +func isRetryableWebSocketHandshakeError(err error) bool { + _, ok := errors.AsType[net.Error](err) + return ok +} + +func isRetryableWebSocketHandshakeStatus(statusCode int) bool { + switch statusCode { + case http.StatusRequestTimeout, + http.StatusTooManyRequests, + http.StatusInternalServerError, + http.StatusBadGateway, + http.StatusServiceUnavailable, + http.StatusGatewayTimeout: + return true + default: + return false } - connection.SetReadLimit(maxWebSocketMessageBytes) - c.connection = connection - c.connectionDone = make(chan struct{}) - go c.keepAlive(connection, c.connectionDone) - return nil } func (c *WebSocketRuntimeSession) keepAlive(connection *websocket.Conn, done <-chan struct{}) { diff --git a/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime_test.go b/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime_test.go index 0c0768048bf..a850dae4341 100644 --- a/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime_test.go +++ b/cli/azd/extensions/azure.ai.rle/internal/project/websocket_runtime_test.go @@ -9,10 +9,13 @@ import ( "encoding/json" "errors" "fmt" + "net" "net/http" "net/http/httptest" "slices" "strings" + "sync" + "sync/atomic" "testing" "time" @@ -24,6 +27,7 @@ func TestRunWebSocketShellUsesPersistentSocketForStatefulOperations(t *testing.T var requests []map[string]any var safePaths []string upgrades := 0 + var captureMu sync.Mutex upgrader := websocket.Upgrader{} server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Query().Get("api-version") != "test-version" { @@ -33,12 +37,16 @@ func TestRunWebSocketShellUsesPersistentSocketForStatefulOperations(t *testing.T t.Errorf("unexpected authorization header %q", r.Header.Get("Authorization")) } if r.URL.Path != "/ws" { + captureMu.Lock() safePaths = append(safePaths, r.URL.Path) + captureMu.Unlock() _, _ = fmt.Fprintf(w, `{"path":%q}`, r.URL.Path) //nolint:gosec // Test response uses JSON encoding. return } + captureMu.Lock() upgrades++ + captureMu.Unlock() connection, err := upgrader.Upgrade(w, r, nil) if err != nil { t.Errorf("upgrade WebSocket: %v", err) @@ -50,7 +58,9 @@ func TestRunWebSocketShellUsesPersistentSocketForStatefulOperations(t *testing.T if err := connection.ReadJSON(&request); err != nil { return } + captureMu.Lock() requests = append(requests, request) + captureMu.Unlock() responseType := "observation" if request["type"] == "state" { responseType = "state" @@ -87,17 +97,22 @@ func TestRunWebSocketShellUsesPersistentSocketForStatefulOperations(t *testing.T t.Fatal(err) } - if upgrades != 1 { - t.Fatalf("expected one persistent WebSocket, got %d", upgrades) + captureMu.Lock() + capturedUpgrades := upgrades + capturedRequests := slices.Clone(requests) + capturedSafePaths := slices.Clone(safePaths) + captureMu.Unlock() + if capturedUpgrades != 1 { + t.Fatalf("expected one persistent WebSocket, got %d", capturedUpgrades) } - if len(requests) != 3 { - t.Fatalf("expected three WebSocket requests, got %#v", requests) + if len(capturedRequests) != 3 { + t.Fatalf("expected three WebSocket requests, got %#v", capturedRequests) } - assertWebSocketRequest(t, requests[0], "reset", map[string]any{"seed": float64(42)}) - assertWebSocketRequest(t, requests[1], "step", map[string]any{"message": "hello"}) - assertWebSocketRequest(t, requests[2], "state", nil) - if !slices.Equal(safePaths, []string{"/health", "/metadata", "/schema"}) { - t.Fatalf("unexpected safe HTTP operations: %v", safePaths) + assertWebSocketRequest(t, capturedRequests[0], "reset", map[string]any{"seed": float64(42)}) + assertWebSocketRequest(t, capturedRequests[1], "step", map[string]any{"message": "hello"}) + assertWebSocketRequest(t, capturedRequests[2], "state", nil) + if !slices.Equal(capturedSafePaths, []string{"/health", "/metadata", "/schema"}) { + t.Fatalf("unexpected safe HTTP operations: %v", capturedSafePaths) } if authorizationCalls != 4 { t.Fatalf("expected one WebSocket and three HTTP authorization calls, got %d", authorizationCalls) @@ -107,6 +122,167 @@ func TestRunWebSocketShellUsesPersistentSocketForStatefulOperations(t *testing.T } } +func TestWebSocketHandshakeRetriesTransientFailures(t *testing.T) { + var attempts atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if attempts.Add(1) < 3 { + http.Error(w, "temporarily unavailable", http.StatusServiceUnavailable) + return + } + connection, err := (&websocket.Upgrader{}).Upgrade(w, r, nil) + if err != nil { + t.Errorf("upgrade WebSocket: %v", err) + return + } + defer connection.Close() + if _, _, err := connection.ReadMessage(); err != nil { + return + } + if err := connection.WriteJSON(map[string]any{"type": "state", "data": map[string]any{}}); err != nil { + t.Errorf("write WebSocket response: %v", err) + } + })) + defer server.Close() + + session := NewWebSocketRuntimeSession(server.URL, 30, nil) + session.handshakeRetryDelays = []time.Duration{time.Millisecond, time.Millisecond} + defer session.Close() + if _, err := session.Call(t.Context(), "state", ""); err != nil { + t.Fatal(err) + } + if attempts.Load() != 3 { + t.Fatalf("expected three handshake attempts, got %d", attempts.Load()) + } +} + +func TestWebSocketHandshakeDoesNotRetryPermanentFailure(t *testing.T) { + var attempts atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + attempts.Add(1) + http.Error(w, "unauthorized", http.StatusUnauthorized) + })) + defer server.Close() + + session := NewWebSocketRuntimeSession(server.URL, 30, nil) + session.handshakeRetryDelays = []time.Duration{time.Millisecond, time.Millisecond} + defer session.Close() + if _, err := session.Call(t.Context(), "state", ""); err == nil { + t.Fatal("expected WebSocket handshake to fail") + } + if attempts.Load() != 1 { + t.Fatalf("expected one handshake attempt, got %d", attempts.Load()) + } +} + +func TestWebSocketHandshakeStopsAfterRetryLimit(t *testing.T) { + var attempts atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + attempts.Add(1) + http.Error(w, "temporarily unavailable", http.StatusServiceUnavailable) + })) + defer server.Close() + + session := NewWebSocketRuntimeSession(server.URL, 30, nil) + session.handshakeRetryDelays = []time.Duration{time.Millisecond, time.Millisecond} + defer session.Close() + if _, err := session.Call(t.Context(), "state", ""); err == nil { + t.Fatal("expected WebSocket handshake to fail") + } + if attempts.Load() != 3 { + t.Fatalf("expected three handshake attempts, got %d", attempts.Load()) + } +} + +func TestWebSocketHandshakeCancellationStopsBackoff(t *testing.T) { + attemptReceived := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + select { + case <-attemptReceived: + default: + close(attemptReceived) + } + http.Error(w, "temporarily unavailable", http.StatusServiceUnavailable) + })) + defer server.Close() + + session := NewWebSocketRuntimeSession(server.URL, 30, nil) + session.handshakeRetryDelays = []time.Duration{time.Hour, time.Hour} + defer session.Close() + ctx, cancel := context.WithCancel(t.Context()) + callDone := make(chan error, 1) + go func() { + _, err := session.Call(ctx, "state", "") + callDone <- err + }() + <-attemptReceived + cancel() + select { + case err := <-callDone: + if !errors.Is(err, context.Canceled) { + t.Fatalf("expected context cancellation, got %v", err) + } + case <-time.After(time.Second): + t.Fatal("expected cancellation to stop handshake backoff") + } +} + +func TestWebSocketHandshakeRetriesShareOperationTimeout(t *testing.T) { + var attempts atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if attempts.Add(1) == 1 { + time.Sleep(600 * time.Millisecond) + http.Error(w, "temporarily unavailable", http.StatusServiceUnavailable) + return + } + connection, err := (&websocket.Upgrader{}).Upgrade(w, r, nil) + if err != nil { + t.Errorf("upgrade WebSocket: %v", err) + return + } + defer connection.Close() + if _, _, err := connection.ReadMessage(); err != nil { + return + } + time.Sleep(500 * time.Millisecond) + _ = connection.WriteJSON(map[string]any{"type": "state", "data": map[string]any{}}) + })) + defer server.Close() + + session := NewWebSocketRuntimeSession(server.URL, 1, nil) + session.handshakeRetryDelays = []time.Duration{time.Millisecond, time.Millisecond} + defer session.Close() + started := time.Now() + if _, err := session.Call(t.Context(), "state", ""); err == nil { + t.Fatal("expected shared operation timeout to expire") + } + if elapsed := time.Since(started); elapsed > 1500*time.Millisecond { + t.Fatalf("expected handshake retries to share the operation timeout, took %s", elapsed) + } +} + +func TestRetryableWebSocketHandshakeStatuses(t *testing.T) { + retryable := []int{408, 429, 500, 502, 503, 504} + for _, statusCode := range retryable { + if !isRetryableWebSocketHandshakeStatus(statusCode) { + t.Errorf("expected status %d to be retryable", statusCode) + } + } + for _, statusCode := range []int{400, 401, 403, 404} { + if isRetryableWebSocketHandshakeStatus(statusCode) { + t.Errorf("expected status %d not to be retryable", statusCode) + } + } +} + +func TestRetryableWebSocketHandshakeErrors(t *testing.T) { + if !isRetryableWebSocketHandshakeError(&net.DNSError{Err: "temporary failure", IsTemporary: true}) { + t.Fatal("expected network error to be retryable") + } + if isRetryableWebSocketHandshakeError(errors.New("invalid proxy configuration")) { + t.Fatal("expected local configuration error not to be retryable") + } +} + func TestRuntimeWebSocketURL(t *testing.T) { tests := []struct { baseURL string @@ -396,7 +572,33 @@ func TestCanceledQueuedCallDoesNotReachWebSocket(t *testing.T) { } } -func TestCallAndDrainUsesBoundedReadWhenTimeoutDisabled(t *testing.T) { +func TestCallAndDrainHasNoImplicitTimeoutWhenDisabled(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + connection, err := (&websocket.Upgrader{}).Upgrade(w, r, nil) + if err != nil { + t.Errorf("upgrade WebSocket: %v", err) + return + } + defer connection.Close() + if _, _, err := connection.ReadMessage(); err != nil { + t.Errorf("read WebSocket request: %v", err) + return + } + time.Sleep(25 * time.Millisecond) + if err := connection.WriteJSON(map[string]any{"type": "state", "data": map[string]any{}}); err != nil { + t.Errorf("write WebSocket response: %v", err) + } + })) + defer server.Close() + + session := NewWebSocketRuntimeSession(server.URL, 0, nil) + defer session.Close() + if _, err := session.CallAndDrain(t.Context(), "state", ""); err != nil { + t.Fatalf("expected timeout-disabled call to wait for its response: %v", err) + } +} + +func TestCallAndDrainBoundsResponseDrainAfterCancellation(t *testing.T) { requestReceived := make(chan struct{}) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { connection, err := (&websocket.Upgrader{}).Upgrade(w, r, nil) @@ -417,11 +619,22 @@ func TestCallAndDrainUsesBoundedReadWhenTimeoutDisabled(t *testing.T) { session := NewWebSocketRuntimeSession(server.URL, 0, nil) session.drainTimeout = 10 * time.Millisecond defer session.Close() - _, err := session.CallAndDrain(t.Context(), "state", "") - if err == nil { - t.Fatal("expected bounded drain to fail when the runtime does not respond") - } + ctx, cancel := context.WithCancel(t.Context()) + callDone := make(chan error, 1) + go func() { + _, err := session.CallAndDrain(ctx, "state", "") + callDone <- err + }() <-requestReceived + cancel() + select { + case err := <-callDone: + if err == nil { + t.Fatal("expected response drain to end after cancellation") + } + case <-time.After(time.Second): + t.Fatal("expected cancellation to bound the response drain") + } } func assertWebSocketRequest(