Azure Omni-Tool MCP Server

import { z } from 'zod'; import { getService, listServiceTypes, isValidServiceType, ServiceType, StorageService, CosmosService, SearchService, KustoService, MonitorService, AppConfigService, KeyVaultService, PostgresService, ServiceResult, } from '../services/index.js'; import { createAuditContext } from '../lib/audit.js'; import { getOperatorInfo } from '../lib/config.js'; import { logger } from '../lib/logger.js'; export const AzureServiceSchema = z.object({ service: z.enum(['storage', 'cosmos', 'search', 'kusto', 'monitor', 'appconfig', 'keyvault', 'postgres']) .describe('Azure service type'), action: z.string().describe('Action: list, get, query, etc.'), params: z.record(z.string()).optional().describe('Action parameters'), }); export type AzureServiceInput = z.infer<typeof AzureServiceSchema>; export interface AzureServiceResponse { service: string; action: string; success: boolean; data?: unknown; count?: number; error?: string; correlation_id: string; operator?: { email?: string; name?: string }; } type ActionHandler = (params: Record<string, string>) => Promise<ServiceResult>; function getActionHandlers(serviceType: ServiceType): Map<string, ActionHandler> { const handlers = new Map<string, ActionHandler>(); switch (serviceType) { case 'storage': { const svc = getService<StorageService>('storage'); handlers.set('list', p => svc.list(p.resourceGroup)); handlers.set('listContainers', p => svc.listContainers(p.accountName)); handlers.set('listBlobs', p => svc.listBlobs(p.accountName, p.containerName)); handlers.set('getContainer', p => svc.getContainerProperties(p.accountName, p.containerName)); handlers.set('listTables', p => svc.listTables(p.accountName)); handlers.set('queryTable', p => svc.queryTable(p.accountName, p.tableName, p.filter)); break; } case 'cosmos': { const svc = getService<CosmosService>('cosmos'); handlers.set('list', p => svc.list(p.resourceGroup)); handlers.set('listDatabases', p => svc.listDatabases(p.accountName, p.resourceGroup)); handlers.set('listContainers', p => svc.listContainers(p.accountName, p.resourceGroup, p.databaseName)); handlers.set('query', p => svc.queryContainer(p.accountName, p.resourceGroup, p.databaseName, p.containerName, p.query)); handlers.set('getContainer', p => svc.getContainer(p.accountName, p.resourceGroup, p.databaseName, p.containerName)); break; } case 'search': { const svc = getService<SearchService>('search'); handlers.set('list', p => svc.list(p.resourceGroup)); handlers.set('listIndexes', p => svc.listIndexes(p.serviceName, p.resourceGroup)); handlers.set('getIndex', p => svc.getIndex(p.serviceName, p.resourceGroup, p.indexName)); handlers.set('query', p => { const top = p.top ? parseInt(p.top, 10) : undefined; if (top !== undefined && isNaN(top)) return Promise.reject(new Error(`Invalid 'top': '${p.top}' is not a number.`)); return svc.queryIndex(p.serviceName, p.resourceGroup, p.indexName, p.searchText, top); }); handlers.set('getService', p => svc.getServiceDetails(p.serviceName, p.resourceGroup)); break; } case 'kusto': { const svc = getService<KustoService>('kusto'); handlers.set('list', p => svc.list(p.resourceGroup)); handlers.set('listDatabases', p => svc.listDatabases(p.clusterName, p.resourceGroup)); handlers.set('listTables', p => svc.listTables(p.clusterName, p.resourceGroup, p.databaseName)); handlers.set('getSchema', p => svc.getTableSchema(p.clusterName, p.resourceGroup, p.databaseName, p.tableName)); handlers.set('sample', p => { const count = parseInt(p.count ?? '10', 10); if (isNaN(count) || count < 1 || !Number.isInteger(count)) return Promise.reject(new Error(`Invalid count: '${p.count}'. Must be positive integer.`)); return svc.sampleTable(p.clusterName, p.resourceGroup, p.databaseName, p.tableName, count); }); handlers.set('query', p => svc.runKql(p.clusterName, p.resourceGroup, p.databaseName, p.query)); break; } case 'monitor': { const svc = getService<MonitorService>('monitor'); handlers.set('list', p => svc.list(p.resourceGroup)); handlers.set('getWorkspace', p => svc.getWorkspace(p.workspaceName, p.resourceGroup)); handlers.set('listTables', p => svc.listTables(p.workspaceName, p.resourceGroup)); handlers.set('query', p => svc.queryLogs(p.workspaceId, p.query, p.timespan)); handlers.set('listMetrics', p => svc.listMetricDefinitions(p.resourceId)); handlers.set('getMetrics', p => svc.getMetrics(p.resourceId, p.metrics, p.interval)); break; } case 'appconfig': { const svc = getService<AppConfigService>('appconfig'); handlers.set('list', p => svc.list(p.resourceGroup)); handlers.set('getStore', p => svc.getStore(p.storeName, p.resourceGroup)); handlers.set('listKeyValues', p => svc.listKeyValues(p.storeName, p.label)); handlers.set('getKeyValue', p => svc.getKeyValue(p.storeName, p.key, p.label)); handlers.set('setKeyValue', p => svc.setKeyValue(p.storeName, p.key, p.value, p.label)); handlers.set('lock', p => svc.lockKeyValue(p.storeName, p.key, p.label)); handlers.set('unlock', p => svc.unlockKeyValue(p.storeName, p.key, p.label)); break; } case 'keyvault': { const svc = getService<KeyVaultService>('keyvault'); handlers.set('list', p => svc.list(p.resourceGroup)); handlers.set('getVault', p => svc.getVault(p.vaultName, p.resourceGroup)); handlers.set('listKeys', p => svc.listKeys(p.vaultName)); handlers.set('getKey', p => svc.getKey(p.vaultName, p.keyName)); handlers.set('createKey', p => svc.createKey(p.vaultName, p.keyName, p.keyType || 'RSA')); handlers.set('listSecrets', p => svc.listSecrets(p.vaultName)); handlers.set('getSecret', p => svc.getSecret(p.vaultName, p.secretName)); handlers.set('listCertificates', p => svc.listCertificates(p.vaultName)); break; } case 'postgres': { const svc = getService<PostgresService>('postgres'); handlers.set('list', p => svc.list(p.resourceGroup)); handlers.set('getServer', p => svc.getServer(p.serverName, p.resourceGroup)); handlers.set('listDatabases', p => svc.listDatabases(p.serverName, p.resourceGroup)); handlers.set('listParameters', p => svc.listParameters(p.serverName, p.resourceGroup)); handlers.set('getParameter', p => svc.getServerParameter(p.serverName, p.resourceGroup, p.paramName)); handlers.set('listTables', p => svc.listTables(p.serverName, p.resourceGroup, p.databaseName)); handlers.set('getTableSchema', p => svc.getTableSchema(p.serverName, p.resourceGroup, p.databaseName, p.tableName)); handlers.set('query', p => svc.executeQuery(p.serverName, p.resourceGroup, p.databaseName, p.query)); break; } } return handlers; } export async function handleAzureService(input: AzureServiceInput): Promise<AzureServiceResponse> { const { service, action, params = {} } = input; const operator = getOperatorInfo(); const audit = createAuditContext(`${service}:${action}`, 'low', 'query'); logger.debug('Service action', { service, action, params }); if (!isValidServiceType(service)) { return { service, action, success: false, error: `Unknown service. Available: ${listServiceTypes().join(', ')}`, correlation_id: audit.correlationId, operator, }; } const handlers = getActionHandlers(service); const handler = handlers.get(action); if (!handler) { return { service, action, success: false, error: `Unknown action "${action}" for ${service}. Available: ${Array.from(handlers.keys()).join(', ')}`, correlation_id: audit.correlationId, operator, }; } try { const result = await handler(params); if (result.success) { await audit.logSuccess(); logger.info('Service action succeeded', { service, action, count: result.count }); } else { await audit.logFailure(result.error || 'Unknown'); logger.warn('Service action failed', { service, action, error: result.error }); } return { service, action, success: result.success, data: result.data, count: result.count, error: result.error, correlation_id: audit.correlationId, operator, }; } catch (err) { const error = err instanceof Error ? err.message : String(err); await audit.logFailure(error); logger.error('Service action error', err instanceof Error ? err : new Error(error)); return { service, action, success: false, error, correlation_id: audit.correlationId, operator, }; } } export const azureServiceTool = { name: 'azure_service', description: `Interact with specific Azure services. SERVICES: storage, cosmos, search, kusto, monitor, appconfig, keyvault, postgres STORAGE actions: list, listContainers, listBlobs, getContainer, listTables, queryTable COSMOS actions: list, listDatabases, listContainers, query, getContainer SEARCH actions: list, listIndexes, getIndex, query, getService KUSTO actions: list, listDatabases, listTables, getSchema, sample, query MONITOR actions: list, getWorkspace, listTables, query, listMetrics, getMetrics APPCONFIG actions: list, getStore, listKeyValues, getKeyValue, setKeyValue, lock, unlock KEYVAULT actions: list, getVault, listKeys, getKey, createKey, listSecrets, getSecret, listCertificates POSTGRES actions: list, getServer, listDatabases, listParameters, getParameter, listTables, getTableSchema, query Pass required params for each action (e.g., accountName, resourceGroup, query).`, inputSchema: { type: 'object', properties: { service: { type: 'string', enum: ['storage', 'cosmos', 'search', 'kusto', 'monitor', 'appconfig', 'keyvault', 'postgres'], description: 'Azure service type' }, action: { type: 'string', description: 'Action to perform' }, params: { type: 'object', description: 'Action parameters (varies by service/action)', additionalProperties: { type: 'string' } }, }, required: ['service', 'action'], }, };