From 34bf6684fe1d6791ab8f9976e73a954ee65e9feb Mon Sep 17 00:00:00 2001
From: chill-czar <aisarthak2@gmail.com>
Date: Tue, 18 Aug 2026 01:15:35 +0530
Subject: [PATCH] fix(api): handle gRPC unavailable and deadline errors in
 volume orchestrator retry loop

---
 packages/api/internal/handlers/volume_util.go | 11 ++-
 .../api/internal/handlers/volume_util_test.go | 75 +++++++++++++++++++
 2 files changed, 85 insertions(+), 1 deletion(-)

diff --git a/packages/api/internal/handlers/volume_util.go b/packages/api/internal/handlers/volume_util.go
index 2ffc100b1f..10d8911459 100644
--- a/packages/api/internal/handlers/volume_util.go
+++ b/packages/api/internal/handlers/volume_util.go
@@ -11,6 +11,7 @@ import (
 	"github.com/gin-gonic/gin"
 	"github.com/google/uuid"
 	"go.uber.org/zap"
+	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 
 	"github.com/e2b-dev/infra/packages/api/internal"
@@ -250,9 +251,17 @@ func isUnknownVolumeTypeError(err error) (string, bool) {
 }
 
 func isRetryableError(err error) bool {
-	if errors.Is(err, net.ErrClosed) || errors.Is(err, context.DeadlineExceeded) {
+	if errors.Is(err, net.ErrClosed) || errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) {
 		return true
 	}
 
+	st, ok := status.FromError(err)
+	if ok {
+		switch st.Code() {
+		case codes.Unavailable, codes.DeadlineExceeded, codes.Canceled:
+			return true
+		}
+	}
+
 	return false
 }
diff --git a/packages/api/internal/handlers/volume_util_test.go b/packages/api/internal/handlers/volume_util_test.go
index 2acb3eeb09..60c6aca487 100644
--- a/packages/api/internal/handlers/volume_util_test.go
+++ b/packages/api/internal/handlers/volume_util_test.go
@@ -1,11 +1,16 @@
 package handlers
 
 import (
+	"context"
+	"errors"
+	"net"
 	"testing"
 
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 
 	"github.com/e2b-dev/infra/packages/api/internal/api"
 	"github.com/e2b-dev/infra/packages/api/internal/cfg"
@@ -151,3 +156,73 @@ func TestFindNodesByVolumeLabel(t *testing.T) {
 		})
 	}
 }
+
+func TestIsRetryableError(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name     string
+		err      error
+		expected bool
+	}{
+		{
+			name:     "nil error is not retryable",
+			err:      nil,
+			expected: false,
+		},
+		{
+			name:     "net.ErrClosed is retryable",
+			err:      net.ErrClosed,
+			expected: true,
+		},
+		{
+			name:     "context.DeadlineExceeded is retryable",
+			err:      context.DeadlineExceeded,
+			expected: true,
+		},
+		{
+			name:     "context.Canceled is retryable",
+			err:      context.Canceled,
+			expected: true,
+		},
+		{
+			name:     "gRPC codes.Unavailable is retryable",
+			err:      status.Error(codes.Unavailable, "connection refused"),
+			expected: true,
+		},
+		{
+			name:     "gRPC codes.DeadlineExceeded is retryable",
+			err:      status.Error(codes.DeadlineExceeded, "context deadline exceeded"),
+			expected: true,
+		},
+		{
+			name:     "gRPC codes.Canceled is retryable",
+			err:      status.Error(codes.Canceled, "request canceled"),
+			expected: true,
+		},
+		{
+			name:     "gRPC codes.NotFound is not retryable",
+			err:      status.Error(codes.NotFound, "volume not found"),
+			expected: false,
+		},
+		{
+			name:     "gRPC codes.InvalidArgument is not retryable",
+			err:      status.Error(codes.InvalidArgument, "bad volume name"),
+			expected: false,
+		},
+		{
+			name:     "generic error is not retryable",
+			err:      errors.New("something went wrong"),
+			expected: false,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			assert.Equal(t, tc.expected, isRetryableError(tc.err))
+		})
+	}
+}
+