diff --git a/src/McpServer/autopilot-monitor-mcp/src/__tests__/search-input-bounds.test.ts b/src/McpServer/autopilot-monitor-mcp/src/__tests__/search-input-bounds.test.ts new file mode 100644 index 00000000..cbbdf4b2 --- /dev/null +++ b/src/McpServer/autopilot-monitor-mcp/src/__tests__/search-input-bounds.test.ts @@ -0,0 +1,341 @@ +/** + * Caller-sized input must not decide what one call costs: one MCP process serves every tenant, so a + * single request may cost no more than a small multiple of its event volume. Pins the compiled concept + * matcher against the per-keyword rule it replaced (seeded differential test), its keyword-count-free + * per-event cost, the cooperative scoring loop, the keyword and error-code caps, and the schema bounds + * on query, keywords, section and continuation. + * + * Pure unit tests: fetch is stubbed and the registered tool handlers are invoked directly. + */ +import { describe, it, expect, vi, afterEach } from 'vitest'; +import { McpServer } from '@modelcontextprotocol/server'; +import { registerTools } from '../tools.js'; +import { buildQuery, runWithCaller } from '../client.js'; +import { + compileConcepts, + dedupeSynonymConcepts, + extractErrorCodeNeedles, + KEYWORD_SYNONYM_GROUPS, + matchConcepts, + MAX_ERROR_CODE_NEEDLES, + MAX_SCORED_KEYWORDS, + MAX_SEARCH_KEYWORD_CHARS, + MAX_SEARCH_KEYWORDS, + MAX_SEARCH_QUERY_CHARS, + SCORED_FIELDS, + scoreEvent, + scoreEvents, +} from '../tools/search.js'; +import type { SearchProvider } from '../search-provider.js'; + +// ── Deterministic inputs ──────────────────────────────────────────────── + +/** mulberry32: seeded, so every run checks the same cases. */ +function prng(seed: number): () => number { + let s = seed; + return () => { + s = (s + 0x6d2b79f5) | 0; + let t = Math.imul(s ^ (s >>> 15), 1 | s); + t = (t + Math.imul(t ^ (t >>> 7), 61 | t)) ^ t; + return ((t ^ (t >>> 14)) >>> 0) / 4294967296; + }; +} + +type Ev = { eventType?: string; severity?: string; source?: string; message?: string; data?: Record }; + +const VOCAB = ['install', 'app', 'download', 'policy', 'device', 'enrollment', 'phase', 'script', 'detection', 'exit', + 'code', 'started', 'completed', 'certificate', 'network', 'agent', 'registry', 'value', 'result', 'status']; + +/** Production-sized events: a ~60-character message and ~600 characters of data, every field filled. */ +function makeEvents(n: number, seed = 1): Ev[] { + const rnd = prng(seed); + const words = (len: number): string => { + let s = ''; + while (s.length < len) s += (s ? ' ' : '') + VOCAB[Math.floor(rnd() * VOCAB.length)]; + return s; + }; + return Array.from({ length: n }, (_, i) => ({ + eventType: `${VOCAB[i % VOCAB.length]}_${VOCAB[(i * 7) % VOCAB.length]}`, + severity: ['Info', 'Warning', 'Error'][i % 3], + source: 'ImeLogTracker', + message: words(60), + data: Object.fromEntries(Array.from({ length: 8 }, (_, k) => [`${VOCAB[k]}Key${k}`, words(60)])), + })); +} + +/** Distinct keywords that occur in no event. */ +const absent = (n: number): string[] => Array.from({ length: n }, (_, i) => `zq${i.toString(36).padStart(5, '0')}`); + +// ── The per-keyword rule the compiled matcher replaced, written out as the oracle ── + +const GROUP_OF = new Map(KEYWORD_SYNONYM_GROUPS.flatMap((group) => group.map((term) => [term, group] as const))); +const WEIGHT = { eventType: 3, message: 2, source: 1.5, severity: 1, data: 0.5 } as const; + +function oracleWordPrefix(text: string, term: string): boolean { + if (term.length < 4) return false; + return text.split(/[\s_\-.:,/]+/).some((word) => { + if (word.length < 4) return false; + const p = Math.min(word.length, term.length, 6); + return word.slice(0, p) === term.slice(0, p); + }); +} + +/** Keyword as substring or word prefix; its synonyms at word level only. */ +function oracleFieldMatch(text: string, keyword: string): boolean { + if (text.includes(keyword) || oracleWordPrefix(text, keyword)) return true; + return (GROUP_OF.get(keyword) ?? []).some((term) => term !== keyword && oracleWordPrefix(text, term)); +} + +/** Per concept: the heaviest field it matches, or null. */ +function oracleBestFields(e: Ev, concepts: readonly string[]): Array { + const texts: Array<[keyof typeof WEIGHT, string]> = [ + ['eventType', (e.eventType ?? '').toLowerCase()], + ['message', (e.message ?? '').toLowerCase()], + ['source', (e.source ?? '').toLowerCase()], + ['severity', (e.severity ?? '').toLowerCase()], + ['data', e.data ? JSON.stringify(e.data).toLowerCase() : ''], + ]; + return concepts.map((keyword) => { + let best: keyof typeof WEIGHT | null = null; + for (const [name, text] of texts) { + if (text && oracleFieldMatch(text, keyword) && (best === null || WEIGHT[name] > WEIGHT[best])) best = name; + } + return best; + }); +} + +// Stems with prefix collisions at every length around the 4-6 character comparison window, synonyms +// and the words they hide in ("in[stall]", "c[hang]e"), short keywords below the prefix minimum. +const STEMS = ['abcd', 'abcde', 'abcdef', 'abcdefg', 'abcdx', 'abc', 'ab', 'abce', 'abcdeq', 'xyzw', 'xyz', 'install', + 'installer', 'installation', 'setup', 'setups', 'reboot', 'rebooting', 'restarts', 'reset', 'resetting', 'stuck', + 'stuckness', 'hang', 'hangs', 'change', 'hung', 'timeout', 'timedout', 'stall', 'stalled', 'download', 'downloads', + 'transfer', 'hello', 'hel', 'whfb', 'failed', 'failure', 'error', 'esp', 'msi', 'auth', 'authentication', 'app', 'do', 'id']; +const ch = String.fromCharCode; +// The old split's delimiters, Unicode whitespace it also splits on, characters it does not split on, +// and nothing at all (two stems glued together). +const SEPARATORS = [' ', '_', '-', '.', ':', ',', '/', ch(9), ch(10), ' ', '__', ' - ', ch(0xa0), ch(0x2028), ch(0x3000), + ch(0xfeff), '(', ')', '"', '=', '']; + +function randomText(rnd: () => number, maxWords: number): string { + const n = Math.floor(rnd() * maxWords); + let s = ''; + for (let i = 0; i < n; i++) { + const stem = STEMS[Math.floor(rnd() * STEMS.length)]; + s += (i ? SEPARATORS[Math.floor(rnd() * SEPARATORS.length)] : '') + (rnd() < 0.2 ? stem.toUpperCase() : stem); + } + return s; +} + +function randomKeywords(rnd: () => number): string[] { + const pick = (): string => STEMS[Math.floor(rnd() * STEMS.length)]; + const out = new Set(); + const n = 1 + Math.floor(rnd() * 8); + for (let i = 0; i < n; i++) { + const r = rnd(); + const keyword = r < 0.6 ? pick() + : r < 0.75 ? pick().slice(0, 2 + Math.floor(rnd() * 6)) + : r < 0.9 ? `${pick()} ${pick()}` + : pick() + ['ing', 's', 'ed', 'x'][Math.floor(rnd() * 4)]; + if (keyword) out.add(keyword); + } + return [...out]; +} + +// ── Registered tools ──────────────────────────────────────────────────── + +type ToolHandler = (args: Record, extra: unknown) => Promise<{ content: Array<{ text: string }> }>; +type ParseableSchema = { safeParse(value: unknown): { success: boolean } }; +type RegisteredTool = { handler: ToolHandler; inputSchema: ParseableSchema & { shape: Record } }; + +const stubDocs: SearchProvider = { name: 'stub', semanticCapable: true, size: 1, index: async () => {}, search: async () => [] }; + +function registeredTools(): Record { + const server = new McpServer({ name: 'test', version: '0.0.0' }); + registerTools(server, undefined, undefined, { vector: stubDocs, sections: ['trust'] }, true, true, false); + return (server as unknown as { _registeredTools: Record })._registeredTools; +} + +const SESSION = 'e259c121-1234-4abc-9def-0123456789ab'; +const GA = { token: 'ga', isGlobalAdmin: true }; +const extra = { signal: new AbortController().signal }; + +function stubSessionEvents(events: Ev[]): void { + vi.stubGlobal('fetch', vi.fn(async () => { + const body = { success: true, sessionId: SESSION, count: events.length, events }; + return { ok: true, status: 200, json: async () => body, text: async () => JSON.stringify(body) } as unknown as Response; + })); +} + +afterEach(() => vi.unstubAllGlobals()); + +// ── Tests ─────────────────────────────────────────────────────────────── + +describe('compiled concept matcher', () => { + it('lists the scored fields strongest first, so the first field a concept matches is its strongest', () => { + expect(SCORED_FIELDS.map((f) => f.name)).toEqual(['eventType', 'message', 'source', 'severity', 'data']); + const weights = SCORED_FIELDS.map((f) => f.weight); + expect(weights).toEqual([...weights].sort((a, b) => b - a)); + }); + + it('matches every concept in the same strongest field as the per-keyword rule it replaced', () => { + const rnd = prng(7); + const names = SCORED_FIELDS.map((f) => f.name); + let compared = 0; + for (let i = 0; i < 10_000; i++) { + const keywords = randomKeywords(rnd); + const e: Ev = { + eventType: randomText(rnd, 4), + message: randomText(rnd, 12), + source: randomText(rnd, 3), + severity: ['Info', 'Error', 'Warning', '', 'Debug'][Math.floor(rnd() * 5)], + data: rnd() < 0.7 ? { a: randomText(rnd, 10), b: randomText(rnd, 6), n: Math.floor(rnd() * 100) } : undefined, + }; + const matcher = compileConcepts(keywords); + expect(matcher.concepts).toEqual(dedupeSynonymConcepts(keywords)); + const got = matchConcepts(e, matcher).map((f) => (f < 0 ? null : names[f])); + expect(got, JSON.stringify({ keywords, e })).toEqual(oracleBestFields(e, matcher.concepts)); + compared += got.length; + } + expect(compared).toBeGreaterThan(20_000); + }); + + it('splits each field once per event, however many keywords the query carries', () => { + const events = makeEvents(200); + const splitsFor = (keywords: string[]): number => { + const matcher = compileConcepts(keywords); + const split = vi.spyOn(String.prototype, 'split'); + try { + for (const e of events) scoreEvent(e, keywords, false, undefined, undefined, matcher); + return split.mock.calls.length; + } finally { + split.mockRestore(); + } + }; + const one = splitsFor(absent(1)); + expect(one).toBe(events.length * SCORED_FIELDS.length); + expect(splitsFor(absent(MAX_SCORED_KEYWORDS))).toBe(one); + }); + + it('scores a maximal request in well under a second', async () => { + const events = makeEvents(3000); + const t0 = performance.now(); + await scoreEvents(events, absent(MAX_SCORED_KEYWORDS), false, new Map(), new Map(), 0.05); + expect(performance.now() - t0).toBeLessThan(1000); + }); +}); + +describe('scoreEvents', () => { + it('returns exactly what the per-event scorer returns', async () => { + const events = makeEvents(700, 3); + const keywords = ['install', 'certificate', 'stuck', 'agent']; + const scored = await scoreEvents(events, keywords, true, new Map(), new Map(), 0.05); + const expected = events.flatMap((e, index) => { + const r = scoreEvent(e, keywords, true, new Map(), new Map()); + return r && r.score >= 0.05 ? [{ ...r, index }] : []; + }); + expect(scored.length).toBeGreaterThan(0); + expect(scored.map(({ index, score, matchedKeywords, bestFields }) => ({ index, score, matchedKeywords, bestFields }))) + .toEqual(expected); + }); + + it('yields to the event loop while it scores, so other requests keep running', async () => { + let otherRequestRan = false; + setImmediate(() => { otherRequestRan = true; }); + await scoreEvents(makeEvents(600), absent(1), false, new Map(), new Map(), 0); + expect(otherRequestRan).toBe(true); + }); + + it('stops once the caller has gone', async () => { + const gone = new AbortController(); + gone.abort(); + await expect(runWithCaller({ token: 't', isGlobalAdmin: false, signal: gone.signal }, + () => scoreEvents(makeEvents(600), absent(1), false, new Map(), new Map(), 0))).rejects.toThrow(); + }); +}); + +describe('search_events keyword cap', () => { + const tools = registeredTools(); + const run = async (args: Record): Promise> => { + stubSessionEvents(makeEvents(5)); + const res = await runWithCaller(GA, () => tools.search_events.handler( + { sessionId: SESSION, depth: 'fast', guaranteedTopRanked: 3, ...args }, extra)); + return JSON.parse(res.content[0].text) as Record; + }; + + it(`scores the first ${MAX_SCORED_KEYWORDS} keywords, caller keywords first, and reports the rest`, async () => { + const query = Array.from({ length: 40 }, (_, i) => `queryword${i}`).join(' '); + const body = await run({ query, keywords: ['callerword'] }); + const used = body.keywordsUsed as string[]; + expect(used).toHaveLength(MAX_SCORED_KEYWORDS); + expect(used[0]).toBe('callerword'); + expect(body.keywordsDropped).toBe(41 - MAX_SCORED_KEYWORDS); + expect(body.keywordNote).toContain(String(MAX_SCORED_KEYWORDS)); + }); + + it('reports nothing dropped for an ordinary query', async () => { + const body = await run({ query: 'app install failed with certificate error' }); + expect(body.keywordsDropped).toBeUndefined(); + expect(body.keywordNote).toBeUndefined(); + }); +}); + +describe('error-code needles', () => { + it(`takes at most ${MAX_ERROR_CODE_NEEDLES} codes, 0x-prefixed ones first`, () => { + const bare = Array.from({ length: 12 }, (_, i) => `abcd${(i + 16).toString(16)}`).join(' '); + const needles = extractErrorCodeNeedles(`${bare} 0x80070002 0x87D1041C`); + expect(needles).toHaveLength(MAX_ERROR_CODE_NEEDLES); + expect(needles.slice(0, 2)).toEqual(['80070002', '87d1041c']); + }); +}); + +describe('schema bounds', () => { + const tools = registeredTools(); + const accepts = (tool: string, args: Record): boolean => tools[tool].inputSchema.safeParse(args).success; + + it('bounds the query of all three search tools', () => { + for (const tool of ['search_events', 'search_knowledge', 'search_docs']) { + expect(accepts(tool, { query: 'x'.repeat(MAX_SEARCH_QUERY_CHARS) }), tool).toBe(true); + expect(accepts(tool, { query: 'x'.repeat(MAX_SEARCH_QUERY_CHARS + 1) }), tool).toBe(false); + } + }); + + it('bounds the caller keywords of search_events in number and length', () => { + const atLimit = Array.from({ length: MAX_SEARCH_KEYWORDS }, () => 'k'.repeat(MAX_SEARCH_KEYWORD_CHARS)); + expect(accepts('search_events', { query: 'q', keywords: atLimit })).toBe(true); + expect(accepts('search_events', { query: 'q', keywords: [...atLimit, 'k'] })).toBe(false); + expect(accepts('search_events', { query: 'q', keywords: ['k'.repeat(MAX_SEARCH_KEYWORD_CHARS + 1)] })).toBe(false); + }); + + it('bounds the search_docs section', () => { + expect(accepts('search_docs', { query: 'q', section: 'trust' })).toBe(true); + expect(accepts('search_docs', { query: 'q', section: 's'.repeat(101) })).toBe(false); + }); + + it('accepts a real nextLink on every paginated tool and refuses a forged one', () => { + const nextLink = `/api/search/sessions?pageSize=50&tenantId=11111111-1111-4111-8111-111111111111&continuation=${'A'.repeat(800)}`; + const forged = `/api/search/sessions?${'fields=&'.repeat(36_000)}`; + const paginated = Object.entries(tools).filter(([, tool]) => 'continuation' in tool.inputSchema.shape); + expect(paginated.map(([name]) => name)).toEqual(expect.arrayContaining(['search_sessions', 'get_session_events', 'get_audit_logs'])); + for (const [name, tool] of paginated) { + expect(tool.inputSchema.shape.continuation.safeParse(nextLink).success, name).toBe(true); + expect(tool.inputSchema.shape.continuation.safeParse(forged).success, name).toBe(false); + } + }); +}); + +describe('buildQuery', () => { + it('builds the same query string as before, skipping empty values', () => { + expect(buildQuery({ a: '1', b: 'x y', c: undefined, d: '', e: 0, f: null })).toBe('?a=1&b=x+y&e=0'); + expect(buildQuery({})).toBe(''); + }); + + it('stays linear for a caller-sized key set', () => { + // Far more keys than a request body carries, so the complexity class is unmistakable: + // measured ~0.1 s linear against ~9 s for a per-key search of the whole list. + const many = Object.fromEntries(Array.from({ length: 64_000 }, (_, i) => [`prop.os_info.k${i}`, 'v'])); + const t0 = performance.now(); + const q = buildQuery(many); + expect(performance.now() - t0).toBeLessThan(2000); + expect(q.split('&')).toHaveLength(64_000); + }); +}); diff --git a/src/McpServer/autopilot-monitor-mcp/src/client.ts b/src/McpServer/autopilot-monitor-mcp/src/client.ts index 597a0a7b..496347a1 100644 --- a/src/McpServer/autopilot-monitor-mcp/src/client.ts +++ b/src/McpServer/autopilot-monitor-mcp/src/client.ts @@ -523,8 +523,10 @@ async function apiFetch(path: string, options: ApiFetchOptions = {}): Promise function buildQuery(params: Record): string { const p = new URLSearchParams(); + // append, not set: object keys are unique, and set() searches the whole list on every call — + // quadratic for a caller-sized key set such as search_sessions' deviceProperties. for (const [k, v] of Object.entries(params)) { - if (v != null && v !== '') p.set(k, String(v)); + if (v != null && v !== '') p.append(k, String(v)); } const s = p.toString(); return s ? `?${s}` : ''; diff --git a/src/McpServer/autopilot-monitor-mcp/src/tools/admin.ts b/src/McpServer/autopilot-monitor-mcp/src/tools/admin.ts index f19b28f2..463d4fdf 100644 --- a/src/McpServer/autopilot-monitor-mcp/src/tools/admin.ts +++ b/src/McpServer/autopilot-monitor-mcp/src/tools/admin.ts @@ -3,7 +3,7 @@ import { z } from 'zod'; import { apiFetch, buildQuery, DEFAULT_FIRST_PAGE_SIZE, effectivePageSize, enforceDelegatedTenant, enforceDelegatedTenantForPage, followNextLink, getCallerUpnDomain, getDelegatedTenantIds, getHomeTenantId, pageSizeForCall, pickGlobalOrTenantPath, scanWithTimeoutFallback, jsonBody } from '../client.js'; import { withToolTelemetry } from '../telemetry.js'; import { getResourceContent, withEventTypeNote, RESOURCE_NAMES } from '../resource-catalog.js'; -import { READ_ONLY, READ_ONLY_OPEN, MUTATING, MAX_RESULT_SIZE_CHARS, LEAN_RAW_EVENT_FIELDS, LEAN_RAW_EVENT_OMISSION, leanFieldSelection, toolResultText, SessionIdSchema, TenantGuidSchema, tenantIdDescription, pageSizeDescription, CONTINUATION_DESCRIPTION, daysDescription } from './shared.js'; +import { READ_ONLY, READ_ONLY_OPEN, MUTATING, MAX_RESULT_SIZE_CHARS, LEAN_RAW_EVENT_FIELDS, LEAN_RAW_EVENT_OMISSION, leanFieldSelection, toolResultText, SessionIdSchema, TenantGuidSchema, tenantIdDescription, pageSizeDescription, CONTINUATION_DESCRIPTION, ContinuationSchema, daysDescription } from './shared.js'; import { toolError } from './error-handler.js'; import { API_BASE_URL } from '../config.js'; import { collectDeploymentState } from '../deployment-state.js'; @@ -430,7 +430,7 @@ export function registerAdminTools(server: McpServer, ga: boolean, strictGa: boo daily: z.boolean().optional().default(false).describe('Return daily aggregated summary instead of per-endpoint breakdown'), pageSize: z.coerce.number().int().min(1).max(2000).optional().default(50) .describe(pageSizeDescription(50, 2000)), - continuation: z.string().optional() + continuation: ContinuationSchema.optional() .describe(CONTINUATION_DESCRIPTION + ' Offset cursor: re-send the same filters with it.'), }, annotations: READ_ONLY, @@ -512,7 +512,7 @@ export function registerAdminTools(server: McpServer, ga: boolean, strictGa: boo .describe(daysDescription(30, 365)), pageSize: z.coerce.number().int().min(1).max(1000).optional().default(50) .describe(pageSizeDescription(50)), - continuation: z.string().optional() + continuation: ContinuationSchema.optional() .describe(CONTINUATION_DESCRIPTION + ' Offset cursor: re-send the same filters with it.'), }, annotations: READ_ONLY, @@ -770,7 +770,7 @@ export function registerAdminTools(server: McpServer, ga: boolean, strictGa: boo .describe(daysDescription(7, 365, 'Applies to the stats and the session list.')), pageSize: z.coerce.number().int().min(1).max(200).optional() .describe(pageSizeDescription(FLEET_OVERVIEW_PAGE_SIZE, 200)), - continuation: z.string().optional() + continuation: ContinuationSchema.optional() .describe(CONTINUATION_DESCRIPTION), }, annotations: READ_ONLY, @@ -835,7 +835,7 @@ export function registerAdminTools(server: McpServer, ga: boolean, strictGa: boo 'included; default: all safe fields. Applied server-side; unknown/secret keys are ignored.'), pageSize: z.coerce.number().int().min(1).max(1000).optional() .describe(pageSizeDescription(100, 1000, 'Ignored in filter mode (query/tenantId).')), - continuation: z.string().optional() + continuation: ContinuationSchema.optional() .describe(CONTINUATION_DESCRIPTION + ' Ignored in filter mode.'), }, annotations: READ_ONLY, @@ -939,7 +939,7 @@ export function registerAdminTools(server: McpServer, ga: boolean, strictGa: boo entityId: z.string().optional().describe('Exact-match filter on the affected entity id (e.g. a tenantId, deviceId, or report id).'), pageSize: z.coerce.number().int().min(1).max(1000).optional() .describe(pageSizeDescription(DEFAULT_FIRST_PAGE_SIZE)), - continuation: z.string().optional() + continuation: ContinuationSchema.optional() .describe(CONTINUATION_DESCRIPTION), }, annotations: READ_ONLY, @@ -995,7 +995,7 @@ export function registerAdminTools(server: McpServer, ga: boolean, strictGa: boo dateTo: z.string().optional().describe('ISO 8601 UTC timestamp — inclusive upper bound of the window.'), pageSize: z.coerce.number().int().min(1).max(1000).optional() .describe(pageSizeDescription(DEFAULT_FIRST_PAGE_SIZE)), - continuation: z.string().optional() + continuation: ContinuationSchema.optional() .describe(CONTINUATION_DESCRIPTION), }, annotations: READ_ONLY, @@ -1082,7 +1082,7 @@ export function registerAdminTools(server: McpServer, ga: boolean, strictGa: boo tenantId: z.string().optional().describe(tenantIdDescription(ga, delegated)), pageSize: z.coerce.number().int().min(1).max(1000).optional() .describe(pageSizeDescription(DEFAULT_FIRST_PAGE_SIZE)), - continuation: z.string().optional() + continuation: ContinuationSchema.optional() .describe(CONTINUATION_DESCRIPTION), }, annotations: READ_ONLY, @@ -1125,7 +1125,7 @@ export function registerAdminTools(server: McpServer, ga: boolean, strictGa: boo .describe('One effective status. "published" lists the approved submissions whose rule is live in the global catalog.'), pageSize: z.coerce.number().int().min(1).max(1000).optional() .describe(pageSizeDescription(DEFAULT_FIRST_PAGE_SIZE)), - continuation: z.string().optional() + continuation: ContinuationSchema.optional() .describe(CONTINUATION_DESCRIPTION), }, annotations: READ_ONLY, @@ -1400,7 +1400,7 @@ export function registerAdminTools(server: McpServer, ga: boolean, strictGa: boo dateTo: z.string().optional().describe('ISO 8601 UTC timestamp — exclusive upper bound on the last-updated time.'), pageSize: z.coerce.number().int().min(1).max(1000).optional() .describe(pageSizeDescription(DEFAULT_FIRST_PAGE_SIZE)), - continuation: z.string().optional() + continuation: ContinuationSchema.optional() .describe(CONTINUATION_DESCRIPTION), }, annotations: READ_ONLY, @@ -1496,7 +1496,7 @@ export function registerAdminTools(server: McpServer, ga: boolean, strictGa: boo .describe('Keep-list of column names; PartitionKey + RowKey always kept. Columns: PartitionKey, RowKey, Timestamp (storage write time), OccurredUtc (event time), EventId, SessionId, TenantId, EventType, Severity (int), Source, Phase (int), Message, Sequence, DataJson (raw string), ReceivedAt, SentAt, OriginalTimestamp, TimestampClamped, CausedByTransitionStepIndex, CausedBySignalOrdinal. Default: all but DataJson on an unfiltered read; the full row on a filtered read (eventType/severity/source).'), pageSize: z.coerce.number().int().min(1).max(1000).optional() .describe(pageSizeDescription(DEFAULT_FIRST_PAGE_SIZE, 1000, 'Index rows walked per call. When the projection includes DataJson the ceiling is ' + RAW_EVENTS_PAYLOAD_PAGE_MAX + ' (default ' + RAW_EVENTS_PAYLOAD_PAGE_DEFAULT + '); a larger value is refused.')), - continuation: z.string().optional() + continuation: ContinuationSchema.optional() .describe(CONTINUATION_DESCRIPTION), }, annotations: READ_ONLY, @@ -1579,7 +1579,7 @@ export function registerAdminTools(server: McpServer, ga: boolean, strictGa: boo fields: z.string().optional().describe('Keep-list of literal column names (case-insensitive, PascalCase, e.g. "Status,OsEdition,ImeAgentVersion,GeoCity"); PartitionKey + RowKey are always kept. Omit for the full raw row.'), pageSize: z.coerce.number().int().min(1).max(1000).optional() .describe(pageSizeDescription(DEFAULT_FIRST_PAGE_SIZE)), - continuation: z.string().optional() + continuation: ContinuationSchema.optional() .describe(CONTINUATION_DESCRIPTION), }, annotations: READ_ONLY, @@ -1654,7 +1654,7 @@ export function registerAdminTools(server: McpServer, ga: boolean, strictGa: boo 'client-side. PartitionKey and RowKey are always kept.'), pageSize: z.coerce.number().int().min(1).max(1000).optional() .describe(pageSizeDescription(DEFAULT_FIRST_PAGE_SIZE)), - continuation: z.string().optional() + continuation: ContinuationSchema.optional() .describe(CONTINUATION_DESCRIPTION), }, annotations: READ_ONLY, @@ -2195,7 +2195,7 @@ export function registerAdminTools(server: McpServer, ga: boolean, strictGa: boo ...inventoryScopeShape, pageSize: z.coerce.number().int().min(1).max(500).optional().default(100) .describe(pageSizeDescription(100, 500)), - continuation: z.string().optional() + continuation: ContinuationSchema.optional() .describe(CONTINUATION_DESCRIPTION + ' Offset cursor: re-send the same filters with it.'), }, annotations: READ_ONLY, diff --git a/src/McpServer/autopilot-monitor-mcp/src/tools/search.ts b/src/McpServer/autopilot-monitor-mcp/src/tools/search.ts index 838662e0..f268e85c 100644 --- a/src/McpServer/autopilot-monitor-mcp/src/tools/search.ts +++ b/src/McpServer/autopilot-monitor-mcp/src/tools/search.ts @@ -1,6 +1,6 @@ import { McpServer } from '@modelcontextprotocol/server'; import { z } from 'zod'; -import { apiFetch, buildQuery, enforceDelegatedTenant, followNextLink, pickGlobalOrTenantPath } from '../client.js'; +import { apiFetch, buildQuery, enforceDelegatedTenant, followNextLink, getCallerSignal, pickGlobalOrTenantPath } from '../client.js'; import { withToolTelemetry, logSearchZeroHit } from '../telemetry.js'; import type { SearchProvider, DocsSearchBundle } from '../search-provider.js'; import { READ_ONLY, MAX_RESULT_SIZE_CHARS, toolResultText, SessionIdSchema, isBenignHealthDetectionReport, tenantIdDescription } from './shared.js'; @@ -160,6 +160,26 @@ const DEPTH_PRESETS = { deep: { budget: DEEP_BUDGET, defaultTopK: 20, defaultMinScore: 0.05 }, } as const; +// ── Input bounds ──────────────────────────────────────────────────────── +// +// One process serves every tenant, and every fetched event is scored against every query keyword, +// so the keyword count multiplies the work of a call. These bounds keep one request's cost a small +// multiple of its event volume; real queries stay far below them. + +/** Longest `query` the three search tools accept. */ +export const MAX_SEARCH_QUERY_CHARS = 1000; +/** Most caller-supplied `keywords` per search_events call. */ +export const MAX_SEARCH_KEYWORDS = 20; +/** Longest caller-supplied keyword. */ +export const MAX_SEARCH_KEYWORD_CHARS = 100; +/** Distinct keywords one search scores (caller keywords first, then the query's words); the rest is reported as dropped. */ +export const MAX_SCORED_KEYWORDS = 32; +const KEYWORDS_DROPPED_NOTE = + `Only the first ${MAX_SCORED_KEYWORDS} distinct keywords were scored (your keywords first, then words from the ` + + 'query); shorten the query to choose which ones count.'; +/** Longest search_docs `section` (a top-level folder of the docs bundle). */ +const MAX_DOCS_SECTION_CHARS = 100; + /** * EventSeverity enum (C# `EnrollmentEvent.EventSeverity`) → name. The raw endpoint stores * Severity as the underlying Int32; the enriched endpoints emit the name string. We map back @@ -561,6 +581,9 @@ export function resolveQueryKeywords(query: string, callerKeywords?: string[]): // ── Error-code fallback ───────────────────────────────────────────────── +/** Error codes one query contributes to the literal fallback. */ +export const MAX_ERROR_CODE_NEEDLES = 10; + /** * Pull opaque error-code tokens (HRESULT / Win32 / NTSTATUS hex) out of a query. These embed * poorly — to a sentence-transformer "0x87D1041C" is near-random noise — so a semantic search can @@ -572,15 +595,22 @@ export function resolveQueryKeywords(query: string, callerKeywords?: string[]): * Bare (un-prefixed) hex must contain at least one a-f letter and be 6-8 chars, so plain decimal * numbers — dates, counts, build numbers — never trigger the fallback. `0x`-prefixed tokens are * always codes, so those accept 4-8 digits (covers short Win32 HRESULTs like 0x801c03ed). + * + * At most MAX_ERROR_CODE_NEEDLES distinct codes, `0x`-prefixed ones first: every needle is scanned + * against every document of the corpus. */ export function extractErrorCodeNeedles(query: string): string[] { const needles = new Set(); // 0x-prefixed hex (0x80070002, 0x87D1041C). The `x` is not in [0-9a-f], so the bare-hex pass // below can never re-match the same digits — no double counting. - for (const m of query.matchAll(/\b0x([0-9a-f]{4,8})\b/gi)) needles.add(m[1].toLowerCase()); + for (const m of query.matchAll(/\b0x([0-9a-f]{4,8})\b/gi)) { + if (needles.size >= MAX_ERROR_CODE_NEEDLES) return [...needles]; + needles.add(m[1].toLowerCase()); + } // Bare HRESULT-shaped hex: 6-8 chars with at least one a-f letter (87D1041C) — the letter // requirement excludes pure decimals that aren't error codes. for (const m of query.matchAll(/\b([0-9a-f]{6,8})\b/gi)) { + if (needles.size >= MAX_ERROR_CODE_NEEDLES) break; if (/[a-f]/i.test(m[1])) needles.add(m[1].toLowerCase()); } return [...needles]; @@ -613,6 +643,8 @@ export function queryHasProblemIntent(keywords: string[]): boolean { } const MIN_PREFIX_LEN = 4; +/** A prefix comparison looks at no more than this many leading characters. */ +const MAX_PREFIX_LEN = MIN_PREFIX_LEN + 2; /** Split text into word tokens for prefix matching (same delimiters everywhere). */ function splitWords(text: string): string[] { @@ -624,7 +656,7 @@ function sharesWordPrefix(words: string[], keyword: string): boolean { if (keyword.length < MIN_PREFIX_LEN) return false; for (const word of words) { if (word.length < MIN_PREFIX_LEN) continue; - const prefixLen = Math.min(word.length, keyword.length, MIN_PREFIX_LEN + 2); + const prefixLen = Math.min(word.length, keyword.length, MAX_PREFIX_LEN); if (word.slice(0, prefixLen) === keyword.slice(0, prefixLen)) return true; } return false; @@ -730,7 +762,7 @@ export function orderBySelectivity(types: string[], keywords: string[], specific * non-canonical words (e.g. "restarted" for system_reboot_detected, "stuck" for a timeout). * Groups are symmetric — every term maps to the whole group. */ -const KEYWORD_SYNONYM_GROUPS: readonly string[][] = [ +export const KEYWORD_SYNONYM_GROUPS: readonly string[][] = [ ['reboot', 'rebooted', 'restart', 'restarted', 'reset'], ['timeout', 'timedout', 'stall', 'stalled', 'stuck', 'hang', 'hung'], ['hello', 'whfb', 'passwordless'], @@ -750,7 +782,7 @@ const SYNONYM_LOOKUP: ReadonlyMap = (() => { /** * Collapse query keywords that fall in the SAME synonym group down to one representative concept * (first occurrence kept), preserving order and keeping non-synonym keywords as-is. Because - * keywordMatchesField already expands a single keyword to its whole group, one representative per + * compileConcepts already expands a single keyword to its whole group, one representative per * group is sufficient — and it stops a query that spells the same concept two ways (e.g. "stuck" * AND "timeout") from counting as two matched keywords, which inflated coverage for off-topic types. * A no-op for queries with no intra-group duplicates. @@ -769,37 +801,6 @@ export function dedupeSynonymConcepts(keywords: string[]): string[] { return out; } -/** - * Match a SYNONYM at word granularity (whole word, or a shared word-prefix of >= MIN_PREFIX_LEN), - * never as a mid-word substring. Synonyms are short (stall/hang/reset/…) and the loose substring - * check used for literal keywords would mis-fire — e.g. "stall" sits inside "in[stall]", "hang" - * inside "c[hang]e" — making a "timeout" query spuriously match every app_install_* event. Word- - * level matching keeps the useful equivalences (reboot↔restart, stuck↔timeout) without those - * false positives. - */ -function synonymMatchesWord(text: string, term: string): boolean { - if (term.length < MIN_PREFIX_LEN) return false; - for (const word of text.split(/[\s_\-.:,/]+/)) { - if (word === term) return true; - if (word.length >= MIN_PREFIX_LEN) { - const prefixLen = Math.min(word.length, term.length, MIN_PREFIX_LEN + 2); - if (word.slice(0, prefixLen) === term.slice(0, prefixLen)) return true; - } - } - return false; -} - -/** True if the keyword (loose, prefix-aware) or any of its synonyms (word-level) matches `text`. */ -function keywordMatchesField(text: string, keyword: string): boolean { - if (prefixAwareMatch(text, keyword)) return true; - const group = SYNONYM_LOOKUP.get(keyword); - if (!group) return false; - for (const term of group) { - if (term !== keyword && synonymMatchesWord(text, term)) return true; - } - return false; -} - /** Field weights for scoring — eventType is the most discriminating field. */ const FIELD_WEIGHTS = { eventType: 3.0, @@ -809,6 +810,92 @@ const FIELD_WEIGHTS = { data: 0.5, } as const; +/** + * The scored fields in DESCENDING weight order: a concept counts once, at the first field it matches, + * which is only its strongest field while this order holds. + */ +export const SCORED_FIELDS: ReadonlyArray<{ name: string; weight: number; text: (e: EventEntry) => string }> = [ + { name: 'eventType', weight: FIELD_WEIGHTS.eventType, text: (e) => (e.eventType ?? '').toLowerCase() }, + { name: 'message', weight: FIELD_WEIGHTS.message, text: (e) => (e.message ?? '').toLowerCase() }, + { name: 'source', weight: FIELD_WEIGHTS.source, text: (e) => (e.source ?? '').toLowerCase() }, + { name: 'severity', weight: FIELD_WEIGHTS.severity, text: (e) => (e.severity ?? '').toLowerCase() }, + { name: 'data', weight: FIELD_WEIGHTS.data, text: (e) => (e.data ? JSON.stringify(e.data).toLowerCase() : '') }, +]; + +/** One word-prefix term of a concept (the keyword or one of its synonyms), cut to MAX_PREFIX_LEN. */ +type PrefixTerm = { concept: number; head: string }; + +/** + * A query's concepts, compiled once per search so that scoring an event splits each field into words + * once — the per-event cost does not grow with the number of keywords. + * + * A concept matches a field when its keyword is a substring of the field, or when a field word shares a + * prefix with one of the concept's terms: the shorter of the two MAX_PREFIX_LEN-character heads is a + * prefix of the longer, both at least MIN_PREFIX_LEN long (the sharesWordPrefix rule). The terms are + * the keyword and, for a synonym, its whole group. Synonyms match at word granularity only, never as a + * mid-word substring: they are short (stall/hang/reset/…), "stall" sits inside "in[stall]" and "hang" + * inside "c[hang]e", so a substring check would make a "timeout" query match every app_install_* event. + */ +export type ConceptMatcher = { + readonly concepts: readonly string[]; + /** Terms by the headKey of their first three characters. */ + readonly termsByKey: ReadonlyMap; +}; + +/** Map key of a string's first three UTF-16 code units, so a field word is looked up without cutting a substring. */ +function headKey(s: string): number { + return s.charCodeAt(0) * 0x1_0000_0000 + s.charCodeAt(1) * 0x1_0000 + s.charCodeAt(2); +} + +export function compileConcepts(queryKeywords: string[]): ConceptMatcher { + const concepts = dedupeSynonymConcepts(queryKeywords); + const termsByKey = new Map(); + concepts.forEach((keyword, concept) => { + for (const term of SYNONYM_LOOKUP.get(keyword) ?? [keyword]) { + if (term.length < MIN_PREFIX_LEN) continue; + const head = term.slice(0, MAX_PREFIX_LEN); + const key = headKey(head); + const terms = termsByKey.get(key); + if (!terms) termsByKey.set(key, [{ concept, head }]); + else if (!terms.some((t) => t.concept === concept && t.head === head)) terms.push({ concept, head }); + } + }); + return { concepts, termsByKey }; +} + +/** + * For each concept, the index into SCORED_FIELDS of the strongest field it matches, or -1. A concept is + * tested only until it matched, so `data` is serialized only while a concept is still open. + */ +export function matchConcepts(e: EventEntry, { concepts, termsByKey }: ConceptMatcher): number[] { + const fieldOf: number[] = new Array(concepts.length).fill(-1); + let open = concepts.length; + for (let f = 0; f < SCORED_FIELDS.length && open > 0; f++) { + const text = SCORED_FIELDS[f].text(e); + if (!text) continue; + for (let c = 0; c < concepts.length; c++) { + if (fieldOf[c] < 0 && text.includes(concepts[c])) { + fieldOf[c] = f; + open--; + } + } + if (open === 0 || termsByKey.size === 0) continue; + for (const word of splitWords(text)) { + if (word.length < MIN_PREFIX_LEN) continue; + const terms = termsByKey.get(headKey(word)); + if (!terms) continue; + for (const { concept, head } of terms) { + if (fieldOf[concept] < 0 && (word.length >= head.length ? word.startsWith(head) : head.startsWith(word))) { + fieldOf[concept] = f; + open--; + } + } + if (open === 0) break; + } + } + return fieldOf; +} + type ScoredEvent = { index: number; score: number; @@ -856,6 +943,9 @@ function severityIntentFactor(e: EventEntry, boostFailures: boolean, entityTypeM * non-problem keyword matches the event TYPE, the event is the specifically-named entity and * earns the ENTITY_BOOST (not the benign-Info damp) on a problem query — see severityIntentFactor. * Omitted → no entity boost (old behavior), which is why the existing unit tests are unaffected. + * + * `matcher` is the query compiled by compileConcepts. It is query-invariant, so scoreEvents compiles it + * once for every fetched event; standalone callers (unit tests) omit it and get it compiled here. */ export function scoreEvent( e: EventEntry, @@ -863,49 +953,30 @@ export function scoreEvent( boostFailures = false, semanticTypeScores?: Map, keywordSpecificity?: Map, - precomputedConcepts?: string[], + matcher: ConceptMatcher = compileConcepts(queryKeywords), ): ScoredEvent | null { - const fields: Array<{ name: string; text: string; weight: number }> = [ - { name: 'eventType', text: (e.eventType ?? '').toLowerCase(), weight: FIELD_WEIGHTS.eventType }, - { name: 'message', text: (e.message ?? '').toLowerCase(), weight: FIELD_WEIGHTS.message }, - { name: 'source', text: (e.source ?? '').toLowerCase(), weight: FIELD_WEIGHTS.source }, - { name: 'severity', text: (e.severity ?? '').toLowerCase(), weight: FIELD_WEIGHTS.severity }, - { name: 'data', text: e.data ? JSON.stringify(e.data).toLowerCase() : '', weight: FIELD_WEIGHTS.data }, - ]; - - // Collapse query keywords that share a synonym group to a single concept so two query words - // expressing ONE idea (e.g. "stuck" + "timeout") don't double-count coverage/score against an - // off-topic type (hello_wait_timeout) and outrank the actually-named entity (app_download_started). - // This is query-invariant, so the cross-session loop computes it once and passes it in to avoid - // recomputing the dedupe (and its Set/Map allocations) for every one of potentially thousands of - // fetched events; standalone callers (unit tests) omit it and get the same value computed here. - const concepts = precomputedConcepts ?? dedupeSynonymConcepts(queryKeywords); + // The concepts are the keywords collapsed to one per synonym group, so two query words expressing ONE + // idea (e.g. "stuck" + "timeout") don't double-count coverage/score against an off-topic type + // (hello_wait_timeout) and outrank the actually-named entity (app_download_started). + const { concepts } = matcher; + const fieldOf = matchConcepts(e, matcher); let totalScore = 0; const matched: string[] = []; const bestFields = new Set(); let entityTypeMatch = false; - for (const kw of concepts) { - let kwBestWeight = 0; - let kwBestField = ''; - for (const field of fields) { - if (field.text && keywordMatchesField(field.text, kw)) { - if (field.weight > kwBestWeight) { - kwBestWeight = field.weight; - kwBestField = field.name; - } - } - } - if (kwBestWeight > 0) { - totalScore += kwBestWeight; - matched.push(kw); - bestFields.add(kwBestField); - // A rare, non-problem keyword matching the event TYPE marks this as the named entity. - if (kwBestField === 'eventType' && !isProblemIntentKeyword(kw) - && (keywordSpecificity?.get(kw) ?? Infinity) <= RARE_THRESHOLD) { - entityTypeMatch = true; - } + for (let c = 0; c < concepts.length; c++) { + if (fieldOf[c] < 0) continue; + const kw = concepts[c]; + const field = SCORED_FIELDS[fieldOf[c]]; + totalScore += field.weight; + matched.push(kw); + bestFields.add(field.name); + // A rare, non-problem keyword matching the event TYPE marks this as the named entity. + if (field.name === 'eventType' && !isProblemIntentKeyword(kw) + && (keywordSpecificity?.get(kw) ?? Infinity) <= RARE_THRESHOLD) { + entityTypeMatch = true; } } @@ -940,6 +1011,36 @@ export function scoreEvent( }; } +/** Events scored between two yields to the event loop. */ +const SCORE_YIELD_INTERVAL = 256; + +/** + * Score fetched events against the query, keeping those at or above `minScore`. Yields to the event + * loop every SCORE_YIELD_INTERVAL events: one process serves every tenant, and a cross-session search + * ranks tens of thousands of events, which in one synchronous pass would stall every other request + * (and /health) until it finished. Stops once the caller has gone, since nobody reads the result. + */ +export async function scoreEvents( + events: EventEntry[], + queryKeywords: string[], + boostFailures: boolean, + semanticTypeScores: Map, + keywordSpecificity: Map, + minScore: number, +): Promise> { + const matcher = compileConcepts(queryKeywords); + const scored: Array = []; + for (let i = 0; i < events.length; i++) { + if (i > 0 && i % SCORE_YIELD_INTERVAL === 0) { + await new Promise((resolve) => setImmediate(resolve)); + getCallerSignal()?.throwIfAborted(); + } + const result = scoreEvent(events[i], queryKeywords, boostFailures, semanticTypeScores, keywordSpecificity, matcher); + if (result && result.score >= minScore) scored.push({ ...result, index: i, event: events[i] }); + } + return scored; +} + /** * Spread the top-scored events across distinct sessions so one session's many * same-type events don't monopolize the list — cross-session search otherwise keeps @@ -1066,7 +1167,7 @@ export function registerSearchTools( 'even when it sits inside another event\'s data — pass a sessionId for a complete scan of every event field. ' + '`truncated: true` means incomplete recall: narrow per `recallNote` or use depth="deep".', inputSchema: { - query: z.string().describe('Natural-language description of what to find (e.g. "machine restarted unexpectedly", "app download stuck", "certificate error")'), + query: z.string().max(MAX_SEARCH_QUERY_CHARS).describe('Natural-language description of what to find (e.g. "machine restarted unexpectedly", "app download stuck", "certificate error")'), sessionId: SessionIdSchema.optional().describe('Search within one session (complete scan). Omitted: cross-session search by event type.'), tenantId: z.string().optional().describe(tenantIdDescription(ga, delegated)), depth: z.enum(['fast', 'deep']).optional().default('fast') @@ -1075,7 +1176,7 @@ export function registerSearchTools( .describe('Number of matching events to return (1-50). Defaults: 10 (fast) / 20 (deep).'), minScore: z.coerce.number().min(0).max(1).optional() .describe('Minimum relevance score (0-1). Defaults: 0.1 (fast) / 0.05 (deep). Lower surfaces weaker matches.'), - keywords: z.array(z.string()).optional() + keywords: z.array(z.string().max(MAX_SEARCH_KEYWORD_CHARS)).max(MAX_SEARCH_KEYWORDS).optional() .describe('Additional exact keywords, merged (case-insensitively) with those auto-extracted from the query.'), guaranteedTopRanked: z.coerce.number().min(0).max(50).optional().default(3) .describe('How many top hits (by pure score) stay at the head in rank order, exempt from cross-session diversification (default 3). Set =topK to disable diversification, 0 for maximum session spread. Ignored with sessionId.'), @@ -1091,8 +1192,11 @@ export function registerSearchTools( const topK = args.topK ?? preset.defaultTopK; const minScore = args.minScore ?? preset.defaultMinScore; - // Extract keywords FIRST so we can use them for index-based pre-filtering - const queryKeywords = resolveQueryKeywords(query, keywords); + // Extract keywords FIRST so we can use them for index-based pre-filtering. Everything below + // sees only the first MAX_SCORED_KEYWORDS; the response reports the rest as dropped. + const allKeywords = resolveQueryKeywords(query, keywords); + const queryKeywords = allKeywords.slice(0, MAX_SCORED_KEYWORDS); + const keywordsDropped = allKeywords.length - queryKeywords.length; if (queryKeywords.length === 0) { logSearchZeroHit('search_events', query, { reason: 'no_keywords' }); return toolResultText( @@ -1108,16 +1212,7 @@ export function registerSearchTools( await fetchSessionEvents(sessionId, tenantId, queryKeywords, preset.budget, query, eventTypeIndex, keywordSpecificity); const boostFailures = queryHasProblemIntent(queryKeywords); - // Synonym-collapse the query keywords ONCE — scoreEvent runs per fetched event (potentially - // thousands), and the concept set is identical across all of them. - const concepts = dedupeSynonymConcepts(queryKeywords); - const scored: Array = []; - for (let i = 0; i < events.length; i++) { - const result = scoreEvent(events[i], queryKeywords, boostFailures, semanticTypeScores, keywordSpecificity, concepts); - if (result && result.score >= minScore) { - scored.push({ ...result, index: i, event: events[i] }); - } - } + const scored = await scoreEvents(events, queryKeywords, boostFailures, semanticTypeScores, keywordSpecificity, minScore); scored.sort((a, b) => b.score - a.score); const results = diversifyBySession(scored, topK, 2, guaranteedTopRanked).map((s) => ({ @@ -1153,6 +1248,7 @@ export function registerSearchTools( depth, searchMethod, keywordsUsed: queryKeywords, + ...(keywordsDropped > 0 ? { keywordsDropped, keywordNote: KEYWORDS_DROPPED_NOTE } : {}), sessionsSearchedCount: sessionIds.length, eventsFetched: events.length, eventsMatched: scored.length, @@ -1186,7 +1282,7 @@ export function registerSearchTools( 'because opaque codes embed poorly. For one known code, lookup_error_code is the direct answer (accepts ' + 'decimal, symbol and enforcement-state input too).', inputSchema: { - query: z.string().describe('Natural language search query (e.g. "app download timeout", "TPM not ready", "ESP stuck")'), + query: z.string().max(MAX_SEARCH_QUERY_CHARS).describe('Natural language search query (e.g. "app download timeout", "TPM not ready", "ESP stuck")'), topK: z.coerce.number().min(1).max(20).optional().default(5).describe('Number of results to return (1-20, default 5)'), type: z.enum(['all', 'analyze-rule', 'gather-rule', 'ime-log-pattern', 'error-code']).optional().default('all') .describe('Filter by document type. Default: search all types.'), @@ -1333,9 +1429,9 @@ export function registerSearchTools( 'Rule of thumb: "how does X work?" → search_docs; "why did this session fail?" → search_knowledge or ' + 'get_session_summary.', inputSchema: { - query: z.string().describe('Natural language question (e.g. "how do I deploy the agent with Intune", "where is my data stored", "which roles exist")'), + query: z.string().max(MAX_SEARCH_QUERY_CHARS).describe('Natural language question (e.g. "how do I deploy the agent with Intune", "where is my data stored", "which roles exist")'), topK: z.coerce.number().min(1).max(20).optional().default(5).describe('Number of results to return (1-20, default 5)'), - section: z.string().optional() + section: z.string().max(MAX_DOCS_SECTION_CHARS).optional() .describe('Restrict to one documentation area, matching the bundle\'s top-level folder ' + `(one of: ${docs.sections.join(', ')}). Omit to search everything.`), minScore: z.coerce.number().min(0).max(1).optional().default(0.25) diff --git a/src/McpServer/autopilot-monitor-mcp/src/tools/sessions.ts b/src/McpServer/autopilot-monitor-mcp/src/tools/sessions.ts index 09b7c29e..88be7be8 100644 --- a/src/McpServer/autopilot-monitor-mcp/src/tools/sessions.ts +++ b/src/McpServer/autopilot-monitor-mcp/src/tools/sessions.ts @@ -2,7 +2,7 @@ import { McpServer } from '@modelcontextprotocol/server'; import { z } from 'zod'; import { ApiError, apiFetch, buildQuery, jsonBody, DEFAULT_FIRST_PAGE_SIZE, effectivePageSize, enforceDelegatedTenant, enforceDelegatedTenantForPage, followNextLink, pageSizeForCall, pickGlobalOrTenantPath, scanWithTimeoutFallback } from '../client.js'; import { withToolTelemetry } from '../telemetry.js'; -import { READ_ONLY, MAX_RESULT_SIZE_CHARS, LEAN_EVENT_FIELDS, LEAN_EVENT_OMISSION, leanFieldSelection, SUMMARY_EVENT_FIELDS, toolResultText, SessionIdSchema, isBenignHealthDetectionReport, foreignSuccessRunIds, isCorrectedForeignFailure, tenantIdDescription, pageSizeDescription, CONTINUATION_DESCRIPTION, daysDescription } from './shared.js'; +import { READ_ONLY, MAX_RESULT_SIZE_CHARS, LEAN_EVENT_FIELDS, LEAN_EVENT_OMISSION, leanFieldSelection, SUMMARY_EVENT_FIELDS, toolResultText, SessionIdSchema, isBenignHealthDetectionReport, foreignSuccessRunIds, isCorrectedForeignFailure, tenantIdDescription, pageSizeDescription, CONTINUATION_DESCRIPTION, ContinuationSchema, daysDescription } from './shared.js'; import { toolError } from './error-handler.js'; import { lookupErrorCode } from '../error-code-catalog.js'; import { withEventTypeNote, assertKnownDevicePropertyKeys } from '../resource-catalog.js'; @@ -174,7 +174,7 @@ export function registerSessionTools(server: McpServer, ga: boolean, delegated: ), pageSize: z.coerce.number().int().min(1).max(1000).optional() .describe(pageSizeDescription(DEFAULT_FIRST_PAGE_SIZE, 1000, 'Raise it for full sweeps.')), - continuation: z.string().optional().describe(CONTINUATION_DESCRIPTION), + continuation: ContinuationSchema.optional().describe(CONTINUATION_DESCRIPTION), }, annotations: READ_ONLY, }, @@ -227,7 +227,7 @@ export function registerSessionTools(server: McpServer, ga: boolean, delegated: tenantId: z.string().optional().describe(tenantIdDescription(ga, delegated)), pageSize: z.coerce.number().int().min(1).max(1000).optional() .describe(pageSizeDescription(DEFAULT_FIRST_PAGE_SIZE, 1000, 'Raise it for full sweeps.')), - continuation: z.string().optional().describe(CONTINUATION_DESCRIPTION), + continuation: ContinuationSchema.optional().describe(CONTINUATION_DESCRIPTION), }, annotations: READ_ONLY, }, @@ -410,7 +410,7 @@ export function registerSessionTools(server: McpServer, ga: boolean, delegated: .describe('Projection. Omitted on an UNFILTERED read = lean default "' + LEAN_EVENT_FIELDS + '" without the multi-KB "data" payload (omittedFields says so); omitted on a FILTERED read = full events incl. "data". List "data" for the whole payload or "data." (e.g. "data.errorCode") for single keys. Further keys: eventId, sessionId, tenantId, receivedAt, sentAt, rowKey, originalTimestamp, timestampClamped, causedByTransitionStepIndex, causedBySignalOrdinal.'), pageSize: z.coerce.number().int().min(1).max(1000).optional() .describe(pageSizeDescription(DEFAULT_FIRST_PAGE_SIZE)), - continuation: z.string().optional().describe(CONTINUATION_DESCRIPTION), + continuation: ContinuationSchema.optional().describe(CONTINUATION_DESCRIPTION), }, annotations: READ_ONLY, }, @@ -749,7 +749,7 @@ export function registerSessionTools(server: McpServer, ga: boolean, delegated: overallRisk: z.enum(['low', 'medium', 'high', 'critical']).optional(), pageSize: z.coerce.number().int().min(1).max(1000).optional() .describe(pageSizeDescription(DEFAULT_FIRST_PAGE_SIZE, 1000, 'Raise it for a full exposure audit.')), - continuation: z.string().optional().describe(CONTINUATION_DESCRIPTION), + continuation: ContinuationSchema.optional().describe(CONTINUATION_DESCRIPTION), }, annotations: READ_ONLY, }, diff --git a/src/McpServer/autopilot-monitor-mcp/src/tools/shared.ts b/src/McpServer/autopilot-monitor-mcp/src/tools/shared.ts index da084c7d..753941bf 100644 --- a/src/McpServer/autopilot-monitor-mcp/src/tools/shared.ts +++ b/src/McpServer/autopilot-monitor-mcp/src/tools/shared.ts @@ -372,6 +372,13 @@ export function pageSizeDescription(defaultSize: number, max = 1000, note?: stri export const CONTINUATION_DESCRIPTION = 'The whole nextLink string from the previous response (preferred — it carries the resolved filters) or its opaque continuation value.'; +/** + * continuation: a server-emitted nextLink or its opaque token; real ones stay under 1 KB. The bound + * matters: a pageSize/fields override re-sets params on the caller's link (withQueryOverrides), and + * URLSearchParams.set drops duplicate keys one at a time — quadratic in a forged link's duplicates. + */ +export const ContinuationSchema = z.string().max(4096); + /** days: a trailing window ending now. */ export function daysDescription(defaultDays: number, max: number, note?: string): string { return `Trailing window in days (1-${max}; default ${defaultDays}).${note ? ` ${note}` : ''}`;