add timeout mcp poll
This commit is contained in:
@@ -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]
|
||||
|
||||
Reference in New Issue
Block a user