add timeout mcp poll

This commit is contained in:
larssand
2026-07-02 21:17:50 +02:00
parent f45c275e20
commit 97d017e43c
5 changed files with 59 additions and 21 deletions

View File

@@ -176,24 +176,30 @@ def build_status(
probe_status = client.probe()
discovery_store = FieldDiscoveryStore(history_path) if history_path else None
catalog_fields_total = 0
def budget_exceeded() -> bool:
return time.monotonic() >= poll_deadline
def skipped_status(stream_id: str, stream_name: str) -> dict[str, object]:
return {
"stream_id": stream_id,
"stream_name": stream_name,
"source": "graylog_mcp",
"events_fetched": 0,
"aggregate_events": 0,
"pages": 0,
"partial": True,
"error": "skipped_poll_budget",
"truncated": False,
"latest_event_time": "",
"raw_sample_limit": min(raw_sample_events, 10_000) if use_aggregate else max_events_per_stream,
"health_detail": "Skipped because the MCP poll time budget was reached before this stream.",
}
for stream_config in stream_configs:
stream_id = str(stream_config["id"])
stream_name = str(stream_config.get("title", "") or stream_titles.get(stream_id) or stream_id)
if time.monotonic() >= poll_deadline:
stream_statuses.append({
"stream_id": stream_id,
"stream_name": stream_name,
"source": "graylog_mcp",
"events_fetched": 0,
"aggregate_events": 0,
"pages": 0,
"partial": True,
"error": "skipped_poll_budget",
"truncated": False,
"latest_event_time": "",
"raw_sample_limit": min(raw_sample_events, 10_000) if use_aggregate else max_events_per_stream,
"health_detail": "Skipped because the MCP poll time budget was reached before this stream.",
})
if budget_exceeded():
stream_statuses.append(skipped_status(stream_id, stream_name))
continue
profile = stream_profiles.get(stream_id)
profile_fields = (
@@ -205,17 +211,20 @@ def build_status(
*tuple(str(getattr(relation, "left", "")) for relation in getattr(profile, "relationship_fields", ())),
*tuple(str(getattr(relation, "right", "")) for relation in getattr(profile, "relationship_fields", ())),
) if profile else ()
if discovery_store:
if discovery_store and not budget_exceeded():
try:
catalog_fields_total += discovery_store.ingest_catalog(stream_id, stream_name, _graylog_fields_from_result(client.call_tool("list_fields", {"streams": [stream_id]})))
except RuntimeError:
pass
aggregate_status: dict[str, object] = {}
if use_aggregate:
aggregate_status = GraylogAggregateSource(client, stream_id, str(runtime_values.get("graylog_query", "*"))).fetch_count(range_seconds=range_seconds, probe_status=probe_status)
if use_aggregate and not budget_exceeded():
aggregate_status = GraylogAggregateSource(client, stream_id, str(runtime_values.get("graylog_query", "*"))).fetch_count(range_seconds=range_seconds, probe_status=probe_status, deadline_monotonic=poll_deadline)
aggregate_events_total += int(aggregate_status.get("aggregate_events", 0) or 0)
raw_limit = min(raw_sample_events, 10_000) if use_aggregate else max_events_per_stream
stream_events, stream_status = GraylogStreamSource(client, stream_id, str(runtime_values.get("graylog_query", "*")), str(runtime_values.get("graylog_field_mapping", "")), stream_name, profile_fields).fetch(max_events=raw_limit, range_seconds=range_seconds, probe_status=probe_status)
if budget_exceeded():
stream_statuses.append({**skipped_status(stream_id, stream_name), **aggregate_status})
continue
stream_events, stream_status = GraylogStreamSource(client, stream_id, str(runtime_values.get("graylog_query", "*")), str(runtime_values.get("graylog_field_mapping", "")), stream_name, profile_fields).fetch(max_events=raw_limit, range_seconds=range_seconds, probe_status=probe_status, deadline_monotonic=poll_deadline)
events.extend(stream_events)
stream_statuses.append({"stream_id": stream_id, "stream_name": stream_name, **aggregate_status, **stream_status, "raw_sample_limit": raw_limit})
sample_limited_streams = [item for item in stream_statuses if item.get("truncated") and use_aggregate]