Azure Assistant MCP

Instructions

Implementation Reference

• src/azure_assistant_mcp/server.py:324-385 (handler)

  Main handler logic for the 'run-kql-template' tool: parses arguments, loads and parameterizes KQL template, determines tenant/scope, executes query, formats and returns results.
  if name == "run-kql-template": tenant_name = arguments.get("tenant_name") subs = arguments.get("subscription_ids") use_all = bool(arguments.get("use_all_subscriptions", True)) top = int(arguments.get("top", 100)) template_name = arguments.get("template_name") params: Dict[str, Any] = arguments.get("params") or {} if not template_name: return [types.TextContent(type="text", text="Error: template_name is required")] kql = load_kql_template(template_name) if not kql: return [types.TextContent(type="text", text=f"Error: template '{template_name}' not found in kql templates.")] # Simple {{key}} replacement try: for k, v in params.items(): kql = kql.replace(f"{{{{{k}}}}}", str(v)) except Exception: pass # Auto-guess tenant if not provided (best effort: look for tenant-like strings is skipped here) cred, default_subs = AZURE_CONFIG.get_credentials(tenant_name) mg = AZURE_CONFIG.get_management_group_id(tenant_name) if subs: pass elif use_all and mg: subs = [] # force MG usage below elif use_all: discovered = _enumerate_subscriptions_for_credential(cred) subs = discovered or default_subs else: subs = default_subs if not subs and not (use_all and mg): return [types.TextContent(type="text", text=( "Error: No subscriptions available; provide subscription_ids, set default_subscription_id in config, " "or ensure subscription discovery is possible (azure-mgmt-subscription installed and SP has access)." ))] if use_all and mg and not subs: result = execute_kql(cred, None, kql, top=top, management_groups=[mg]) else: result = execute_kql(cred, subs, kql, top=top) if result["status"] != "success": return [types.TextContent(type="text", text=f"ARG query failed: {result.get('error','unknown error')}")] rows = result["results"] scope_line = ( f"Scope: managementGroup={mg}" if (use_all and mg and not subs) else f"Subscriptions used: {len(subs)}" ) body = [ f"KQL Template: {template_name}", f"Rows: {result['result_count']}", f"Tenant: {tenant_name or AZURE_CONFIG.get_default_tenant().get('name')}", scope_line, "", _format_rows(rows), "", "KQL:", kql[:1000], ] return [types.TextContent(type="text", text="\n".join(body))]
• src/azure_assistant_mcp/server.py:137-152 (schema)

  Tool schema definition including input parameters validation and description.
  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"], }, ),
• src/azure_assistant_mcp/server.py:95-216 (registration)

  The list_tools function decorated with @server.list_tools() that returns the list of available tools, including 'run-kql-template'.
  @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/kql_loader.py:20-44 (helper)

  Helper function to load KQL templates from files in the kql directory, supporting .kql and .md formats.
  def load_kql_template(name: str) -> Optional[str]: """Load a KQL template by name. Supports either plain `.kql` files or Markdown with a fenced ```kql block. Returns None if no template is found or parse fails. """ base = _base_dir() # Try .kql first kql_path = base / f"{name}.kql" if kql_path.is_file(): try: return kql_path.read_text(encoding="utf-8").strip() except Exception: pass # Try .md with fenced code block md_path = base / f"{name}.md" if md_path.is_file(): try: text = md_path.read_text(encoding="utf-8") m = _FENCE_RE.search(text) if m: return m.group(1).strip() except Exception: pass return None
• src/azure_assistant_mcp/arg.py:7-49 (helper)

  Core helper function to execute the KQL query on Azure Resource Graph, handling scopes and error cases.
  def execute_kql( credential, subscriptions: Optional[List[str]], kql_query: str, top: int = 100, management_groups: Optional[List[str]] = None, ) -> Dict[str, Any]: """Execute a KQL query against Azure Resource Graph. Returns a dict with keys: status, results, result_count, warnings, query """ client = ResourceGraphClient(credential=credential) # Build typed request for reliability options = QueryRequestOptions(result_format="objectArray", top=top) # Either subscriptions or management_groups must be provided. request = QueryRequest( subscriptions=subscriptions, management_groups=management_groups, query=kql_query, options=options, ) try: response = client.resources(request) # response.data is a list[dict] when result_format=objectArray rows = list(response.data or []) return { "status": "success", "results": rows, "result_count": len(rows), "warnings": [], "query": kql_query, } except Exception as e: return { "status": "error", "error": str(e), "results": [], "result_count": 0, "warnings": [], "query": kql_query, }