Merge

Author: MarceloRGonc

.devcontainer/Dockerfile

@@ -18,6 +18,8 @@ RUN <<-```
 ENV LANG en_US.UTF-8
 ENV LANGUAGE en_US:en
 ENV LC_ALL en_US.UTF-8
+ENV AZURE_CONFIG_DIR="/tmp/azure"
+ENV AZURE_EXTENSION_DIR="/opt/azure/config/cliextensions"
 
 RUN <<-```
   # Install yq
@@ -38,6 +40,17 @@ RUN <<-```
 
   # Install Azure CLI
   curl -sL https://aka.ms/InstallAzureCLIDeb | bash
+  mkdir -p /tmp/azure
+  mkdir -p /opt/azure/config/cliextensions
+
+  # Configure Azure CLI extensions
+  az config set extension.use_dynamic_install=yes_without_prompt
+  az extension add --name reservation --only-show-errors || true
+  az extension add --name resource-graph --only-show-errors || true
+  az extension add --name costmanagement --only-show-errors || true
+  az extension add --name billing-benefits --only-show-errors || true
+  az extension add --name quotas --only-show-errors || true
+  az extension add --name ssh --only-show-errors || true
 
   # Install Google Cloud CLI
   curl -sSL https://dl.google.com/dl/cloudsdk/channels/rapid/downloads/google-cloud-cli-516.0.0-linux-x86_64.tar.gz -o /tmp/gcloud.tar.gz

Dockerfile

@@ -77,4 +77,4 @@ COPY docker-entrypoint.sh .
 RUN chmod +x docker-entrypoint.sh
 ENTRYPOINT ["./docker-entrypoint.sh"]
 
-EXPOSE 80
+EXPOSE 80

engine.Dockerfile

@@ -40,9 +40,23 @@ RUN <<-```
 
 ENV LD_LIBRARY_PATH=""
 ENV AZURE_CONFIG_DIR="/tmp/azure"
+ENV AZURE_EXTENSION_DIR="/opt/azure/config/cliextensions"
 
-RUN pip3 install azure-cli==2.74.0 && \
+RUN <<-```
+    set -ex
+    pip3 install azure-cli==2.74.0
     mkdir -p /tmp/azure
+    mkdir -p /opt/azure/config/cliextensions
+    if command -v az >/dev/null 2>&1; then
+        az config set extension.use_dynamic_install=yes_without_prompt
+        az extension add --name reservation --only-show-errors || true
+        az extension add --name resource-graph --only-show-errors || true
+        az extension add --name costmanagement --only-show-errors || true
+        az extension add --name billing-benefits --only-show-errors || true
+        az extension add --name quotas --only-show-errors || true
+        az extension add --name ssh --only-show-errors || true
+    fi
+```
 
 ENV CLOUDSDK_CONFIG="/tmp/gcloud"
 RUN <<-```
@@ -97,4 +111,4 @@ RUN ./tools/link-packages.sh
 ENV NODE_OPTIONS=--no-node-snapshot
 
 ENTRYPOINT [ "/bin/bash", "-c" ]
-CMD [ "cp -r /var/tmp-base/. /tmp/ && /lambda-entrypoint.sh main.handler; exit $?" ]
+CMD [ "cp -r /var/tmp-base/. /tmp/ && /lambda-entrypoint.sh main.handler; exit $?" ]

packages/blocks/azure/src/lib/azure-cli.ts

@@ -3,6 +3,13 @@ import { mkdtempSync } from 'fs';
 import { tmpdir } from 'os';
 import { join } from 'path';
 
+function setEnvIfExists(envVars: Record<string, string>, key: string): void {
+  const value = process.env[key];
+  if (value) {
+    envVars[key] = value;
+  }
+}
+
 export async function runCommand(
   command: string,
   credentials: any,
@@ -13,10 +20,8 @@ export async function runCommand(
     PATH: process.env['PATH'] || '',
   };
 
-  const processAzureConfigDir = process.env['AZURE_CONFIG_DIR'];
-  if (processAzureConfigDir) {
-    envVars['AZURE_CONFIG_DIR'] = processAzureConfigDir;
-  }
+  setEnvIfExists(envVars, 'AZURE_CONFIG_DIR');
+  setEnvIfExists(envVars, 'AZURE_EXTENSION_DIR');
 
   if (!shouldUseHostCredentials) {
     const azureConfigDir = mkdtempSync(join(tmpdir(), 'azure-cli'));

packages/blocks/azure/test/azure-cli.test.ts

@@ -31,6 +31,7 @@ describe('azureCli', () => {
       {
         PATH: process.env['PATH'],
         AZURE_CONFIG_DIR: expect.any(String),
+        AZURE_EXTENSION_DIR: expect.any(String),
       },
     );
     expect(commonMock.runCliCommand).toHaveBeenCalledWith(
@@ -39,6 +40,7 @@ describe('azureCli', () => {
       {
         PATH: process.env['PATH'],
         AZURE_CONFIG_DIR: expect.any(String),
+        AZURE_EXTENSION_DIR: expect.any(String),
       },
     );
   });
@@ -64,6 +66,7 @@ describe('azureCli', () => {
       {
         PATH: process.env['PATH'],
         AZURE_CONFIG_DIR: expect.any(String),
+        AZURE_EXTENSION_DIR: expect.any(String),
       },
     );
     expect(commonMock.runCliCommand).toHaveBeenNthCalledWith(
@@ -73,6 +76,7 @@ describe('azureCli', () => {
       {
         PATH: process.env['PATH'],
         AZURE_CONFIG_DIR: expect.any(String),
+        AZURE_EXTENSION_DIR: expect.any(String),
       },
     );
     expect(commonMock.runCliCommand).toHaveBeenNthCalledWith(
@@ -82,6 +86,7 @@ describe('azureCli', () => {
       {
         PATH: process.env['PATH'],
         AZURE_CONFIG_DIR: expect.any(String),
+        AZURE_EXTENSION_DIR: expect.any(String),
       },
     );
   });
@@ -105,6 +110,7 @@ describe('azureCli', () => {
       'az',
       {
         PATH: process.env['PATH'],
+        AZURE_EXTENSION_DIR: expect.any(String),
       },
     );
     expect(commonMock.runCliCommand).toHaveBeenNthCalledWith(
@@ -113,6 +119,7 @@ describe('azureCli', () => {
       'az',
       {
         PATH: process.env['PATH'],
+        AZURE_EXTENSION_DIR: expect.any(String),
       },
     );
   });
@@ -132,6 +139,7 @@ describe('azureCli', () => {
       'az',
       {
         PATH: process.env['PATH'],
+        AZURE_EXTENSION_DIR: expect.any(String),
       },
     );
     process.env = originalEnv;
@@ -143,6 +151,7 @@ describe('azureCli', () => {
       ...originalEnv,
       PATH: '/mock/path',
       AZURE_CONFIG_DIR: '/mock/config/dir',
+      AZURE_EXTENSION_DIR: '/mock/config/dir/extensions',
     };
 
     commonMock.runCliCommand.mockResolvedValue('mock result');
@@ -157,6 +166,44 @@ describe('azureCli', () => {
       {
         PATH: process.env['PATH'],
         AZURE_CONFIG_DIR: process.env['AZURE_CONFIG_DIR'],
+        AZURE_EXTENSION_DIR: '/mock/config/dir/extensions',
+      },
+    );
+    process.env = originalEnv;
+  });
+
+  test('should use temp config dir but preserve extension dir when useHostSession is false and AZURE_CONFIG_DIR is set', async () => {
+    const originalEnv = process.env;
+    process.env = {
+      ...originalEnv,
+      PATH: '/mock/path',
+      AZURE_CONFIG_DIR: '/mock/config/dir',
+      AZURE_EXTENSION_DIR: '/mock/config/dir/extensions',
+    };
+
+    commonMock.runCliCommand.mockResolvedValueOnce('login result');
+    commonMock.runCliCommand.mockResolvedValueOnce('mock result');
+
+    const result = await runCommand('some command', credentials, false);
+
+    expect(result).toBe('mock result');
+    expect(commonMock.runCliCommand).toHaveBeenCalledTimes(2);
+    expect(commonMock.runCliCommand).toHaveBeenCalledWith(
+      `login --service-principal --username ${credentials.clientId} --password ${credentials.clientSecret} --tenant ${credentials.tenantId}`,
+      'az',
+      {
+        PATH: process.env['PATH'],
+        AZURE_CONFIG_DIR: expect.any(String),
+        AZURE_EXTENSION_DIR: '/mock/config/dir/extensions',
+      },
+    );
+    expect(commonMock.runCliCommand).toHaveBeenCalledWith(
+      'some command',
+      'az',
+      {
+        PATH: process.env['PATH'],
+        AZURE_CONFIG_DIR: expect.any(String),
+        AZURE_EXTENSION_DIR: '/mock/config/dir/extensions',
       },
     );
     process.env = originalEnv;
@@ -176,6 +223,7 @@ describe('azureCli', () => {
       {
         PATH: process.env['PATH'],
         AZURE_CONFIG_DIR: expect.any(String),
+        AZURE_EXTENSION_DIR: expect.any(String),
       },
     );
   });
@@ -198,6 +246,7 @@ describe('azureCli', () => {
       {
         PATH: process.env['PATH'],
         AZURE_CONFIG_DIR: expect.any(String),
+        AZURE_EXTENSION_DIR: expect.any(String),
       },
     );
   });

packages/react-ui/src/app/features/ai/lib/assistant-ui-chat-hook.ts

@@ -6,7 +6,7 @@ import { useChat } from '@ai-sdk/react';
 import { AssistantRuntime } from '@assistant-ui/react';
 import { useAISDKRuntime } from '@assistant-ui/react-ai-sdk';
 import { toast } from '@openops/components/ui';
-import { flowHelper } from '@openops/shared';
+import { flowHelper, FlowVersion } from '@openops/shared';
 import { getFrontendToolDefinitions } from '@openops/ui-kit';
 import { useQuery, useQueryClient } from '@tanstack/react-query';
 import { DefaultChatTransport, ToolSet, UIMessage } from 'ai';
@@ -135,14 +135,10 @@ export const useAssistantChat = ({
           stepDetails &&
           flowId
         ) {
-          const stepId =
-            flowHelper.getStep(context.flowVersion, context.selectedStep)?.id ??
-            context.selectedStep;
-
           return await aiChatApi.open(
             flowId,
             getBlockName(stepDetails),
-            stepId,
+            getStepId(context.flowVersion, context.selectedStep),
             getActionName(stepDetails),
           );
         }
@@ -337,7 +333,16 @@ export const useAssistantChat = ({
           chatMode === ChatMode.StepSettings
         ) {
           await aiChatApi.delete(oldChatId);
+
+          const conversation = await aiChatApi.open(
+            context.flowVersion.flowId,
+            getBlockName(stepDetails),
+            getStepId(context.flowVersion, context.selectedStep),
+            getActionName(stepDetails),
+          );
+
           chat.setMessages([]);
+          onChatIdChange(conversation.chatId);
         } else {
           onChatIdChange(null);
         }
@@ -351,7 +356,7 @@ export const useAssistantChat = ({
         `There was an error canceling the current run and invalidating queries while creating a new chat: ${error}`,
       );
     }
-  }, [chatId, chat, getBuilderState, chatMode, onChatIdChange]);
+  }, [chatId, chat, getBuilderState, chatMode, stepDetails, onChatIdChange]);
 
   return {
     runtime,
@@ -363,3 +368,7 @@ export const useAssistantChat = ({
     chatStatus: chat.status,
   };
 };
+
+const getStepId = (flowVersion: FlowVersion, selectedStep: string) => {
+  return flowHelper.getStep(flowVersion, selectedStep)?.id ?? selectedStep;
+};