diff --git a/server/aws-lsp-codewhisperer/src/shared/utils.test.ts b/server/aws-lsp-codewhisperer/src/shared/utils.test.ts
index e83d04bbb3..85cd8817a7 100644
--- a/server/aws-lsp-codewhisperer/src/shared/utils.test.ts
+++ b/server/aws-lsp-codewhisperer/src/shared/utils.test.ts
@@ -134,26 +134,36 @@ describe('getClientName', () => {
 })
 
 describe('getOriginFromClientInfo', () => {
-    it('returns MD_IDE for SMUS-IDE client name', () => {
+    it('returns MD_IDE for client names starting with SMUS-IDE prefix', () => {
         const result = getOriginFromClientInfo('AmazonQ-For-SMUS-IDE-1.0.0')
         assert.strictEqual(result, 'MD_IDE')
     })
 
-    it('returns MD_IDE for SMUS-CE client name', () => {
+    it('returns MD_IDE for client names starting with SMUS-CE prefix', () => {
         const result = getOriginFromClientInfo('AmazonQ-For-SMUS-CE-1.0.0')
         assert.strictEqual(result, 'MD_IDE')
     })
 
-    it('returns MD_IDE for client names starting with SMUS-IDE prefix', () => {
+    it('returns MD_IDE for client names starting with SMAI-CE prefix', () => {
+        const result = getOriginFromClientInfo('AmazonQ-For-SMAI-CE-1.0.0')
+        assert.strictEqual(result, 'MD_IDE')
+    })
+
+    it('returns MD_IDE for SMUS-IDE client name', () => {
         const result = getOriginFromClientInfo('AmazonQ-For-SMUS-IDE')
         assert.strictEqual(result, 'MD_IDE')
     })
 
-    it('returns MD_IDE for client names starting with SMUS-CE prefix', () => {
+    it('returns MD_IDE for SMUS-CE client name', () => {
         const result = getOriginFromClientInfo('AmazonQ-For-SMUS-CE')
         assert.strictEqual(result, 'MD_IDE')
     })
 
+    it('returns MD_IDE for SMAI-CE client name', () => {
+        const result = getOriginFromClientInfo('AmazonQ-For-SMAI-CE')
+        assert.strictEqual(result, 'MD_IDE')
+    })
+
     it('returns IDE for non-SMUS client name', () => {
         const result = getOriginFromClientInfo('VSCode-Extension')
         assert.strictEqual(result, 'IDE')
diff --git a/server/aws-lsp-codewhisperer/src/shared/utils.ts b/server/aws-lsp-codewhisperer/src/shared/utils.ts
index e178de4a51..a7a95d8801 100644
--- a/server/aws-lsp-codewhisperer/src/shared/utils.ts
+++ b/server/aws-lsp-codewhisperer/src/shared/utils.ts
@@ -381,7 +381,12 @@ export function getClientName(lspParams: InitializeParams | undefined): string |
 }
 
 export function getOriginFromClientInfo(clientName: string | undefined): Origin {
-    if (clientName?.startsWith('AmazonQ-For-SMUS-IDE') || clientName?.startsWith('AmazonQ-For-SMUS-CE')) {
+    // TODO: Update with a new origin for SMAI case, as a short-term solution Sagemaker AI CE is using same origin as that of Sagemaker Unified Studio's IDE and CE
+    if (
+        clientName?.startsWith('AmazonQ-For-SMUS-IDE') ||
+        clientName?.startsWith('AmazonQ-For-SMUS-CE') ||
+        clientName?.startsWith('AmazonQ-For-SMAI-CE')
+    ) {
         return 'MD_IDE'
     }
     return 'IDE'