containerized aks agent

Author: mainred

src/aks-agent/azext_aks_agent/__init__.py

@@ -3,27 +3,26 @@
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
 
-
-import os
+from azext_aks_agent._client_factory import CUSTOM_MGMT_AKS
 
 # pylint: disable=unused-import
-import azext_aks_agent._help
-from azext_aks_agent._consts import (
-    CONST_AGENT_CONFIG_PATH_DIR_ENV_KEY,
-    CONST_AGENT_NAME,
-    CONST_AGENT_NAME_ENV_KEY,
-    CONST_DISABLE_PROMETHEUS_TOOLSET_ENV_KEY,
-    CONST_PRIVACY_NOTICE_BANNER,
-    CONST_PRIVACY_NOTICE_BANNER_ENV_KEY,
-)
 from azure.cli.core import AzCommandsLoader
-from azure.cli.core.api import get_config_dir
+from azure.cli.core.profiles import register_resource_type
+
+
+def register_aks_agent_resource_type():
+    register_resource_type(
+        "latest",
+        CUSTOM_MGMT_AKS,
+        None,
+    )
 
 
 class ContainerServiceCommandsLoader(AzCommandsLoader):
 
     def __init__(self, cli_ctx=None):
         from azure.cli.core.commands import CliCommandType
+        register_aks_agent_resource_type()
 
         aks_agent_custom = CliCommandType(operations_tmpl='azext_aks_agent.custom#{}')
         super().__init__(
@@ -44,14 +43,3 @@ def load_arguments(self, command):
 
 
 COMMAND_LOADER_CLS = ContainerServiceCommandsLoader
-
-
-# NOTE(mainred): holmesgpt leverages the environment variables to customize its behavior.
-def customize_holmesgpt():
-    os.environ[CONST_DISABLE_PROMETHEUS_TOOLSET_ENV_KEY] = "true"
-    os.environ[CONST_AGENT_CONFIG_PATH_DIR_ENV_KEY] = get_config_dir()
-    os.environ[CONST_AGENT_NAME_ENV_KEY] = CONST_AGENT_NAME
-    os.environ[CONST_PRIVACY_NOTICE_BANNER_ENV_KEY] = CONST_PRIVACY_NOTICE_BANNER
-
-
-customize_holmesgpt()

src/aks-agent/azext_aks_agent/_client_factory.py

@@ -0,0 +1,26 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+from azure.cli.core.commands.client_factory import get_mgmt_service_client
+from azure.cli.core.commands.parameters import get_resources_in_subscription
+from azure.cli.core.profiles import CustomResourceType, ResourceType
+from azure.mgmt.msi import ManagedServiceIdentityClient
+from knack.util import CLIError
+
+CUSTOM_MGMT_AKS = CustomResourceType('azext_aks_agent.vendored_sdks.azure_mgmt_containerservice.2025_10_01',
+                                     'ContainerServiceClient')
+
+# Note: cf_xxx, as the client_factory option value of a command group at command declaration, it should ignore
+# parameters other than cli_ctx; get_xxx_client is used as the client of other services in the command implementation,
+# and usually accepts subscription_id as a parameter to reconfigure the subscription when sending the request
+
+
+# container service clients
+def get_container_service_client(cli_ctx, subscription_id=None):
+    return get_mgmt_service_client(cli_ctx, CUSTOM_MGMT_AKS, subscription_id=subscription_id)
+
+
+def cf_managed_clusters(cli_ctx, *_):
+    return get_container_service_client(cli_ctx).managed_clusters

src/aks-agent/azext_aks_agent/_consts.py

@@ -30,6 +30,17 @@
 CONST_MCP_GITHUB_REPO = "Azure/aks-mcp"
 CONST_MCP_BINARY_DIR = "bin"
 
-# Color constants for terminal output
-HELP_COLOR = "cyan"  # same as AI_COLOR for now
-ERROR_COLOR = "red"
+# Kubernetes WebSocket exec protocol constants
+RESIZE_CHANNEL = 4  # WebSocket channel for terminal resize messages
+# WebSocket heartbeat configuration (matching kubectl client-go)
+# Based on kubernetes/client-go/tools/remotecommand/websocket.go#L59-L65
+# pingPeriod = 5 * time.Second
+# pingReadDeadline = (pingPeriod * 12) + (1 * time.Second)
+# The read deadline is calculated to allow up to 12 missed pings plus 1 second buffer
+# This provides tolerance for network delays while detecting actual connection failures
+HEARTBEAT_INTERVAL = 5.0                              # pingPeriod: 5 seconds between pings
+HEARTBEAT_TIMEOUT = (HEARTBEAT_INTERVAL * 12) + 1    # pingReadDeadline: 61 seconds total timeout
+
+AGENT_NAMESPACE = "kube-system"
+AGENT_LABEL_SELECTOR = "app.kubernetes.io/name=aks-agent"
+AKS_MCP_LABEL_SELECTOR = "app.kubernetes.io/name=aks-mcp"

src/aks-agent/azext_aks_agent/_help.py

@@ -5,17 +5,15 @@
 # --------------------------------------------------------------------------------------------
 
 # pylint: disable=too-many-lines
-
 from knack.help_files import helps
 
 helps[
     "aks agent"
 ] = """
     type: command
-    short-summary: Run AI assistant to analyze and troubleshoot Kubernetes clusters.
+    short-summary: Run AI assistant to analyze and troubleshoot Azure Kubernetes Service (AKS) clusters.
     long-summary: |-
       This command allows you to ask questions about your Azure Kubernetes cluster and get answers using AI models.
-      No need to manually set environment variables! All model and credential information can be configured interactively using `az aks agent-init` or via a config file.
     parameters:
         - name: --name -n
           type: string
@@ -33,12 +31,6 @@
             Each provider may require different environment variables and model naming conventions.
             For a full list of supported providers, model patterns, and required environment variables, see https://docs.litellm.ai/docs/providers.
             Note: For Azure OpenAI, it is recommended to set the deployment name as the model name until https://github.com/BerriAI/litellm/issues/13950 is resolved.
-        - name: --api-key
-          type: string
-          short-summary: API key to use for the LLM (if not given, uses environment variables AZURE_API_KEY, OPENAI_API_KEY). (Deprecated)
-        - name: --config-file
-          type: string
-          short-summary: Path to configuration file.
         - name: --max-steps
           type: int
           short-summary: Maximum number of steps the LLM can take to investigate the issue.
@@ -56,82 +48,35 @@
           short-summary: Refresh the toolsets status.
         - name: --status
           type: bool
-          short-summary: Show AKS agent configuration and status information.
-        - name: --aks-mcp
-          type: bool
-          short-summary: Enable AKS MCP integration for enhanced capabilities. Traditional mode is the default.
-
+          short-summary: Show AKS agent deployment status including helm release, deployments, and pod information.
     examples:
+        - name: Ask about pod issues in the cluster with OpenAI
+          text: |-
+            az aks agent "Why are my pods not starting?" --model gpt-4o --resource-group myResourceGroup --name myAKSCluster
         - name: Ask about pod issues in the cluster with last configured model
           text: |-
-            az aks agent "Why are my pods not starting?" --name MyManagedCluster --resource-group MyResourceGroup
-        - name: Ask about pod issues in the cluster with Azure OpenAI
+            az aks agent "Why are my pods not starting?" --resource-group myResourceGroup --name myAKSCluster
+          - name: Check AKS agent deployment status
           text: |-
-            az aks agent "Why are my pods not starting?" --name MyManagedCluster --resource-group MyResourceGroup --model azure/gpt-4.1
-        - name: Ask about pod issues in the cluster with OpenAI
+            az aks agent --status --resource-group myResourceGroup --name myAKSCluster
+        - name: Ask about pod issues in the cluster with Azure OpenAI
           text: |-
-            az aks agent "Why are my pods not starting?" --name MyManagedCluster --resource-group MyResourceGroup --model gpt-4o
-        - name: Run agent with config file
-          text: |
-            az aks agent "Check kubernetes pod resource usage" --config-file /path/to/custom.yaml --name MyManagedCluster --resource-group MyResourceGroup
-            Here is an example of config file:
-            ```json
-            llms:
-            - provider: azure
-              MODEL_NAME: gpt-4.1
-              AZURE_API_KEY: *******
-              AZURE_API_BASE: https://{azure-openai-service-name}.openai.azure.com/
-              AZURE_API_VERSION: 2025-04-01-preview
-            # define a list of mcp servers, mcp server can be defined
-            mcp_servers:
-              aks_mcp:
-                description: "The AKS-MCP is a Model Context Protocol (MCP) server that enables AI assistants to interact with Azure Kubernetes Service (AKS) clusters"
-                url: "http://localhost:8003/sse"
-
-            # try adding your own tools or toggle the built-in toolsets here
-            # e.g. query company-specific data, fetch logs from your existing observability tools, etc
-            # To check how to add a customized toolset, please refer to https://docs.robusta.dev/master/configuration/holmesgpt/custom_toolsets.html#custom-toolsets
-            # To find all built-in toolsets, please refer to https://docs.robusta.dev/master/configuration/holmesgpt/builtin_toolsets.html
-            toolsets:
-              # add a new json processor toolset
-              json_processor:
-                description: "A toolset for processing JSON data using jq"
-                prerequisites:
-                  - command: "jq --version"  # Ensure jq is installed
-                tools:
-                  - name: "process_json"
-                    description: "A tool that uses jq to process JSON input"
-                    command: "echo '{{ json_input }}' | jq '.'"  # Example jq command to format JSON
-              # disable a built-in toolsets
-              aks/core:
-                enabled: false
-              ```
+            az aks agent "Why are my pods not starting?" --model azure/gpt-4.1 --resource-group myResourceGroup --name myAKSCluster
         - name: Run in interactive mode without a question
-          text: az aks agent "Check the pod status in my cluster" --name MyManagedCluster --resource-group MyResourceGroup --model azure/gpt-4.1 --api-key "sk-xxx"
-        - name: Run in non-interactive batch mode
-          text: az aks agent "Diagnose networking issues" --no-interactive --max-steps 15 --model azure/gpt-4.1
-        - name: Show detailed tool output during analysis
-          text: az aks agent "Why is my service workload unavailable in namespace workload-ns?" --show-tool-output --model azure/gpt-4.1
-        - name: Use custom configuration file
-          text: az aks agent "Check kubernetes pod resource usage" --config-file /path/to/custom.yaml --model azure/gpt-4.1
-        - name: Run agent with no echo of the original question
-          text: az aks agent "What is the status of my cluster?" --no-echo-request --model azure/gpt-4.1
-        - name: Refresh toolsets to get the latest available tools
-          text: az aks agent "What is the status of my cluster?" --refresh-toolsets --model azure/gpt-4.1
-        - name: Show agent status (MCP readiness)
-          text: az aks agent --status
-        - name: Run in interactive mode without a question
-          text: az aks agent "Check the pod status in my cluster" --name MyManagedCluster --resource-group MyResourceGroup --model azure/my-gpt4.1-deployment --api-key "sk-xxx"
+          text: |-
+            az aks agent "Check the pod status in my cluster" --model azure/gpt-4.1 --resource-group myResourceGroup --name myAKSCluster
         - name: Run in non-interactive batch mode
-          text: az aks agent "Diagnose networking issues" --no-interactive --max-steps 15 --model azure/my-gpt4.1-deployment
+          text: |-
+            az aks agent "Diagnose networking issues" --no-interactive --max-steps 15 --model azure/gpt-4.1 --resource-group myResourceGroup --name myAKSCluster
         - name: Show detailed tool output during analysis
-          text: az aks agent "Why is my service workload unavailable in namespace workload-ns?" --show-tool-output --model azure/my-gpt4.1-deployment
-        - name: Use custom configuration file
-          text: az aks agent "Check kubernetes pod resource usage" --config-file /path/to/custom.yaml --model azure/my-gpt4.1-deployment
+          text: |-
+            az aks agent "Why is my service workload unavailable in namespace workload-ns?" --show-tool-output --model azure/gpt-4.1 --resource-group myResourceGroup --name myAKSCluster
         - name: Run agent with no echo of the original question
-          text: az aks agent "What is the status of my cluster?" --no-echo-request --model azure/my-gpt4.1-deployment
+          text: |-
+            az aks agent "What is the status of my cluster?" --no-echo-request --model azure/gpt-4.1 --resource-group myResourceGroup --name myAKSCluster
         - name: Refresh toolsets to get the latest available tools
-          text: az aks agent "What is the status of my cluster?" --refresh-toolsets --model azure/my-gpt4.1-deployment
+          text: |-
+            az aks agent "What is the status of my cluster?" --refresh-toolsets --model azure/gpt-4.1 --resource-group myResourceGroup --name myAKSCluster
 """
 
 helps[
@@ -142,7 +87,34 @@
     long-summary: |-
       This command interactively guides you to select an LLM provider and model, validates the connection, and saves the configuration for later use.
       You can run this command multiple times to add or update different model configurations.
+    parameters:
+        - name: --name -n
+          type: string
+          short-summary: Name of the managed cluster.
+        - name: --resource-group -g
+          type: string
+          short-summary: Name of the resource group.
+    examples:
+        - name: Initialize and deploy AKS agent to a cluster
+          text: |-
+            az aks agent-init --resource-group myResourceGroup --name myAKSCluster
+"""
+
+helps[
+    "aks agent-cleanup"
+] = """
+    type: command
+    short-summary: Cleanup and uninstall AKS agent from the cluster.
+    long-summary: |-
+      This command removes the AKS agent and deletes all associated resources from the cluster.
+    parameters:
+        - name: --name -n
+          type: string
+          short-summary: Name of the managed cluster.
+        - name: --resource-group -g
+          type: string
+          short-summary: Name of the resource group.
     examples:
-        - name: Initialize configuration for Azure OpenAI, OpenAI or other llms
-          text: az aks agent-init
+        - name: Cleanup and uninstall AKS agent from the cluster
+          text: az aks agent-cleanup --resource-group myResourceGroup --name myAKSCluster
 """

src/aks-agent/azext_aks_agent/_params.py

@@ -4,14 +4,6 @@
 # --------------------------------------------------------------------------------------------
 
 # pylint: disable=too-many-statements,too-many-lines
-import os.path
-
-from azext_aks_agent._consts import CONST_AGENT_CONFIG_FILE_NAME
-from azext_aks_agent._validators import validate_agent_config_file
-from azure.cli.core.api import get_config_dir
-from azure.cli.core.commands.parameters import get_three_state_flag
-
-
 def load_arguments(self, _):
     with self.argument_context("aks agent") as c:
         c.positional(
@@ -24,13 +16,11 @@ def load_arguments(self, _):
             "resource_group_name",
             options_list=["--resource-group", "-g"],
             help="Name of resource group.",
-            required=False,
         )
         c.argument(
-            "name",
+            "cluster_name",
             options_list=["--name", "-n"],
             help="Name of the managed cluster.",
-            required=False,
         )
         c.argument(
             "max_steps",
@@ -39,25 +29,12 @@ def load_arguments(self, _):
             required=False,
             help="Maximum number of steps the LLM can take to investigate the issue.",
         )
-        c.argument(
-            "config_file",
-            default=os.path.join(get_config_dir(), CONST_AGENT_CONFIG_FILE_NAME),
-            validator=validate_agent_config_file,
-            required=False,
-            help="Path to the config file.",
-        )
         c.argument(
             "model",
             help=" Specify the LLM provider and model or deployment to use for the AI assistant.",
             required=False,
             type=str,
         )
-        c.argument(
-            "api_key",
-            help="API key to use for the LLM (if not given, uses environment variables AZURE_API_KEY, OPENAI_API_KEY)",
-            required=False,
-            type=str,
-        )
         c.argument(
             "no_interactive",
             help="Disable interactive mode. When set, the agent will not prompt for input and will run in batch mode.",
@@ -84,14 +61,26 @@ def load_arguments(self, _):
             action="store_true",
             help="Show AKS agent configuration and status information.",
         )
+
+    with self.argument_context("aks agent-init") as c:
         c.argument(
-            "use_aks_mcp",
-            options_list=["--aks-mcp"],
-            default=False,
-            arg_type=get_three_state_flag(),
-            help=(
-                "Enable AKS MCP integration for enhanced capabilities. "
-                "Traditional mode is the default. Use --aks-mcp to enable MCP mode, or "
-                "--no-aks-mcp to explicitly disable it."
-            ),
+            "resource_group_name",
+            options_list=["--resource-group", "-g"],
+            help="Name of resource group.",
+        )
+        c.argument(
+            "cluster_name",
+            options_list=["--name", "-n"],
+            help="Name of the managed cluster.",
+        )
+    with self.argument_context("aks agent-cleanup") as c:
+        c.argument(
+            "resource_group_name",
+            options_list=["--resource-group", "-g"],
+            help="Name of resource group.",
+        )
+        c.argument(
+            "cluster_name",
+            options_list=["--name", "-n"],
+            help="Name of the managed cluster.",
         )