diff --git a/src/kolide_mcp/server.py b/src/kolide_mcp/server.py index 291d5cc..87277e1 100644 --- a/src/kolide_mcp/server.py +++ b/src/kolide_mcp/server.py @@ -99,30 +99,65 @@ def _extract_device_id(record: dict[str, Any]) -> str | None: async def _enrich_device_owners(records: list[dict[str, Any]]) -> None: - """Enrich records in-place with owner_name and owner_email fields.""" - owner_cache: dict[str, dict[str, str]] = {} + """Enrich records in-place with owner_name and owner_email fields. - for record in records: - dev_id = _extract_device_id(record) - if not dev_id: + Fetches device and person data concurrently in two batched rounds + instead of issuing sequential calls per record. + """ + import asyncio + + unique_device_ids = { + did for r in records if (did := _extract_device_id(r)) + } + if not unique_device_ids: + return + + semaphore = asyncio.Semaphore(10) + + async def _fetch_device(dev_id: str) -> tuple[str, dict[str, Any] | None]: + async with semaphore: + try: + return (dev_id, await client.request("GET", f"/devices/{dev_id}")) + except KolideAPIError: + return (dev_id, None) + + device_results = await asyncio.gather( + *[_fetch_device(did) for did in unique_device_ids] + ) + + person_to_devices: dict[str, list[str]] = {} + for dev_id, device in device_results: + if device is None: continue + person_id = device.get("registered_owner_info", {}).get("identifier") + if person_id: + person_to_devices.setdefault(str(person_id), []).append(dev_id) - if dev_id not in owner_cache: + async def _fetch_person(person_id: str) -> tuple[str, dict[str, Any] | None]: + async with semaphore: try: - device = await client.request("GET", f"/devices/{dev_id}") - person_id = device.get("registered_owner_info", {}).get("identifier") - if person_id: - person = await client.request("GET", f"/people/{person_id}") - owner_cache[dev_id] = { - "owner_name": person.get("name", ""), - "owner_email": person.get("email", ""), - } - else: - owner_cache[dev_id] = {} + return (person_id, await client.request("GET", f"/people/{person_id}")) except KolideAPIError: - owner_cache[dev_id] = {} + return (person_id, None) + + person_results = await asyncio.gather( + *[_fetch_person(pid) for pid in person_to_devices] + ) + + owner_cache: dict[str, dict[str, str]] = {} + for person_id, person in person_results: + if person is None: + continue + owner_info = { + "owner_name": person.get("name", ""), + "owner_email": person.get("email", ""), + } + for dev_id in person_to_devices[person_id]: + owner_cache[dev_id] = owner_info - if owner_cache[dev_id]: + for record in records: + dev_id = _extract_device_id(record) + if dev_id and dev_id in owner_cache: record.update(owner_cache[dev_id])