comments addressed

Author: mgupta-om

packages/core/src/awsService/sagemaker/detached-server/kubectlClientStub.ts

@@ -24,6 +24,8 @@ export interface HyperpodCluster {
     clusterName: string
     clusterArn: string
     status: string
+    eksClusterName?: string
+    eksClusterArn?: string
     regionCode: string
 }
 
@@ -34,24 +36,30 @@ export interface WorkspaceConnectionResult {
     sessionId: string
 }
 
-/** Refresh the token 1 minute before it expires. */
+/** Buffer time (ms) before token expiry to trigger a proactive refresh, avoiding mid-request expirations. */
 const tokenRefreshBufferMs = 60_000
 
+export interface EksClusterInfo {
+    name?: string
+    endpoint?: string
+    certificateAuthority?: { data?: string }
+}
+
 export class KubectlClient {
     private kubeConfig: k8s.KubeConfig
-    private k8sApi!: k8s.CustomObjectsApi
+    private k8sApi: k8s.CustomObjectsApi | undefined
     private tokenExpiry: number = 0
 
-    private constructor(
-        private readonly eksCluster: any,
+    protected constructor(
+        private readonly eksCluster: EksClusterInfo,
         private readonly hyperpodCluster: HyperpodCluster,
         private readonly credentials: AwsCredentialIdentity | Provider<AwsCredentialIdentity>
     ) {
         this.kubeConfig = new k8s.KubeConfig()
     }
 
-    static async create(
-        eksCluster: any,
+    static async createForCluster(
+        eksCluster: EksClusterInfo,
         hyperpodCluster: HyperpodCluster,
         credentials: AwsCredentialIdentity | Provider<AwsCredentialIdentity>
     ): Promise<KubectlClient> {
@@ -60,10 +68,17 @@ export class KubectlClient {
         return client
     }
 
-    getEksCluster(): any {
+    getEksCluster(): EksClusterInfo {
         return this.eksCluster
     }
 
+    protected getApi(): k8s.CustomObjectsApi {
+        if (!this.k8sApi) {
+            throw new Error('[Hyperpod] KubectlClient not initialized — call createForCluster()')
+        }
+        return this.k8sApi
+    }
+
     async createWorkspaceConnection(devSpace: HyperpodDevSpace): Promise<WorkspaceConnectionResult> {
         try {
             await this.ensureValidToken()
@@ -84,7 +99,7 @@ export class KubectlClient {
                 },
             }
 
-            const response = await this.k8sApi.createNamespacedCustomObject(
+            const response = await this.getApi().createNamespacedCustomObject(
                 group,
                 version,
                 devSpace.namespace,
@@ -96,6 +111,8 @@ export class KubectlClient {
             const status = body.status
             const presignedUrl = status?.workspaceConnectionUrl
             const connectionType = status?.workspaceConnectionType
+            // tokenValue and sessionId may not be present in all API responses.
+            // When empty, the existing connection flow in model.ts falls back to parsing these from the presigned URL.
             const token = status?.tokenValue ?? ''
             const sessionId = status?.sessionId ?? ''
 
@@ -107,7 +124,7 @@ export class KubectlClient {
         }
     }
 
-    private async initKubeConfig(): Promise<void> {
+    protected async initKubeConfig(): Promise<void> {
         if (!this.eksCluster.name || !this.eksCluster.endpoint) {
             return
         }
@@ -147,13 +164,13 @@ export class KubectlClient {
         this.k8sApi = this.kubeConfig.makeApiClient(k8s.CustomObjectsApi)
     }
 
-    private async ensureValidToken(): Promise<void> {
+    protected async ensureValidToken(): Promise<void> {
         if (Date.now() >= this.tokenExpiry - tokenRefreshBufferMs) {
             await this.refreshToken()
         }
     }
 
-    private async refreshToken(): Promise<void> {
+    protected async refreshToken(): Promise<void> {
         if (!this.eksCluster.name) {
             return
         }

packages/core/src/awsService/sagemaker/explorer/sagemakerHyperpodNode.ts

@@ -100,7 +100,7 @@ export class SagemakerHyperpodNode extends AWSTreeNodeBase {
                         }
                         return creds as AwsCredentialIdentity
                     }
-                    kcClient = await KubectlClient.create(eksCluster, cluster, credentialsProvider)
+                    kcClient = await KubectlClient.createForCluster(eksCluster, cluster, credentialsProvider)
                     this.kubectlClients.set(cluster.clusterName, kcClient)
                 }
                 const spacesPerCluster = await kcClient.getSpacesForCluster(eksCluster)

packages/core/src/awsService/sagemaker/hyperpodCommands.ts

@@ -65,7 +65,6 @@ export async function connectToHyperPodDevSpace(node: SagemakerDevSpaceNode): Pr
         // Call createWorkspaceConnection to get presigned URL with credentials
         const workspaceConnection = await kubectlClient.createWorkspaceConnection(node.devSpace)
         const connectionUrl = workspaceConnection.url
-        const connectionToken = workspaceConnection.token
 
         const remoteEnv = await prepareDevEnvConnection({
             spaceArn: '',
@@ -81,7 +80,6 @@ export async function connectToHyperPodDevSpace(node: SagemakerDevSpaceNode): Pr
             eksEndpoint: eksCluster?.endpoint,
             eksCertAuthData: eksCluster?.certificateAuthority?.data,
             wsUrl: connectionUrl,
-            token: connectionToken,
         })
 
         await startVscodeRemote(

packages/core/src/awsService/sagemaker/model.ts

@@ -181,17 +181,12 @@ export async function prepareDevEnvConnection(opts: DevEnvConnectionOptions) {
     const sshPrefix = getSshPrefix(connectionType)
     let hostname: string
     if (connectionType === 'sm_hp') {
-        // Always construct hostname from workspace components for a meaningful window title.
-        // The `session` param is the SSM sessionId (e.g. "eks-sagemaker-jupyter"), not a hostname.
-        let hpSession: string | undefined
-        if (workspaceName && namespace && clusterName && region && accountId) {
+        let hpSession = session
+        if (!hpSession) {
             const proposedSession = `${workspaceName}_${namespace}_${clusterName}_${region}_${accountId}`
             hpSession = isValidSshHostname(proposedSession)
                 ? proposedSession
-                : createValidSshSession(workspaceName, namespace, clusterName, region, accountId)
-        }
-        if (!hpSession) {
-            hpSession = session || 'hyperpod'
+                : createValidSshSession(workspaceName!, namespace!, clusterName!, region!, accountId!)
         }
         hostname = `${sshPrefix}${hpSession}`
     } else {