diff --git a/infra/mcp-server.bicep b/infra/mcp-server.bicep index 9676596e..ca378728 100644 --- a/infra/mcp-server.bicep +++ b/infra/mcp-server.bicep @@ -32,6 +32,9 @@ param imageTag string = 'latest' @description('Entra ID client secret for OAuth token exchange') param entraClientSecret string = '' +@description('Enable structured tool-call logging to stderr (queryable via Container App Logs)') +param toolLoggingEnabled bool = false + // --- Azure Container Registry --- resource acr 'Microsoft.ContainerRegistry/registries@2023-07-01' = { @@ -120,6 +123,7 @@ resource containerApp 'Microsoft.App/containerApps@2024-03-01' = { { name: 'AUTOPILOT_API_URL', value: apiUrl } { name: 'PORT', value: '8080' } { name: 'AUTOPILOT_ENTRA_CLIENT_SECRET', secretRef: 'entra-client-secret' } + { name: 'MCP_TOOL_LOGGING', value: toolLoggingEnabled ? 'true' : 'false' } ] probes: [ { diff --git a/src/Backend/AutopilotMonitor.Functions/Middleware/AuthenticationMiddleware.cs b/src/Backend/AutopilotMonitor.Functions/Middleware/AuthenticationMiddleware.cs index 7223e63b..c30e30a1 100644 --- a/src/Backend/AutopilotMonitor.Functions/Middleware/AuthenticationMiddleware.cs +++ b/src/Backend/AutopilotMonitor.Functions/Middleware/AuthenticationMiddleware.cs @@ -250,6 +250,9 @@ public class AuthenticationMiddleware : IFunctionsWorkerMiddleware var upn = principal.GetUserPrincipalName() ?? "unknown"; var tid = principal.GetTenantId() ?? ""; var normalizedEndpoint = EndpointNormalizer.Normalize(requestPath); + var mcpToolName = httpContext.Request.Headers["X-MCP-Tool-Name"].FirstOrDefault(); + if (!string.IsNullOrEmpty(mcpToolName)) + normalizedEndpoint = $"{mcpToolName}:{normalizedEndpoint}"; var repo = _userUsageRepo; var logger = _logger; _ = Task.Run(async () => diff --git a/src/Backend/AutopilotMonitor.Functions/Middleware/RequestTelemetryMiddleware.cs b/src/Backend/AutopilotMonitor.Functions/Middleware/RequestTelemetryMiddleware.cs index 76aab969..590a8c43 100644 --- a/src/Backend/AutopilotMonitor.Functions/Middleware/RequestTelemetryMiddleware.cs +++ b/src/Backend/AutopilotMonitor.Functions/Middleware/RequestTelemetryMiddleware.cs @@ -78,6 +78,14 @@ public class RequestTelemetryMiddleware : IFunctionsWorkerMiddleware requestTelemetry.Properties["HttpMethod"] = method; requestTelemetry.Properties["HttpPath"] = httpContext.Request.Path.Value ?? ""; + var clientSource = httpContext.Request.Headers["X-Client-Source"].FirstOrDefault(); + if (!string.IsNullOrEmpty(clientSource)) + requestTelemetry.Properties["ClientSource"] = clientSource; + + var mcpToolName = httpContext.Request.Headers["X-MCP-Tool-Name"].FirstOrDefault(); + if (!string.IsNullOrEmpty(mcpToolName)) + requestTelemetry.Properties["McpToolName"] = mcpToolName; + if (context.Items.TryGetValue("CorrelationId", out var corrId) && corrId is string correlationId) requestTelemetry.Properties["CorrelationId"] = correlationId; diff --git a/src/McpServer/autopilot-monitor-mcp/src/client.ts b/src/McpServer/autopilot-monitor-mcp/src/client.ts index fa445919..cca1e7e2 100644 --- a/src/McpServer/autopilot-monitor-mcp/src/client.ts +++ b/src/McpServer/autopilot-monitor-mcp/src/client.ts @@ -50,6 +50,20 @@ export function getCurrentToken(): string | undefined { return tokenStore.getStore(); } +/** + * Per-request tool name store. Carries the MCP tool name through the async + * context so apiFetch can send it as X-MCP-Tool-Name header to the backend. + */ +const toolNameStore = new AsyncLocalStorage(); + +export function runWithToolName(toolName: string, fn: () => T | Promise): T | Promise { + return toolNameStore.run(toolName, fn); +} + +export function getCurrentToolName(): string | undefined { + return toolNameStore.getStore(); +} + async function apiFetch(path: string, options: RequestInit = {}): Promise { const token = getCurrentToken(); if (!token) { @@ -64,6 +78,9 @@ async function apiFetch(path: string, options: RequestInit = {}): Promise) ?? {}), }; + const toolName = getCurrentToolName(); + if (toolName) headers['X-MCP-Tool-Name'] = toolName; + // Apply timeout to prevent hanging on unresponsive backend const signal = options.signal ?? AbortSignal.timeout(API_TIMEOUT_MS); diff --git a/src/McpServer/autopilot-monitor-mcp/src/telemetry.ts b/src/McpServer/autopilot-monitor-mcp/src/telemetry.ts new file mode 100644 index 00000000..c9c10c1c --- /dev/null +++ b/src/McpServer/autopilot-monitor-mcp/src/telemetry.ts @@ -0,0 +1,33 @@ +import { runWithToolName } from './client.js'; + +const toolLoggingEnabled = process.env.MCP_TOOL_LOGGING === 'true'; + +/** + * Wraps an MCP tool handler to: + * 1. Always: propagate the tool name via AsyncLocalStorage so apiFetch sends + * the X-MCP-Tool-Name header to the backend (tracked in App Insights). + * 2. Optionally (MCP_TOOL_LOGGING=true): emit structured JSON to stderr, + * queryable via Container App Logs in Azure Monitor. + */ +export async function withToolTelemetry(toolName: string, fn: () => T | Promise): Promise { + if (!toolLoggingEnabled) { + return runWithToolName(toolName, fn) as Promise; + } + + const start = Date.now(); + let isError = false; + try { + return await runWithToolName(toolName, fn); + } catch (err) { + isError = true; + throw err; + } finally { + console.error(JSON.stringify({ + type: 'tool_call', + tool: toolName, + durationMs: Date.now() - start, + isError, + timestamp: new Date().toISOString(), + })); + } +} diff --git a/src/McpServer/autopilot-monitor-mcp/src/tools/admin.ts b/src/McpServer/autopilot-monitor-mcp/src/tools/admin.ts index 2a84c86d..8ed291ec 100644 --- a/src/McpServer/autopilot-monitor-mcp/src/tools/admin.ts +++ b/src/McpServer/autopilot-monitor-mcp/src/tools/admin.ts @@ -1,6 +1,7 @@ import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { z } from 'zod'; import { apiFetch, buildQuery } from '../client.js'; +import { withToolTelemetry } from '../telemetry.js'; import { READ_ONLY, READ_ONLY_OPEN } from './shared.js'; import { toolError } from './error-handler.js'; @@ -19,7 +20,7 @@ export function registerAdminTools(server: McpServer): void { daily: z.boolean().optional().default(false).describe('Return daily aggregated summary instead of per-endpoint breakdown'), }, READ_ONLY, - async (args) => { + async (args) => withToolTelemetry('get_api_usage', async () => { try { let data: unknown; const params: Record = { tenantId: args.tenantId, dateFrom: args.dateFrom, dateTo: args.dateTo }; @@ -34,7 +35,7 @@ export function registerAdminTools(server: McpServer): void { } catch (error: unknown) { return toolError('get_api_usage', args, error); } - } + }) ); // Tool 12: get_geographic_metrics @@ -50,7 +51,7 @@ export function registerAdminTools(server: McpServer): void { .describe('Geographic grouping level (default: "city")'), }, READ_ONLY, - async (args) => { + async (args) => withToolTelemetry('get_geographic_metrics', async () => { try { const { tenantId, ...rest } = args; const params: Record = { ...rest }; @@ -61,7 +62,7 @@ export function registerAdminTools(server: McpServer): void { } catch (error: unknown) { return toolError('get_geographic_metrics', args, error); } - } + }) ); // Tool 13: get_geographic_sessions @@ -79,7 +80,7 @@ export function registerAdminTools(server: McpServer): void { days: z.coerce.number().optional().default(30).describe('Time range in days (default: 30)'), }, READ_ONLY, - async (args) => { + async (args) => withToolTelemetry('get_geographic_sessions', async () => { try { const { tenantId, country, region, city, ...rest } = args; const params: Record = { ...rest }; @@ -95,7 +96,7 @@ export function registerAdminTools(server: McpServer): void { } catch (error: unknown) { return toolError('get_geographic_sessions', args, error); } - } + }) ); // Tool 14: get_platform_metrics @@ -106,7 +107,7 @@ export function registerAdminTools(server: McpServer): void { 'top sessions by CPU/memory, and per-agent-version breakdown. Global Admin only.', {}, READ_ONLY, - async (args) => { + async (args) => withToolTelemetry('get_platform_metrics', async () => { try { type SessionMetric = { sessionId: string; tenantId: string; deviceName?: string; model?: string; status?: string; @@ -172,7 +173,7 @@ export function registerAdminTools(server: McpServer): void { } catch (error: unknown) { return toolError('get_platform_metrics', args, error); } - } + }) ); // Tool 15: get_usage_metrics @@ -184,7 +185,7 @@ export function registerAdminTools(server: McpServer): void { tenantId: z.string().optional().describe('Tenant ID for tenant-specific metrics. Omit for platform-wide overview (Global Admin only).'), }, READ_ONLY, - async (args) => { + async (args) => withToolTelemetry('get_usage_metrics', async () => { try { const { tenantId } = args; if (tenantId) { @@ -196,7 +197,7 @@ export function registerAdminTools(server: McpServer): void { } catch (error: unknown) { return toolError('get_usage_metrics', args, error); } - } + }) ); // Tool 16: get_audit_logs @@ -208,7 +209,7 @@ export function registerAdminTools(server: McpServer): void { tenantId: z.string().optional().describe('Tenant ID for tenant-scoped audit log. Omit for cross-tenant view (Global Admin only).'), }, READ_ONLY, - async (args) => { + async (args) => withToolTelemetry('get_audit_logs', async () => { try { const { tenantId } = args; const endpoint = tenantId @@ -219,7 +220,7 @@ export function registerAdminTools(server: McpServer): void { } catch (error: unknown) { return toolError('get_audit_logs', args, error); } - } + }) ); // Tool 17: get_ops_events @@ -233,7 +234,7 @@ export function registerAdminTools(server: McpServer): void { maxResults: z.coerce.number().min(1).max(500).optional().default(200).describe('Max events to return (default: 200)'), }, READ_ONLY, - async (args) => { + async (args) => withToolTelemetry('get_ops_events', async () => { try { const params: Record = {}; if (args.category) params.category = args.category; @@ -243,7 +244,7 @@ export function registerAdminTools(server: McpServer): void { } catch (error: unknown) { return toolError('get_ops_events', args, error); } - } + }) ); // Tool 18: list_session_reports @@ -253,14 +254,14 @@ export function registerAdminTools(server: McpServer): void { 'Global Admin only — returns reports across all tenants.', {}, READ_ONLY, - async (args) => { + async (args) => withToolTelemetry('list_session_reports', async () => { try { const data = await apiFetch('/api/global/session-reports'); return { content: [{ type: 'text' as const, text: JSON.stringify(data, null, 2) }] }; } catch (error: unknown) { return toolError('list_session_reports', args, error); } - } + }) ); // ── Raw Data Tools ──────────────────────────────────────────────────── @@ -284,7 +285,7 @@ export function registerAdminTools(server: McpServer): void { limit: z.coerce.number().min(1).max(500).optional().default(100), }, READ_ONLY, - async (args) => { + async (args) => withToolTelemetry('query_raw_events', async () => { try { const { tenantId, ...rest } = args; const params: Record = { ...rest }; @@ -295,7 +296,7 @@ export function registerAdminTools(server: McpServer): void { } catch (error: unknown) { return toolError('query_raw_events', args, error); } - } + }) ); // Tool 19: query_raw_sessions @@ -314,7 +315,7 @@ export function registerAdminTools(server: McpServer): void { limit: z.coerce.number().min(1).max(200).optional().default(50), }, READ_ONLY, - async (args) => { + async (args) => withToolTelemetry('query_raw_sessions', async () => { try { const { tenantId, ...rest } = args; const params: Record = { ...rest }; @@ -325,7 +326,7 @@ export function registerAdminTools(server: McpServer): void { } catch (error: unknown) { return toolError('query_raw_sessions', args, error); } - } + }) ); // ── Admin Diagnostic Tools ──────────────────────────────────────────── @@ -340,14 +341,14 @@ export function registerAdminTools(server: McpServer): void { 'List all available Azure Table Storage tables that can be queried via query_table. Global Admin only.', {}, READ_ONLY, - async (args) => { + async (args) => withToolTelemetry('list_tables', async () => { try { const data = await apiFetch('/api/global/raw/tables'); return { content: [{ type: 'text' as const, text: JSON.stringify(data, null, 2) }] }; } catch (error: unknown) { return toolError('list_tables', args, error); } - } + }) ); // Tool 21: query_table @@ -363,7 +364,7 @@ export function registerAdminTools(server: McpServer): void { limit: z.coerce.number().min(1).max(500).optional().default(100), }, READ_ONLY, - async (args) => { + async (args) => withToolTelemetry('query_table', async () => { try { const { tableName, ...rest } = args; const data = await apiFetch(`/api/global/raw/tables/${encodeURIComponent(tableName)}${buildQuery(rest)}`); @@ -371,7 +372,7 @@ export function registerAdminTools(server: McpServer): void { } catch (error: unknown) { return toolError('query_table', args, error); } - } + }) ); // Tool 22: query_backend_logs @@ -384,7 +385,7 @@ export function registerAdminTools(server: McpServer): void { timespan: z.string().optional().default('PT1H').describe('ISO 8601 duration (default: PT1H = last 1 hour). Examples: PT30M, PT6H, P1D'), }, READ_ONLY_OPEN, - async (args) => { + async (args) => withToolTelemetry('query_backend_logs', async () => { try { const data = await apiFetch('/api/global/raw/logs', { method: 'POST', @@ -394,7 +395,7 @@ export function registerAdminTools(server: McpServer): void { } catch (error: unknown) { return toolError('query_backend_logs', args, error); } - } + }) ); // Tool: get_rule_stats @@ -411,7 +412,7 @@ export function registerAdminTools(server: McpServer): void { endDate: z.string().optional().describe('End date (YYYY-MM-DD). Defaults to today.'), }, READ_ONLY, - async (args) => { + async (args) => withToolTelemetry('get_rule_stats', async () => { try { const params = buildQuery({ startDate: args.startDate, @@ -426,6 +427,6 @@ export function registerAdminTools(server: McpServer): void { } catch (error: unknown) { return toolError('get_rule_stats', args, error); } - } + }) ); } diff --git a/src/McpServer/autopilot-monitor-mcp/src/tools/search.ts b/src/McpServer/autopilot-monitor-mcp/src/tools/search.ts index 58d4fe74..66bd2b6f 100644 --- a/src/McpServer/autopilot-monitor-mcp/src/tools/search.ts +++ b/src/McpServer/autopilot-monitor-mcp/src/tools/search.ts @@ -1,6 +1,7 @@ import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { z } from 'zod'; import { apiFetch, buildQuery } from '../client.js'; +import { withToolTelemetry } from '../telemetry.js'; import type { SearchProvider } from '../search-provider.js'; import { READ_ONLY } from './shared.js'; import { toolError } from './error-handler.js'; @@ -11,18 +12,74 @@ type EventEntry = { eventType?: string; severity?: string; source?: string; message?: string; timestamp?: string; phase?: string; data?: Record; _sessionId?: string; + sessionId?: string; }; -/** Fetch events for a single session or across recent failed sessions. */ +/** All known event type strings — used for index-based pre-filtering. */ +const KNOWN_EVENT_TYPES = [ + 'phase_transition', 'esp_state_change', 'completion_check', + 'enrollment_complete', 'enrollment_failed', 'desktop_arrived', + 'app_install_started', 'app_install_completed', 'app_install_failed', + 'app_download_started', 'app_install_skipped', + 'network_state_change', 'network_connectivity_check', + 'os_info', 'hardware_spec', 'tpm_status', 'autopilot_profile', + 'secureboot_status', 'bitlocker_status', 'network_adapters', + 'network_interface_info', 'aad_join_status', 'enrollment_type_detected', + 'error_detected', 'software_inventory_analysis', 'vulnerability_report', + 'performance_snapshot', 'log_entry', 'gather_result', + 'script_completed', 'script_failed', 'ime_agent_version', + 'do_telemetry', 'download_progress', 'agent_started', + 'esp_phase_changed', 'shadow_discrepancy', +]; + +/** Match query keywords against known event types using prefix-aware matching. */ +function extractEventTypeCandidates(keywords: string[]): string[] { + return KNOWN_EVENT_TYPES.filter((et) => + keywords.some((kw) => prefixAwareMatch(et, kw)), + ); +} + +/** Fetch events for a single session or across sessions using index-based search. */ async function fetchSessionEvents( sessionId: string | undefined, tenantId: string | undefined, -): Promise<{ events: EventEntry[]; sessionIds: string[] }> { + queryKeywords?: string[], +): Promise<{ events: EventEntry[]; sessionIds: string[]; searchMethod: string }> { + // Single-session: direct fetch (unchanged) if (sessionId) { const q = buildQuery({ tenantId } as Record); const data = await apiFetch(`/api/sessions/${sessionId}/events${q}`) as { events?: EventEntry[] }; - return { events: data?.events ?? [], sessionIds: [sessionId] }; + return { events: data?.events ?? [], sessionIds: [sessionId], searchMethod: 'direct-session' }; } + + // Multi-session: try index-based search first if we have keyword candidates + if (queryKeywords && queryKeywords.length > 0) { + const candidates = extractEventTypeCandidates(queryKeywords); + if (candidates.length > 0 && candidates.length <= 10) { + try { + const params: Record = { + eventTypes: candidates.join(','), + limit: 200, + sessionLimit: 10, + }; + if (tenantId) params.tenantId = tenantId; + const basePath = tenantId ? '/api/raw/events/search' : '/api/global/raw/events/search'; + const data = await apiFetch(`${basePath}${buildQuery(params)}`) as { + events?: EventEntry[]; + }; + const events = (data?.events ?? []).map((e) => ({ + ...e, + _sessionId: e.sessionId ?? e._sessionId, + })); + const sessionIds = [...new Set(events.map((e) => e._sessionId).filter(Boolean) as string[])]; + return { events, sessionIds, searchMethod: 'index-based' }; + } catch { + // Fall through to legacy approach + } + } + } + + // Fallback: search recent failed sessions + fetch all events (legacy N+1 path) const searchParams: Record = { status: 'Failed', limit: 5 }; if (tenantId) searchParams.tenantId = tenantId; const searchQ = buildQuery(searchParams); @@ -38,9 +95,9 @@ async function fetchSessionEvents( const d = await apiFetch(`/api/sessions/${sid}/events${q}`) as { events?: EventEntry[] }; return (d?.events ?? []).map((e) => ({ ...e, _sessionId: sid })); } catch { return [] as EventEntry[]; } - }) + }), ); - return { events: allEvents.flat(), sessionIds: ids }; + return { events: allEvents.flat(), sessionIds: ids, searchMethod: 'legacy-failed-sessions' }; } const KEYWORD_STOP_WORDS = new Set([ @@ -191,11 +248,11 @@ export function registerSearchTools(server: McpServer, knowledgeBase?: SearchPro .describe('Minimum relevance score (0-1, default 0.1). Events matching at least one keyword in any field pass this threshold.'), }, READ_ONLY, - async (args) => { + async (args) => withToolTelemetry('search_events_semantic', async () => { try { const { query, sessionId, tenantId, topK, minScore } = args; - const { events, sessionIds } = await fetchSessionEvents(sessionId, tenantId); + // Extract keywords FIRST so we can use them for index-based pre-filtering const queryKeywords = extractKeywords(query); if (queryKeywords.length === 0) { return { @@ -206,6 +263,8 @@ export function registerSearchTools(server: McpServer, knowledgeBase?: SearchPro }; } + const { events, sessionIds, searchMethod } = await fetchSessionEvents(sessionId, tenantId, queryKeywords); + const scored: Array = []; for (let i = 0; i < events.length; i++) { const result = scoreEvent(events[i], queryKeywords); @@ -234,6 +293,7 @@ export function registerSearchTools(server: McpServer, knowledgeBase?: SearchPro text: JSON.stringify({ query, searchBackend: 'weighted-keyword', + searchMethod, keywordsUsed: queryKeywords, sessionsSearched: sessionIds, eventsScanned: events.length, @@ -246,7 +306,7 @@ export function registerSearchTools(server: McpServer, knowledgeBase?: SearchPro } catch (error: unknown) { return toolError('search_events_semantic', args, error); } - } + }) ); // Tool 10: search_knowledge (unchanged — vector, pre-indexed at startup) @@ -265,7 +325,7 @@ export function registerSearchTools(server: McpServer, knowledgeBase?: SearchPro .describe('Minimum similarity score threshold (0-1, default 0.3). Lower = more results, higher = stricter matching.'), }, READ_ONLY, - async (args) => { + async (args) => withToolTelemetry('search_knowledge', async () => { try { const { query, topK, type, minScore } = args; if (!knowledgeBase || knowledgeBase.size === 0) { @@ -310,7 +370,7 @@ export function registerSearchTools(server: McpServer, knowledgeBase?: SearchPro } catch (error: unknown) { return toolError('search_knowledge', args, error); } - } + }) ); // Tool 23: deep_search_events — same scoring + exhaustive data scan @@ -336,10 +396,22 @@ export function registerSearchTools(server: McpServer, knowledgeBase?: SearchPro .describe('Additional exact keywords for matching. Auto-extracted from query if omitted.'), }, READ_ONLY, - async (args) => { + async (args) => withToolTelemetry('deep_search_events', async () => { try { const { query, sessionId, tenantId, topK, minScore, keywords } = args; - const { events, sessionIds } = await fetchSessionEvents(sessionId, tenantId); + + // Extract keywords FIRST for index-based pre-filtering + const queryKeywords = keywords ?? extractKeywords(query); + if (queryKeywords.length === 0) { + return { + content: [{ + type: 'text' as const, + text: JSON.stringify({ query, resultCount: 0, results: [], note: 'No searchable keywords extracted from query.' }), + }], + }; + } + + const { events, sessionIds, searchMethod } = await fetchSessionEvents(sessionId, tenantId, queryKeywords); if (events.length === 0) { return { @@ -353,16 +425,6 @@ export function registerSearchTools(server: McpServer, knowledgeBase?: SearchPro }; } - const queryKeywords = keywords ?? extractKeywords(query); - if (queryKeywords.length === 0) { - return { - content: [{ - type: 'text' as const, - text: JSON.stringify({ query, resultCount: 0, results: [], note: 'No searchable keywords extracted from query.' }), - }], - }; - } - const scored: Array = []; for (let i = 0; i < events.length; i++) { const result = scoreEvent(events[i], queryKeywords); @@ -391,6 +453,7 @@ export function registerSearchTools(server: McpServer, knowledgeBase?: SearchPro text: JSON.stringify({ query, searchBackend: 'weighted-keyword', + searchMethod, keywordsUsed: queryKeywords, sessionsSearched: sessionIds, totalEventsScanned: events.length, @@ -403,6 +466,6 @@ export function registerSearchTools(server: McpServer, knowledgeBase?: SearchPro } catch (error: unknown) { return toolError('deep_search_events', args, error); } - } + }) ); } diff --git a/src/McpServer/autopilot-monitor-mcp/src/tools/sessions.ts b/src/McpServer/autopilot-monitor-mcp/src/tools/sessions.ts index d69ac442..6f237b31 100644 --- a/src/McpServer/autopilot-monitor-mcp/src/tools/sessions.ts +++ b/src/McpServer/autopilot-monitor-mcp/src/tools/sessions.ts @@ -1,6 +1,7 @@ import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { z } from 'zod'; import { apiFetch, buildQuery } from '../client.js'; +import { withToolTelemetry } from '../telemetry.js'; import { READ_ONLY } from './shared.js'; import { toolError } from './error-handler.js'; @@ -65,7 +66,7 @@ export function registerSessionTools(server: McpServer): void { limit: z.coerce.number().min(1).max(100).optional().default(50).describe('Maximum number of results (1-100, default 50)'), }, READ_ONLY, - async (args) => { + async (args) => withToolTelemetry('search_sessions', async () => { try { const { deviceProperties, tenantId, ...rest } = args; const queryParams: Record = { ...rest }; @@ -81,7 +82,7 @@ export function registerSessionTools(server: McpServer): void { } catch (error: unknown) { return toolError('search_sessions', args, error); } - } + }) ); // Tool 2: search_sessions_by_event @@ -97,7 +98,7 @@ export function registerSessionTools(server: McpServer): void { limit: z.coerce.number().min(1).max(100).optional().default(50), }, READ_ONLY, - async (args) => { + async (args) => withToolTelemetry('search_sessions_by_event', async () => { try { const { tenantId, ...rest } = args; const params: Record = { ...rest }; @@ -108,7 +109,7 @@ export function registerSessionTools(server: McpServer): void { } catch (error: unknown) { return toolError('search_sessions_by_event', args, error); } - } + }) ); // Tool 3: get_session @@ -121,7 +122,7 @@ export function registerSessionTools(server: McpServer): void { includeAnalysis: z.boolean().optional().default(false).describe('Include rule analysis results (failure explanations and remediation steps)'), }, READ_ONLY, - async (args) => { + async (args) => withToolTelemetry('get_session', async () => { try { const { sessionId, tenantId, includeAnalysis } = args; const q = buildQuery({ tenantId } as Record); @@ -138,7 +139,7 @@ export function registerSessionTools(server: McpServer): void { } catch (error: unknown) { return toolError('get_session', args, error); } - } + }) ); // Tool 4: get_session_events @@ -157,7 +158,7 @@ export function registerSessionTools(server: McpServer): void { source: z.string().optional().describe('Filter by event source/app name (e.g. "MicrosoftTeams")'), }, READ_ONLY, - async (args) => { + async (args) => withToolTelemetry('get_session_events', async () => { try { const { sessionId, tenantId, ...filters } = args; const q = buildQuery({ tenantId } as Record); @@ -178,7 +179,7 @@ export function registerSessionTools(server: McpServer): void { } catch (error: unknown) { return toolError('get_session_events', args, error); } - } + }) ); // Tool 5: get_session_summary @@ -195,7 +196,7 @@ export function registerSessionTools(server: McpServer): void { tenantId: z.string().optional().describe('Tenant ID. If omitted, auto-resolved from the session (Global Admin can access any tenant).'), }, READ_ONLY, - async (args) => { + async (args) => withToolTelemetry('get_session_summary', async () => { try { const { sessionId, tenantId } = args; const q = buildQuery({ tenantId } as Record); @@ -312,7 +313,7 @@ export function registerSessionTools(server: McpServer): void { } catch (error: unknown) { return toolError('get_session_summary', args, error); } - } + }) ); // Tool 6: get_metrics @@ -325,7 +326,7 @@ export function registerSessionTools(server: McpServer): void { days: z.coerce.number().refine(v => [7, 30, 90].includes(v), { message: 'days must be 7, 30, or 90' }).optional().default(30), }, READ_ONLY, - async (args) => { + async (args) => withToolTelemetry('get_metrics', async () => { try { const { tenantId, ...rest } = args; const params: Record = { ...rest }; @@ -340,7 +341,7 @@ export function registerSessionTools(server: McpServer): void { } catch (error: unknown) { return toolError('get_metrics', args, error); } - } + }) ); // Tool 7: search_sessions_by_cve @@ -357,7 +358,7 @@ export function registerSessionTools(server: McpServer): void { limit: z.coerce.number().min(1).max(100).optional().default(50), }, READ_ONLY, - async (args) => { + async (args) => withToolTelemetry('search_sessions_by_cve', async () => { try { const { tenantId, ...rest } = args; const params: Record = { ...rest }; @@ -368,7 +369,7 @@ export function registerSessionTools(server: McpServer): void { } catch (error: unknown) { return toolError('search_sessions_by_cve', args, error); } - } + }) ); // Tool 8: list_blocked_devices @@ -380,7 +381,7 @@ export function registerSessionTools(server: McpServer): void { tenantId: z.string().optional().describe('Tenant ID to scope results. Omit for cross-tenant listing (Global Admin only).'), }, READ_ONLY, - async (args) => { + async (args) => withToolTelemetry('list_blocked_devices', async () => { try { const { tenantId } = args; const endpoint = tenantId @@ -391,7 +392,7 @@ export function registerSessionTools(server: McpServer): void { } catch (error: unknown) { return toolError('list_blocked_devices', args, error); } - } + }) ); // Tool: get_ime_version_history @@ -403,13 +404,13 @@ export function registerSessionTools(server: McpServer): void { 'Available to all tenant members (no tenantId needed, data is global).', {}, READ_ONLY, - async (args) => { + async (args) => withToolTelemetry('get_ime_version_history', async () => { try { const data = await apiFetch('/api/metrics/ime-versions'); return { content: [{ type: 'text' as const, text: JSON.stringify(data, null, 2) }] }; } catch (error: unknown) { return toolError('get_ime_version_history', args, error); } - } + }) ); }