{AAZ} Added functionality to AAZ MCP Server

.devcontainer/Dockerfile

@@ -1,7 +1,11 @@
 FROM mcr.microsoft.com/devcontainers/python:3.12
 
-# Copy binaries from the official
+# Copy uv binaries for virtual environment management
 COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/
 
-# Remove after production
-RUN pip install --no-cache-dir mcp
+RUN pip install --no-cache-dir mcp \
+    && mkdir -p /workspaces-src \
+    && chown -R vscode:vscode /workspaces-src
+
+# Ensure bash is default shell
+SHELL ["/bin/bash", "-c"]

.devcontainer/devcontainer.json

@@ -1,17 +1,12 @@
 {
   "name": "Azure CLI Dev Container",
   "build": {
-    "dockerfile": "Dockerfile"
+    "dockerfile": "Dockerfile",
+    "context": "."
   },
   "features": {
     "ghcr.io/devcontainers/features/github-cli:1": {}
   },
-  "workspaceFolder": "/workspaces",
-  "onCreateCommand": "uv venv .venv --seed",
-  "postCreateCommand": "REPO_NAME=$(basename $GITHUB_REPOSITORY) && cat $REPO_NAME/.devcontainer/init.sh >> ~/.bashrc && mkdir .vscode && cp $REPO_NAME/.devcontainer/mcp.json .vscode/",
-  "hostRequirements": {
-    "cpus": 16,
-    "memory": "64gb",
-    "storage": "128gb"
-  }
-}
+  "workspaceFolder": "/workspaces/azure-cli",
+  "postCreateCommand": "mkdir -p /workspaces/azure-cli/.vscode && cp /workspaces/azure-cli/.devcontainer/mcp.json /workspaces/azure-cli/.vscode/ && /workspaces/azure-cli/.devcontainer/init.sh"
+}

.devcontainer/init.sh

@@ -1,74 +1,89 @@
-
-source .venv/bin/activate
-
-# Logout default account
-export GITHUB_TOKEN=
-
-# Check `repo` scope exists or not
-if gh auth status -a 2>/dev/null | grep "Token scopes: " | grep -q "repo"; then
-    echo "You now have access to GitHub."
-else
+#!/bin/bash
+set -euo pipefail
+
+echo "======================================"
+echo " Azure CLI Devcontainer Setup Started "
+echo "======================================"
+
+# -----------------------------
+# Remove existing virtual environments
+# -----------------------------
+VENV_DIRS=(
+    "$HOME/.venv"
+    "/workspaces/.venv"
+)
+for VENV in "${VENV_DIRS[@]}"; do
+    if [ -d "$VENV" ]; then
+        echo "Removing existing virtual environment at $VENV..."
+        rm -rf "$VENV"
+    fi
+done
+
+# -----------------------------
+# GitHub authentication
+# -----------------------------
+export GITHUB_TOKEN=${GITHUB_TOKEN:-}
+export GITHUB_USER=$(gh api user --jq .login 2>/dev/null || echo "")
+
+if [ -z "$GITHUB_USER" ]; then
+    echo "No GitHub user detected, running gh auth login..."
     gh auth login -p https -w
+    export GITHUB_USER=$(gh api user --jq .login || echo "unknown")
 fi
-
-# Check `aaz-dev` is available or not
-if ! command -v aaz-dev &> /dev/null; then
-    GREEN="\033[0;32m"
-    YELLOW="\033[0;33m"
-    NC="\033[0m"  # no color
-
-    set_or_add_remote() {
-        local REPO=$1
-        local REMOTE=$2
-        local DIR="/workspaces/$REPO"
-        local OWNER=$([ "$REMOTE" = "origin" ] && echo "$GITHUB_USER" || echo "Azure")
-        local URL="https://github.com/$OWNER/$REPO.git"
-
-        git -C "$DIR" remote get-url "$REMOTE" &>/dev/null || git -C "$DIR" remote add "$REMOTE" "$URL"
-        git -C "$DIR" remote set-url "$REMOTE" "$URL"
-    }
-
-    setup_repo() {
-        local REPO=$1
-        local DIR="/workspaces/$REPO"
-
-        echo
+echo "Using GitHub user: $GITHUB_USER"
+
+# -----------------------------
+# Install aaz-dev
+# -----------------------------
+pip install --upgrade pip
+pip install --no-cache-dir aaz-dev
+
+# -----------------------------
+# Repo paths
+# -----------------------------
+export AZDEV_CLI_PATH="/workspaces/azure-cli"
+export CLI_PATH="/workspaces/azure-cli"
+export CLI_EXTENSION_PATH="/workspaces-src/azure-cli-extensions"
+export AAZ_PATH="/workspaces-src/aaz"
+export SWAGGER_PATH="/workspaces-src/azure-rest-api-specs"
+
+# -----------------------------
+# Ensure /workspaces-src exists
+# -----------------------------
+mkdir -p /workspaces-src
+
+# -----------------------------
+# Clone repos only if they don't exist
+# -----------------------------
+OTHER_REPOS=("azure-cli-extensions" "aaz" "azure-rest-api-specs")
+
+for REPO in "${OTHER_REPOS[@]}"; do
+    DIR="/workspaces-src/$REPO"
+    if [ -d "$DIR/.git" ]; then
+        echo "$REPO already exists, skipping clone."
+    else
+        echo "Cloning $REPO..."
         gh repo fork "Azure/$REPO" --clone=false --default-branch-only
-
-        if [ -d "$DIR" ]; then
-            set_or_add_remote "$REPO" origin
-            set_or_add_remote "$REPO" upstream
-        else
-            git clone "https://github.com/$GITHUB_USER/$REPO.git" --single-branch --no-tags
-            set_or_add_remote "$REPO" upstream
-
-            # Synchronize with upstream
-            BRANCH=$(git -C "$DIR" remote show upstream | grep "HEAD branch" | awk '{print $NF}')
-            git -C "$DIR" pull -r upstream "$BRANCH"
-        fi
-    }
-
-    SECONDS=0
-
-    echo
-    uv pip install aaz-dev --link-mode=copy
-
-    # `azdev` repositories
-    setup_repo "azure-cli"
-    setup_repo "azure-cli-extensions"
-
-    azdev setup -c -r ./azure-cli-extensions
-
-    # `aaz-dev` repositories
-    setup_repo "aaz"
-    setup_repo "azure-rest-api-specs"
-
-    ELAPSED_TIME=$SECONDS
-
-    echo -e "\n${YELLOW}Elapsed time: $((ELAPSED_TIME / 60))m $((ELAPSED_TIME % 60))s.${NC}"
-    echo -e "\n${GREEN}Finished setup! Please launch the codegen tool via:${NC}"
-    echo -e "${GREEN}\$ aaz-dev run -c azure-cli -e azure-cli-extensions -s azure-rest-api-specs -a aaz${NC}\n"
-else
-    echo -e "\nPlease launch the codegen tool via:"
-    echo -e "$ aaz-dev run -c azure-cli -e azure-cli-extensions -s azure-rest-api-specs -a aaz\n"
-fi
+        git clone "https://github.com/$GITHUB_USER/$REPO.git" "$DIR" --single-branch --no-tags
+        git -C "$DIR" remote add upstream "https://github.com/Azure/$REPO.git"
+        git -C "$DIR" pull -r upstream main || true
+    fi
+done
+
+# -----------------------------
+# Install aaz-flow using uv (no separate venv)
+# -----------------------------
+echo "Installing aaz-flow and dependencies using uv..."
+cd "/workspaces/azure-cli/tools/aaz-flow"
+sudo uv pip install -e . --system
+cd -
+
+# -----------------------------
+# Final setup
+# -----------------------------
+echo "Running azdev setup..."
+python3 -m venv /workspaces-src/.azdev-env
+source /workspaces-src/.azdev-env/bin/activate
+pip install --upgrade pip
+pip install azdev
+azdev setup -c /workspaces/azure-cli -r /workspaces-src/azure-cli-extensions

.devcontainer/mcp.json

@@ -1,38 +1,16 @@
 {
-  "inputs": [
-    {
-      "type": "promptString",
-      "id": "extension_or_module_name",
-      "description": "The name of the Azure CLI extension or module to generate commands for."
-    },
-    {
-      "type": "promptString",
-      "id": "swagger_module_path",
-      "description": "The file path to the Swagger module."
-    },
-    {
-      "type": "promptString",
-      "id": "resource_provider",
-      "description": "The Azure resource provider for which to generate commands."
-    },
-    {
-      "type": "promptString",
-      "id": "swagger_tag",
-      "description": "The Swagger tag to use for command generation."
-    }
-  ],
   "servers": {
     "AAZ Flow": {
       "command": "python",
       "args": [
-        "azure-cli/tools/aaz-flow/main.py"
+        "tools/aaz-flow/main.py"
       ],
       "env": {
-        "AAZ_PATH": "/workspaces/aaz",
-        "CLI_PATH": "/workspaces/azure-cli",
-        "CLI_EXTENSION_PATH": "/workspaces/azure-cli-extensions",
-        "SWAGGER_PATH": "/workspaces/azure-rest-api-specs"
+          "AAZ_PATH": "/workspaces-src/aaz",
+          "CLI_PATH": "/workspaces-src/azure-cli",
+          "CLI_EXTENSION_PATH": "/workspaces-src/azure-cli-extensions",
+          "SWAGGER_PATH": "/workspaces-src/azure-rest-api-specs"
       }
     }
   }
-}
+}

tools/aaz-flow/README.md

@@ -4,9 +4,11 @@ MCP Server for the AAZ API, enabling prune command-line interface, implement cus
 Please note that AAZ Flow is currently in early development. The functionality and available tools are subject to change and expansion as we continue to develop and improve the server.
 
 ### Tools
-1. `generate_code`
-   - xx
-   - Inputs:
-     - xx
-     - xx
-   - Returns:
+`generate_code`
+
+## Implementation
+1. Performs elicitation for user input to perform workflow
+2. Generates content using llm sampling
+3. Executes AAZ Flow commands directly
+4. Generates tests using llm sampling
+5. Uses tool transformation to make the internal tooling more friendly for llms

tools/aaz-flow/__init__.py

@@ -0,0 +1,6 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+# AAZ Flow package