feat(demand-control): allow maxCost to accept a dynamic function

.changeset/demand-control-dynamic-max-cost.md

@@ -0,0 +1,18 @@
+---
+'@graphql-hive/gateway-runtime': minor
+---
+
+Allow `maxCost` in `DemandControlPluginOptions` to accept a function for dynamic cost limiting.
+
+Previously, `maxCost` only accepted a static `number`. It now also accepts a synchronous or
+asynchronous function `(payload: DemandControlMaxCostPayload) => MaybePromise<number>`, where
+`payload` contains:
+
+- `operationCost` – the estimated cost of the current subgraph operation
+- `totalCost` – the accumulated cost for the whole request context so far
+- `subgraphName` – the name of the subgraph being executed
+- `executionRequest` – the full execution request object
+
+This lets you implement per-user rate limits, per-subgraph budgets, or any other context-aware
+cost policy. The `DemandControlMaxCostPayload` interface is exported from the package for use
+when typing your `maxCost` function.

packages/runtime/src/plugins/useDemandControl.ts

@@ -1,9 +1,12 @@
 import { EMPTY_OBJECT } from '@graphql-tools/delegate';
 import {
   createGraphQLError,
+  ExecutionRequest,
   isAsyncIterable,
   mapAsyncIterator,
+  MaybePromise,
 } from '@graphql-tools/utils';
+import { handleMaybePromise } from '@whatwg-node/promise-helpers';
 import { getNodeEnv } from '~internal/env';
 import {
   FieldNode,
@@ -15,13 +18,37 @@ import {
 import { GatewayPlugin } from '../types';
 import { createCalculateCost } from './demand-control/calculateCost';
 
+export interface DemandControlMaxCostPayload {
+  /**
+   * The estimated cost of the current subgraph operation.
+   */
+  operationCost: number;
+  /**
+   * The total estimated cost accumulated for the whole request context so far.
+   */
+  totalCost: number;
+  /**
+   * The name of the subgraph being executed.
+   */
+  subgraphName: string;
+  /**
+   * The execution request being processed.
+   */
+  executionRequest: ExecutionRequest<any, any>;
+}
+
 export interface DemandControlPluginOptions {
   /**
-   * 	The maximum cost of an accepted operation. An operation with a higher cost than this is rejected.
-   *  If not provided, no maximum cost is enforced.
-   *  @default Infinity
+   * The maximum cost of an accepted operation. An operation with a higher cost than this is rejected.
+   * Can be a static number or a function (sync or async) that receives cost details and returns the
+   * maximum cost allowed. Use a function to implement dynamic cost limits based on the context,
+   * subgraph, or estimated cost.
+   * If not provided, no maximum cost is enforced.
+   * @default Infinity
    */
-  maxCost?: number;
+  maxCost?:
+    | number
+    | ((payload: DemandControlMaxCostPayload) => MaybePromise<number>);
   /**
    * The assumed maximum size of a list for fields that return lists.
    * @default 0
@@ -67,9 +94,16 @@ export function useDemandControl<TContext extends Record<string, any>>({
     fieldCost,
     typeCost,
   });
+  const maxCostFn =
+    maxCost == null
+      ? null
+      : typeof maxCost === 'function'
+        ? maxCost
+        : () => maxCost;
   const costByContextMap = new WeakMap<any, number>();
+  const resolvedMaxCostByContextMap = new WeakMap<any, number>();
   return {
-    onSubgraphExecute({ subgraph, executionRequest, log }) {
+    onSubgraphExecute({ subgraph, subgraphName, executionRequest, log }) {
       if (!subgraph) {
         return;
       }
@@ -92,24 +126,44 @@ export function useDemandControl<TContext extends Record<string, any>>({
         },
         '[useDemandControl]',
       );
-      if (maxCost != null && costByContext > maxCost) {
-        throw createGraphQLError(
-          `Operation estimated cost ${costByContext} exceeded configured maximum ${maxCost}`,
-          {
-            extensions: {
-              code: 'COST_ESTIMATED_TOO_EXPENSIVE',
-              cost: {
-                estimated: costByContext,
-                max: maxCost,
-              },
-            },
+      if (maxCostFn != null) {
+        return handleMaybePromise(
+          () =>
+            maxCostFn({
+              operationCost,
+              totalCost: costByContext,
+              subgraphName,
+              executionRequest,
+            }),
+          (resolvedMaxCost) => {
+            if (executionRequest.context) {
+              resolvedMaxCostByContextMap.set(
+                executionRequest.context,
+                resolvedMaxCost,
+              );
+            }
+            if (costByContext > resolvedMaxCost) {
+              throw createGraphQLError(
+                `Operation estimated cost ${costByContext} exceeded configured maximum ${resolvedMaxCost}`,
+                {
+                  extensions: {
+                    code: 'COST_ESTIMATED_TOO_EXPENSIVE',
+                    cost: {
+                      estimated: costByContext,
+                      max: resolvedMaxCost,
+                    },
+                  },
+                },
+              );
+            }
           },
         );
       }
     },
     onExecutionResult({ result, setResult, context }) {
       if (includeExtensionMetadata) {
         const costByContext = costByContextMap.get(context) || 0;
+        const resolvedMaxCost = resolvedMaxCostByContextMap.get(context);
         if (isAsyncIterable(result)) {
           setResult(
             mapAsyncIterator(result, (value) => ({
@@ -118,7 +172,7 @@ export function useDemandControl<TContext extends Record<string, any>>({
                 ...(value.extensions || {}),
                 cost: {
                   estimated: costByContext,
-                  max: maxCost,
+                  ...(resolvedMaxCost != null ? { max: resolvedMaxCost } : {}),
                 },
               },
             })),
@@ -130,7 +184,7 @@ export function useDemandControl<TContext extends Record<string, any>>({
               ...(result?.extensions || {}),
               cost: {
                 estimated: costByContext,
-                max: maxCost,
+                ...(resolvedMaxCost != null ? { max: resolvedMaxCost } : {}),
               },
             },
           });

packages/runtime/src/types.ts

@@ -56,6 +56,8 @@ export type { UnifiedGraphHandler, UnifiedGraphPlugin };
 export type { TransportEntryAdditions, UnifiedGraphConfig };
 export type { CircuitBreakerConfiguration };
 
+export type { DemandControlMaxCostPayload } from './plugins/useDemandControl';
+
 export type GatewayConfig<
   TContext extends Record<string, any> = Record<string, any>,
 > =

packages/runtime/tests/demand-control.test.ts

@@ -1163,6 +1163,108 @@ describe('Demand Control', () => {
           });
         }
       });
+
+      it('accepts a function for maxCost that receives cost payload', async () => {
+        const subgraph = buildSubgraphSchema({
+          typeDefs: parse(/* GraphQL */ `
+            type Query {
+              foo: Foo
+            }
+
+            type Foo {
+              id: ID
+            }
+          `),
+          resolvers: {
+            Query: {
+              foo: async () => ({ id: 'foo' }),
+            },
+          },
+        });
+        const costPayloads: Parameters<
+          NonNullable<
+            Exclude<
+              Parameters<typeof useDemandControl>[0]['maxCost'],
+              number | undefined
+            >
+          >
+        >[] = [];
+        await using gateway = createTestGateway(mode, subgraph, {
+          includeExtensionMetadata: true,
+          maxCost: (payload) => {
+            costPayloads.push([payload]);
+            return Infinity; // never reject
+          },
+        });
+        const query = /* GraphQL */ `
+          query FooQuery {
+            foo {
+              id
+            }
+          }
+        `;
+        const response = await gateway.fetch('http://localhost:4000/graphql', {
+          method: 'POST',
+          headers: {
+            'Content-Type': 'application/json',
+          },
+          body: JSON.stringify({ query }),
+        });
+        const result = await response.json();
+        expect(result.errors).toBeUndefined();
+        expect(costPayloads.length).toBeGreaterThan(0);
+        const [payload] = costPayloads[0]!;
+        expect(typeof payload.operationCost).toBe('number');
+        expect(typeof payload.totalCost).toBe('number');
+        expect(typeof payload.subgraphName).toBe('string');
+        expect(payload.executionRequest).toBeDefined();
+      });
+
+      it('rejects when async maxCost function returns a limit below total cost', async () => {
+        const subgraph = buildSubgraphSchema({
+          typeDefs: parse(/* GraphQL */ `
+            type Query {
+              foo: Foo
+            }
+
+            type Foo {
+              id: ID
+            }
+          `),
+          resolvers: {
+            Query: {
+              foo: async () => ({ id: 'foo' }),
+            },
+          },
+        });
+        await using gateway = createTestGateway(mode, subgraph, {
+          includeExtensionMetadata: true,
+          maxCost: async ({ totalCost }) => {
+            // simulate async work (e.g. look up per-user rate limit)
+            await new Promise((resolve) => setTimeout(resolve, 0));
+            return totalCost - 1; // always below totalCost → always reject
+          },
+        });
+        const query = /* GraphQL */ `
+          query FooQuery {
+            foo {
+              id
+            }
+          }
+        `;
+        const response = await gateway.fetch('http://localhost:4000/graphql', {
+          method: 'POST',
+          headers: {
+            'Content-Type': 'application/json',
+          },
+          body: JSON.stringify({ query }),
+        });
+        const result = await response.json();
+        expect(result.errors).toBeDefined();
+        expect(result.errors[0].extensions?.code).toBe(
+          'COST_ESTIMATED_TOO_EXPENSIVE',
+        );
+      });
     });
   });
 });