fix: refuse Prometheus endpoint responses over 8 MB

.changeset/mcp-promql-response-cap.md

@@ -0,0 +1,5 @@
+---
+'@hyperdx/api': patch
+---
+
+fix: Refuse Prometheus endpoint responses over 8 MB in MCP PromQL label lookups and queries, instead of buffering them whole

packages/api/src/controllers/__tests__/promql.test.ts

@@ -1,4 +1,8 @@
-import { buildPrometheusRequest, runPromqlQuery } from '@/controllers/promql';
+import {
+  buildPrometheusRequest,
+  PROMETHEUS_MAX_RESPONSE_BYTES,
+  runPromqlQuery,
+} from '@/controllers/promql';
 
 const mockFetch = jest.mocked(global.fetch);
 
@@ -182,4 +186,49 @@ describe('runPromqlQuery', () => {
     ).rejects.toMatchObject({ category: 'user' });
     expect(mockFetch).not.toHaveBeenCalled();
   });
+
+  describe('response size cap', () => {
+    it('reports a response over the cap as a server error naming the cap', async () => {
+      mockFetch.mockResolvedValueOnce(
+        new Response('{', {
+          headers: {
+            'content-length': String(PROMETHEUS_MAX_RESPONSE_BYTES + 1),
+          },
+        }),
+      );
+
+      await expect(runPromqlQuery(prometheus, range)).rejects.toMatchObject({
+        message:
+          'Prometheus at http://prom:9090/api/v1/query_range returned more than 8 MB. Narrow the query or time range.',
+        category: 'server',
+      });
+    });
+
+    it('reads a chunked body under the cap as before', async () => {
+      const json = JSON.stringify({ status: 'success', data: matrix });
+      const [head, tail] = [json.slice(0, 10), json.slice(10)];
+      const body = new ReadableStream<Uint8Array>({
+        start(controller) {
+          controller.enqueue(new TextEncoder().encode(head));
+          controller.enqueue(new TextEncoder().encode(tail));
+          controller.close();
+        },
+      });
+      mockFetch.mockResolvedValueOnce(new Response(body));
+
+      await expect(runPromqlQuery(prometheus, range)).resolves.toEqual(matrix);
+    });
+
+    it('reports a body that is not JSON by its HTTP status', async () => {
+      mockFetch.mockResolvedValueOnce(
+        new Response('<html>Bad gateway</html>', { status: 502 }),
+      );
+
+      await expect(runPromqlQuery(prometheus, range)).rejects.toMatchObject({
+        message:
+          'Prometheus at http://prom:9090/api/v1/query_range returned HTTP 502',
+        category: 'server',
+      });
+    });
+  });
 });

packages/api/src/controllers/promql.ts

@@ -4,6 +4,7 @@ import {
 } from '@hyperdx/common-utils/dist/clickhouse';
 import { getCurrentQueryAttribution } from '@hyperdx/common-utils/dist/clickhouse/node';
 import type { PromqlQuery } from '@hyperdx/common-utils/dist/core/promqlChart';
+import { parseJSON } from '@hyperdx/common-utils/dist/core/utils';
 import { SourceKind, type SQLInterval } from '@hyperdx/common-utils/dist/types';
 import { z } from 'zod';
 
@@ -20,12 +21,28 @@ import {
   queryRangeViaTableFunction,
 } from '@/controllers/timeseriesEngine';
 import { withSpan } from '@/utils/instrumentation';
+import {
+  readBoundedBody,
+  ResponseTooLargeError,
+} from '@/utils/readBoundedBody';
 
 type PromqlBackend = 'prometheus' | 'clickhouse';
 
 const PROMETHEUS_FETCH_TIMEOUT_MS = 30_000;
 const PROMETHEUS_CH_TIMEOUT_MS = 30_000;
 export const PROMETHEUS_MAX_RESOLUTION = 11_000;
+/**
+ * Largest Prometheus response body fetchPrometheusData will read into memory.
+ *
+ * Queries to ClickHouse are already capped by `max_result_rows`
+ * (PROMETHEUS_MAX_RESULT_ROWS). An external Prometheus has no equivalent:
+ * `/query` and `/query_range` take no row limit, and older servers ignore
+ * `limit` on the label APIs. So the only way to bound the memory a response
+ * can take is to stop reading at a byte limit. 8 MB fits roughly
+ * PROMETHEUS_MAX_RESULT_ROWS label values, so both backends allow about the
+ * same amount of data.
+ */
+export const PROMETHEUS_MAX_RESPONSE_BYTES = 8 * 1024 * 1024;
 
 // ClickHouse serves the Prometheus HTTP API under this prefix on its main HTTP
 // port when the `prometheus_api_v1` handler is configured (26.6+).
@@ -176,9 +193,20 @@ async function fetchPrometheusData<T extends z.ZodTypeAny>(
     );
   }
 
-  const body = prometheusBodySchema.safeParse(
-    await resp.json().catch(() => undefined),
-  );
+  let text: string;
+  try {
+    text = await readBoundedBody(resp, PROMETHEUS_MAX_RESPONSE_BYTES);
+  } catch (err) {
+    const tooLargeMessage = `Prometheus at ${redactedTarget} returned more than ${PROMETHEUS_MAX_RESPONSE_BYTES / (1024 * 1024)} MB. Narrow the query or time range.`;
+    const genericMessage = `Failed to read the response from Prometheus at ${redactedTarget}`;
+    throw new PromQLQueryError(
+      err instanceof ResponseTooLargeError ? tooLargeMessage : genericMessage,
+      'server',
+    );
+  }
+  // A body that is not JSON (a proxy's HTML error page, say) is reported by
+  // its HTTP status below.
+  const body = prometheusBodySchema.safeParse(parseJSON(text));
   if (!resp.ok || !body.success || body.data.status !== 'success') {
     const isUserError =
       resp.status >= 400 &&

packages/api/src/mcp/__tests__/promqlQueryTool.int.test.ts

@@ -1,6 +1,7 @@
 import { Metadata } from '@hyperdx/common-utils/dist/core/metadata';
 import { SourceKind } from '@hyperdx/common-utils/dist/types';
 
+import { PROMETHEUS_MAX_RESPONSE_BYTES } from '@/controllers/promql';
 import {
   DEFAULT_DATABASE,
   dropTimeSeriesTable,
@@ -307,6 +308,27 @@ describe('MCP clickstack_promql', () => {
     expect(getFirstText(result)).not.toContain('secret');
   });
 
+  it('reports a response too large to buffer as a server error', async () => {
+    mockFetch.mockResolvedValueOnce(
+      prometheusResponse(
+        { resultType: 'matrix', result: [] },
+        {
+          headers: {
+            'content-length': String(PROMETHEUS_MAX_RESPONSE_BYTES + 1),
+          },
+        },
+      ),
+    );
+    const source = await createPrometheusSource();
+
+    const result = await promql({ sourceId: source._id.toString() });
+
+    expect(result.isError).toBe(true);
+    expect(getFirstText(result)).toContain(
+      'returned more than 8 MB. Narrow the query or time range.',
+    );
+  });
+
   it('rejects a result a chart cannot draw as a user error', async () => {
     mockFetch.mockResolvedValueOnce(
       prometheusResponse({ resultType: 'string', result: [START, 'hello'] }),

packages/api/src/mcp/__tests__/promqlSources.int.test.ts

@@ -2,6 +2,7 @@ import { SourceKind } from '@hyperdx/common-utils/dist/types';
 import { Client } from '@modelcontextprotocol/sdk/client/index.js';
 
 import * as config from '@/config';
+import { PROMETHEUS_MAX_RESPONSE_BYTES } from '@/controllers/promql';
 import {
   DEFAULT_DATABASE,
   dropTimeSeriesTable,
@@ -18,8 +19,8 @@ const mockFetch = jest.mocked(global.fetch);
 
 const TABLE = 'mcp_promql_sources';
 
-const prometheusResponse = (data: unknown) =>
-  new Response(JSON.stringify({ status: 'success', data }));
+const prometheusResponse = (data: unknown, init?: ResponseInit) =>
+  new Response(JSON.stringify({ status: 'success', data }), init);
 
 describe('MCP PromQL sources', () => {
   const server = getServer();
@@ -228,6 +229,26 @@ describe('MCP PromQL sources', () => {
       expect(output.error).toContain('http://prom.example.com/api/v1/');
       expect(output.error).not.toContain('secret');
     });
+
+    it('still returns the source config when a label list is too large to buffer', async () => {
+      const source = await createPrometheusSource();
+      mockFetch.mockImplementation(async () =>
+        prometheusResponse(['up'], {
+          headers: {
+            'content-length': String(PROMETHEUS_MAX_RESPONSE_BYTES + 1),
+          },
+        }),
+      );
+
+      const result = await callTool(client, 'clickstack_describe_source', {
+        sourceId: source._id.toString(),
+      });
+
+      expect(result.isError).toBeFalsy();
+      const output = JSON.parse(getFirstText(result));
+      expect(output.partial).toBe(true);
+      expect(output.error).toContain('returned more than 8 MB');
+    });
   });
 
   describe('with PromQL disabled', () => {

packages/api/src/utils/__tests__/readBoundedBody.test.ts

@@ -0,0 +1,77 @@
+import {
+  readBoundedBody,
+  ResponseTooLargeError,
+} from '@/utils/readBoundedBody';
+
+const MAX = 4 * 1024;
+
+const streamOf = (
+  pull: (controller: ReadableStreamDefaultController<Uint8Array>) => void,
+  onCancel: () => void,
+) =>
+  new ReadableStream<Uint8Array>({
+    pull,
+    cancel() {
+      onCancel();
+    },
+  });
+
+describe('readBoundedBody', () => {
+  it('refuses a declared Content-Length over the cap without reading the body', async () => {
+    let cancelled = false;
+    const body = streamOf(
+      c => c.enqueue(new TextEncoder().encode('{')),
+      () => (cancelled = true),
+    );
+    const resp = new Response(body, {
+      headers: { 'content-length': String(MAX + 1) },
+    });
+
+    await expect(readBoundedBody(resp, MAX)).rejects.toBeInstanceOf(
+      ResponseTooLargeError,
+    );
+    expect(cancelled).toBe(true);
+    expect(resp.bodyUsed).toBe(true);
+  });
+
+  it('stops reading an undeclared-length body at the cap', async () => {
+    const chunk = new Uint8Array(1024).fill(0x20);
+    let pulled = 0;
+    let cancelled = false;
+    const body = streamOf(
+      c => {
+        pulled++;
+        c.enqueue(chunk);
+      },
+      () => (cancelled = true),
+    );
+
+    await expect(
+      readBoundedBody(new Response(body), MAX),
+    ).rejects.toBeInstanceOf(ResponseTooLargeError);
+    expect(cancelled).toBe(true);
+    // One chunk per KB, plus the one that crossed the cap and the stream's
+    // read-ahead.
+    expect(pulled).toBeLessThanOrEqual(MAX / chunk.byteLength + 2);
+  });
+
+  it('joins a chunked body under the cap into one string', async () => {
+    const text = 'héllo wörld';
+    const bytes = new TextEncoder().encode(text);
+    const body = new ReadableStream<Uint8Array>({
+      start(c) {
+        c.enqueue(bytes.slice(0, 3));
+        c.enqueue(bytes.slice(3));
+        c.close();
+      },
+    });
+
+    await expect(readBoundedBody(new Response(body), MAX)).resolves.toBe(text);
+  });
+
+  it('returns an empty string for a response with no body', async () => {
+    await expect(
+      readBoundedBody(new Response(null, { status: 204 }), MAX),
+    ).resolves.toBe('');
+  });
+});

packages/api/src/utils/readBoundedBody.ts

@@ -0,0 +1,41 @@
+/** Thrown by readBoundedBody when a response is larger than its cap. */
+export class ResponseTooLargeError extends Error {
+  constructor(maxBytes: number) {
+    super(`Response body exceeds ${maxBytes} bytes`);
+    this.name = 'ResponseTooLargeError';
+  }
+}
+
+/**
+ * Reads a fetch Response body as text without ever holding more than
+ * `maxBytes` of it. A declared Content-Length over the cap is refused before
+ * any of the body is read; an undeclared (chunked) body is read until the cap
+ * and then cancelled. Use it for bodies from hosts we do not control, where
+ * `resp.text()` / `resp.json()` would buffer whatever the server sends.
+ */
+export async function readBoundedBody(
+  resp: Response,
+  maxBytes: number,
+): Promise<string> {
+  const declared = Number(resp.headers.get('content-length'));
+  if (declared > maxBytes) {
+    await resp.body?.cancel();
+    throw new ResponseTooLargeError(maxBytes);
+  }
+  if (!resp.body) return '';
+
+  const reader = resp.body.getReader();
+  const chunks: Uint8Array[] = [];
+  let received = 0;
+  for (;;) {
+    const { done, value } = await reader.read();
+    if (done) break;
+    received += value.byteLength;
+    if (received > maxBytes) {
+      await reader.cancel();
+      throw new ResponseTooLargeError(maxBytes);
+    }
+    chunks.push(value);
+  }
+  return Buffer.concat(chunks).toString('utf8');
+}