Return correct status code for non existing sandbox in orchestrator in keep alive (#1405)

Author: jakubno

packages/api/internal/orchestrator/keep_alive.go

@@ -61,6 +61,10 @@ func (o *Orchestrator) KeepAliveFor(ctx context.Context, sandboxID string, durat
 	}
 	err = o.UpdateSandbox(ctx, sandboxID, sbx.EndTime, sbx.ClusterID, sbx.NodeID)
 	if err != nil {
+		if errors.Is(err, ErrSandboxNotFound) {
+			return &api.APIError{Code: http.StatusNotFound, ClientMsg: "Sandbox not found", Err: err}
+		}
+
 		return &api.APIError{Code: http.StatusInternalServerError, ClientMsg: "Error when setting sandbox timeout", Err: err}
 	}
 

packages/api/internal/orchestrator/update_instance.go

@@ -9,6 +9,8 @@ import (
 	"github.com/google/uuid"
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/trace"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/types/known/timestamppb"
 
 	"github.com/e2b-dev/infra/packages/api/internal/utils"
@@ -41,9 +43,14 @@ func (o *Orchestrator) UpdateSandbox(
 			EndTime:   timestamppb.New(endTime),
 		},
 	)
-
-	err = utils.UnwrapGRPCError(err)
 	if err != nil {
+		grpcErr, ok := status.FromError(err)
+		if ok && grpcErr.Code() == codes.NotFound {
+			return ErrSandboxNotFound
+		}
+
+		err = utils.UnwrapGRPCError(err)
+
 		return fmt.Errorf("failed to update sandbox '%s': %w", sandboxID, err)
 	}