Azure Assistant MCP

Instructions

Implementation Reference

• src/azure_assistant_mcp/server.py:493-540 (handler)

  The handler function within call_tool that implements the core logic for the 'vm-count-by-tenant' tool. It iterates over configured tenants (optionally filtered), determines the query scope (subscriptions or management group), executes the VM count KQL query, aggregates counts, and formats the response.
  if name == "vm-count-by-tenant": # Optional: restrict to certain tenants wanted = set(arguments.get("tenant_names", []) or []) try_all = bool(arguments.get("use_all_subscriptions", True)) tenants = AZURE_CONFIG.get_tenants() if wanted: tenants = [t for t in tenants if t.get("name") in wanted] if not tenants: return [types.TextContent(type="text", text="No matching tenants found.")] rows: List[str] = [] total = 0 for t in tenants: tname = t.get("name") or t.get("id") cred, default_subs = AZURE_CONFIG.get_credentials(tname) mg = AZURE_CONFIG.get_management_group_id(tname) subs = list(default_subs) if try_all and mg: subs = [] # force MG usage below elif try_all: discovered = _enumerate_subscriptions_for_credential(cred) if discovered: subs = discovered if not subs: rows.append(f"- {tname}: 0 (no subscriptions)") continue kql = vm_count_kql() if try_all and mg and not subs: result = execute_kql(cred, None, kql, top=1, management_groups=[mg]) else: result = execute_kql(cred, subs, kql, top=1) if result["status"] != "success" or not result.get("results"): rows.append(f"- {tname}: error ({result.get('error','query failed')})") continue count = int(result["results"][0].get("VMCount", 0)) total += count rows.append(f"- {tname}: {count}") body = [ "VM counts by tenant:", *rows, "", f"Total VMs across tenants: {total}", ] return [types.TextContent(type="text", text="\n".join(body))]
• src/azure_assistant_mcp/server.py:153-164 (schema)

  The input schema definition and tool description for 'vm-count-by-tenant', provided in the list_tools() response for tool discovery.
  types.Tool( name="vm-count-by-tenant", description="Count virtual machines per configured tenant (uses all subscriptions when available).", inputSchema={ "type": "object", "properties": { "tenant_names": {"type": "array", "items": {"type": "string"}, "description": "Optional subset of tenant names to include"}, "use_all_subscriptions": {"type": "boolean", "description": "Try to include all subscriptions in each tenant (fallbacks to default)", "default": True} }, "required": [], }, ),
• src/azure_assistant_mcp/server.py:95-216 (registration)

  The list_tools() function decorated with @server.list_tools(), which registers and returns the list of available tools including 'vm-count-by-tenant'.
  @server.list_tools() async def list_tools() -> List[types.Tool]: tools: List[types.Tool] = [ types.Tool( name="ask-azure", description="Answer a question by generating and running an Azure Resource Graph KQL query.", inputSchema={ "type": "object", "properties": { "question": {"type": "string", "description": "Your natural-language question about Azure resources"}, "tenant_name": {"type": "string", "description": "Optional configured tenant name"}, "subscription_ids": {"type": "array", "items": {"type": "string"}, "description": "Optional explicit subscription IDs"}, "use_all_subscriptions": {"type": "boolean", "description": "If no subscriptions are provided, attempt to auto-discover all accessible subscriptions (default: true)", "default": True}, "auto_execute": {"type": "boolean", "description": "Execute the generated KQL automatically (default: true)", "default": True} }, "required": ["question"], }, ), types.Tool( name="list-tenants", description="List configured tenants from azure-config.json, including optional management group and default subscription info.", inputSchema={ "type": "object", "properties": {}, "required": [], }, ), types.Tool( name="run-arg-kql", description="Run a provided KQL query against Azure Resource Graph.", inputSchema={ "type": "object", "properties": { "kql_query": {"type": "string", "description": "KQL to execute (must reference valid ARG tables)"}, "tenant_name": {"type": "string", "description": "Optional configured tenant name"}, "subscription_ids": {"type": "array", "items": {"type": "string"}, "description": "Optional explicit subscription IDs"}, "use_all_subscriptions": {"type": "boolean", "description": "If no subscriptions are provided, attempt to auto-discover all accessible subscriptions (default: true)", "default": True}, "top": {"type": "integer", "description": "Max rows to return (default: 100)", "default": 100} }, "required": ["kql_query"], }, ), types.Tool( name="run-kql-template", description="Execute a KQL template from kql/ by name, with optional {{param}} replacements.", inputSchema={ "type": "object", "properties": { "template_name": {"type": "string", "description": "Template filename without extension (looks for .md or .kql in kql folder)"}, "params": {"type": "object", "description": "Optional key/value replacements for {{key}} placeholders in the template"}, "tenant_name": {"type": "string", "description": "Optional configured tenant name"}, "subscription_ids": {"type": "array", "items": {"type": "string"}, "description": "Optional explicit subscription IDs"}, "use_all_subscriptions": {"type": "boolean", "description": "If no subscriptions are provided, attempt to auto-discover all accessible subscriptions (default: true)", "default": True}, "top": {"type": "integer", "description": "Max rows to return (default: 100)", "default": 100} }, "required": ["template_name"], }, ), types.Tool( name="vm-count-by-tenant", description="Count virtual machines per configured tenant (uses all subscriptions when available).", inputSchema={ "type": "object", "properties": { "tenant_names": {"type": "array", "items": {"type": "string"}, "description": "Optional subset of tenant names to include"}, "use_all_subscriptions": {"type": "boolean", "description": "Try to include all subscriptions in each tenant (fallbacks to default)", "default": True} }, "required": [], }, ), types.Tool( name="list-subscriptions", description="List subscriptions accessible to the configured service principal for a tenant.", inputSchema={ "type": "object", "properties": { "tenant_name": {"type": "string", "description": "Optional configured tenant name"} }, "required": [], }, ), types.Tool( name="arg-tables", description="Overview of common Azure Resource Graph tables, their purpose, and typical use cases.", inputSchema={ "type": "object", "properties": {}, "required": [], }, ), types.Tool( name="arg-examples", description="Sample KQL snippets for common scenarios across ARG tables.", inputSchema={ "type": "object", "properties": { "topic": {"type": "string", "description": "Optional topic filter, e.g., 'subscriptions', 'policy', 'advisor', 'health', 'changes', 'resourcegroups'"} }, "required": [], }, ), ] # Add diagnostics when debug is enabled in azure-config.json try: if AZURE_CONFIG and AZURE_CONFIG.is_debug_enabled(): tools.append( types.Tool( name="diagnostics", description="Diagnostics for scoping: shows resolved tenant, config path, ARM enumeration vs ARG MG coverage.", inputSchema={ "type": "object", "properties": { "tenant_name": {"type": "string", "description": "Optional configured tenant name"} }, "required": [], }, ) ) except Exception: pass return tools
• src/azure_assistant_mcp/queries.py:39-40 (helper)

  Helper function that loads and returns the KQL template for counting virtual machines ('vm_count'). Used by the handler.
  def vm_count_kql() -> str: return _tmpl("vm_count")