start building realistic urls for CONNECT

pkg/common/url.go

@@ -42,3 +42,16 @@ func BuildServeURL(externalUrl, urlType string, stub *types.StubWithRelated) str
 	}
 	return fmt.Sprintf("%s://%s/%s/id/%s", parsedUrl.Scheme, parsedUrl.Host, stub.Type.Kind(), stub.ExternalId)
 }
+
+func BuildShellURL(externalUrl, urlType string, stub *types.StubWithRelated) string {
+	parsedUrl, err := url.Parse(externalUrl)
+	if err != nil {
+		return ""
+	}
+
+	if urlType == InvokeUrlTypeHost {
+		return fmt.Sprintf("%s://%s.%s", parsedUrl.Scheme, stub.ExternalId, parsedUrl.Host)
+	}
+
+	return fmt.Sprintf("%s://%s/%s/%s", parsedUrl.Scheme, parsedUrl.Host, "shell", stub.ExternalId)
+}

pkg/gateway/services/stub.go

@@ -255,13 +255,19 @@ func (gws *GatewayService) GetURL(ctx context.Context, in *pb.GetURLRequest) (*p
 		in.UrlType = gws.appConfig.GatewayService.InvokeURLType
 	}
 
-	// Get URL for Serves
+	// Get URL for Serves or Shells
 	if stub.Type.IsServe() {
 		invokeUrl := common.BuildServeURL(gws.appConfig.GatewayService.HTTP.GetExternalURL(), in.UrlType, stub)
 		return &pb.GetURLResponse{
 			Ok:  true,
 			Url: invokeUrl,
 		}, nil
+	} else if stub.Type.IsShell() {
+		invokeUrl := common.BuildShellURL(gws.appConfig.GatewayService.HTTP.GetExternalURL(), in.UrlType, stub)
+		return &pb.GetURLResponse{
+			Ok:  true,
+			Url: invokeUrl,
+		}, nil
 	}
 
 	// Get URL for Deployments

sdk/src/beta9/abstractions/endpoint.py

@@ -27,6 +27,7 @@
     StartEndpointServeResponse,
     StopEndpointServeRequest,
 )
+from ..clients.gateway import GetUrlRequest, GetUrlResponse
 from ..clients.shell import CreateShellRequest, ShellServiceStub
 from ..env import is_local
 from ..type import Autoscaler, GpuType, GpuTypeAlias, QueueDepthAutoscaler, TaskPolicy
@@ -609,7 +610,7 @@ def notify(*_, **__):
         terminal.warn("Endpoint serve timed out. Container has been stopped.")
 
     @with_grpc_error_handling
-    def shell(self, timeout: int = 0):
+    def shell(self, timeout: int = 0, url_type: str = ""):
         stub_type = ENDPOINT_SHELL_STUB_TYPE
 
         if not self.parent.prepare_runtime(
@@ -626,6 +627,18 @@ def shell(self, timeout: int = 0):
         )
         print(r)
 
+        res: GetUrlResponse = self.parent.gateway_stub.get_url(
+            GetUrlRequest(
+                stub_id=self.parent.stub_id,
+                deployment_id=getattr(self, "deployment_id", ""),
+                url_type=url_type,
+            )
+        )
+        if not res.ok:
+            return terminal.error("Failed to get shell connection URL", exit=False)
+
+        print(res.url)
+
         # # Then, issue connection request to the container using CONNECT http method
         # import http.client
 

sdk/src/beta9/cli/shell.py

@@ -44,13 +44,19 @@ def common(**_):
     default=0,
     help="The inactivity timeout for the shell instance in seconds. Set to -1 for no timeout. Set to 0 to use default timeout (10 minutes)",
 )
+@click.option(
+    "--url-type",
+    help="The type of URL to get back. [default is determined by the server] ",
+    type=click.Choice(["host", "path"]),
+)
 @extraclick.pass_service_client
 @click.pass_context
 def shell(
     ctx: click.Context,
     service: ServiceClient,
     entrypoint: str,
     timeout: Optional[int] = None,
+    url_type: str = "path",
 ):
     current_dir = os.getcwd()
     if current_dir not in sys.path:
@@ -78,4 +84,4 @@ def shell(
     if hasattr(user_obj, "set_handler"):
         user_obj.set_handler(f"{module_name}:{obj_name}")
 
-    user_obj.shell(timeout=int(timeout))  # type:ignore
+    user_obj.shell(timeout=int(timeout), url_type=url_type)  # type:ignore