Azure Assistant MCP

Instructions

Implementation Reference

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

  Main execution logic for the 'run-kql-template' tool within the call_tool dispatcher. Handles argument parsing, template loading and parameterization, scope resolution, query execution, and result formatting.
  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 (registration)

  Registration of the 'run-kql-template' tool in the list_tools() handler, including name, description, and full input schema.
  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/kql_loader.py:20-44 (helper)

  Helper function specifically used by 'run-kql-template' to load KQL templates from the 'kql/' directory, supporting .kql files or fenced blocks in .md files.
  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)

  Shared helper function called by 'run-kql-template' (and other tools) to execute the KQL query against Azure Resource Graph service.
  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, }