3 files changed, 51 insertions, 11 deletions
diff --git a/packages/cli/src/ui/commands/ideCommand.test.ts b/packages/cli/src/ui/commands/ideCommand.test.ts
index 104eafb3..9b73acc9 100644
--- a/packages/cli/src/ui/commands/ideCommand.test.ts
+++ b/packages/cli/src/ui/commands/ideCommand.test.ts
@@ -227,7 +227,7 @@ describe('ideCommand', () => {
         }),
         expect.any(Number),
       );
-    });
+    }, 10000);
 
     it('should show an error if installation fails', async () => {
       mockInstall.mockResolvedValue({
diff --git a/packages/cli/src/ui/commands/ideCommand.ts b/packages/cli/src/ui/commands/ideCommand.ts
index 49766b8d..19a8090e 100644
--- a/packages/cli/src/ui/commands/ideCommand.ts
+++ b/packages/cli/src/ui/commands/ideCommand.ts
@@ -187,10 +187,6 @@ export const ideCommand = (config: Config | null): SlashCommand | null => {
       );
 
       const result = await installer.install();
-      if (result.success) {
-        config.setIdeMode(true);
-        context.services.settings.setValue(SettingScope.User, 'ideMode', true);
-      }
       context.ui.addItem(
         {
           type: result.success ? 'info' : 'error',
@@ -198,6 +194,39 @@ export const ideCommand = (config: Config | null): SlashCommand | null => {
         },
         Date.now(),
       );
+      if (result.success) {
+        context.services.settings.setValue(SettingScope.User, 'ideMode', true);
+        // Poll for up to 5 seconds for the extension to activate.
+        for (let i = 0; i < 10; i++) {
+          await config.setIdeModeAndSyncConnection(true);
+          if (
+            ideClient.getConnectionStatus().status ===
+            IDEConnectionStatus.Connected
+          ) {
+            break;
+          }
+          await new Promise((resolve) => setTimeout(resolve, 500));
+        }
+
+        const { messageType, content } = getIdeStatusMessage(ideClient);
+        if (messageType === 'error') {
+          context.ui.addItem(
+            {
+              type: messageType,
+              text: `Failed to automatically enable IDE integration. To fix this, run the CLI in a new terminal window.`,
+            },
+            Date.now(),
+          );
+        } else {
+          context.ui.addItem(
+            {
+              type: messageType,
+              text: content,
+            },
+            Date.now(),
+          );
+        }
+      }
     },
   };
 
diff --git a/packages/core/src/ide/ide-client.ts b/packages/core/src/ide/ide-client.ts
index e4d5f0ba..d6b1d0d2 100644
--- a/packages/core/src/ide/ide-client.ts
+++ b/packages/core/src/ide/ide-client.ts
@@ -104,8 +104,13 @@ export class IdeClient {
 
     this.setState(IDEConnectionStatus.Connecting);
 
+    const ideInfoFromFile = await this.getIdeInfoFromFile();
+    const workspacePath =
+      ideInfoFromFile.workspacePath ??
+      process.env['GEMINI_CLI_IDE_WORKSPACE_PATH'];
+
     const { isValid, error } = IdeClient.validateWorkspacePath(
-      process.env['GEMINI_CLI_IDE_WORKSPACE_PATH'],
+      workspacePath,
       this.currentIdeDisplayName,
       process.cwd(),
     );
@@ -115,7 +120,7 @@ export class IdeClient {
       return;
     }
 
-    const portFromFile = await this.getPortFromFile();
+    const portFromFile = ideInfoFromFile.port;
     if (portFromFile) {
       const connected = await this.establishConnection(portFromFile);
       if (connected) {
@@ -311,7 +316,10 @@ export class IdeClient {
     return port;
   }
 
-  private async getPortFromFile(): Promise<string | undefined> {
+  private async getIdeInfoFromFile(): Promise<{
+    port?: string;
+    workspacePath?: string;
+  }> {
     try {
       const ideProcessId = await getIdeProcessId();
       const portFile = path.join(
@@ -319,10 +327,13 @@ export class IdeClient {
         `gemini-ide-server-${ideProcessId}.json`,
       );
       const portFileContents = await fs.promises.readFile(portFile, 'utf8');
-      const port = JSON.parse(portFileContents).port;
-      return port.toString();
+      const ideInfo = JSON.parse(portFileContents);
+      return {
+        port: ideInfo?.port?.toString(),
+        workspacePath: ideInfo?.workspacePath,
+      };
     } catch (_) {
-      return undefined;
+      return {};
     }
   }