From f18ba2ba8cbaebf12beb763facdb678f3f708bde Mon Sep 17 00:00:00 2001 From: Travis Fischer Date: Tue, 10 Jun 2025 10:43:30 +0700 Subject: [PATCH] feat: add caching for mcp origin tool calls --- .../src/__snapshots__/http-e2e.test.ts.snap | 12 ++++ apps/e2e/src/http-fixtures.ts | 39 ++++++++++- apps/e2e/src/mcp-e2e.test.ts | 2 +- apps/gateway/src/app.ts | 1 + apps/gateway/src/lib/get-request-cache-key.ts | 8 +-- .../is-cache-control-publicly-cacheable.ts | 19 ++++++ .../src/lib/resolve-http-edge-request.ts | 5 +- .../src/lib/resolve-origin-tool-call.ts | 68 ++++++++++++++++--- apps/gateway/src/lib/update-origin-request.ts | 8 ++- packages/types/src/openapi.d.ts | 3 +- packages/types/src/tools.ts | 15 +++- readme.md | 5 +- 12 files changed, 159 insertions(+), 26 deletions(-) create mode 100644 apps/gateway/src/lib/is-cache-control-publicly-cacheable.ts diff --git a/apps/e2e/src/__snapshots__/http-e2e.test.ts.snap b/apps/e2e/src/__snapshots__/http-e2e.test.ts.snap index af0552d5..8547fd58 100644 --- a/apps/e2e/src/__snapshots__/http-e2e.test.ts.snap +++ b/apps/e2e/src/__snapshots__/http-e2e.test.ts.snap @@ -197,3 +197,15 @@ et est aut quod aut provident voluptas autem voluptas", "userId": 1, } `; + +exports[`Bypass caching > 2.4: GET @dev/test-basic-openapi@fc856666/get_post?postId=9 1`] = ` +{ + "body": "consectetur animi nesciunt iure dolore +enim quia ad +veniam autem ut quam aut nobis +et est aut quod aut provident voluptas autem voluptas", + "id": 9, + "title": "nesciunt iure omnis dolorem tempora et accusantium", + "userId": 1, +} +`; diff --git a/apps/e2e/src/http-fixtures.ts b/apps/e2e/src/http-fixtures.ts index a14a07a2..15bf850a 100644 --- a/apps/e2e/src/http-fixtures.ts +++ b/apps/e2e/src/http-fixtures.ts @@ -246,7 +246,7 @@ export const fixtureSuites: E2ETestFixtureSuite[] = [ } }, { - // ensure we bypass the cache for requests with `pragma: no-cache` + // ensure we bypass the cache for requests with `cache-control: no-cache` path: '@dev/test-basic-openapi@fc856666/getPost?postId=9', request: { headers: { @@ -285,6 +285,19 @@ export const fixtureSuites: E2ETestFixtureSuite[] = [ 'cf-cache-status': 'BYPASS' } } + }, + { + path: '@dev/test-basic-openapi@fc856666/get_post?postId=9', + request: { + headers: { + 'cache-control': 'private, max-age=3600, must-revalidate' + } + }, + response: { + headers: { + 'cf-cache-status': 'BYPASS' + } + } } ] }, @@ -297,6 +310,10 @@ export const fixtureSuites: E2ETestFixtureSuite[] = [ // first request to ensure the cache is populated path: '@dev/test-basic-openapi@fc856666/getPost', request: { + headers: { + 'cache-control': + 'public, max-age=3600, s-maxage=3600, stale-while-revalidate=3600' + }, searchParams: { postId: 9 } @@ -305,6 +322,12 @@ export const fixtureSuites: E2ETestFixtureSuite[] = [ { // second request should hit the cache path: '@dev/test-basic-openapi@fc856666/getPost?postId=9', + request: { + headers: { + 'cache-control': + 'public, max-age=3600, s-maxage=3600, stale-while-revalidate=3600' + } + }, response: { headers: { 'cf-cache-status': 'HIT' @@ -312,8 +335,14 @@ export const fixtureSuites: E2ETestFixtureSuite[] = [ } }, { - // normalized request should also hit the cache + // normalized request with different path should also hit the cache path: '@dev/test-basic-openapi@fc856666/get_post?postId=9', + request: { + headers: { + 'cache-control': + 'public, max-age=3600, s-maxage=3600, stale-while-revalidate=3600' + } + }, response: { headers: { 'cf-cache-status': 'HIT' @@ -331,6 +360,9 @@ export const fixtureSuites: E2ETestFixtureSuite[] = [ // first request to ensure the cache is populated path: '@dev/test-basic-openapi@fc856666/get_post', request: { + headers: { + 'cache-control': 'public, s-max-age=3600' + }, method: 'POST', json: { postId: 13 @@ -342,6 +374,9 @@ export const fixtureSuites: E2ETestFixtureSuite[] = [ path: '@dev/test-basic-openapi@fc856666/get_post', request: { method: 'POST', + headers: { + 'cache-control': 'public, s-max-age=3600' + }, json: { postId: 13 } diff --git a/apps/e2e/src/mcp-e2e.test.ts b/apps/e2e/src/mcp-e2e.test.ts index 5fb8bd97..8d1557c0 100644 --- a/apps/e2e/src/mcp-e2e.test.ts +++ b/apps/e2e/src/mcp-e2e.test.ts @@ -53,7 +53,7 @@ for (const [i, fixtureSuite] of fixtureSuites.entries()) { // eslint-disable-next-line no-loop-func async () => { const { tools } = await client.listTools() - console.log('tools', tools) + // console.log('tools', tools) expect(tools.map((t) => t.name)).toContain(fixture.request.name) const result = await client.callTool({ diff --git a/apps/gateway/src/app.ts b/apps/gateway/src/app.ts index 743270b1..c44b8feb 100644 --- a/apps/gateway/src/app.ts +++ b/apps/gateway/src/app.ts @@ -76,6 +76,7 @@ app.all(async (ctx) => { deployment: resolvedEdgeRequest.deployment, consumer: resolvedEdgeRequest.consumer, pricingPlan: resolvedEdgeRequest.pricingPlan, + cacheControl: resolvedEdgeRequest.cacheControl, sessionId: ctx.get('sessionId')!, ip: ctx.get('ip'), env: ctx.env, diff --git a/apps/gateway/src/lib/get-request-cache-key.ts b/apps/gateway/src/lib/get-request-cache-key.ts index 375e997a..d9346696 100644 --- a/apps/gateway/src/lib/get-request-cache-key.ts +++ b/apps/gateway/src/lib/get-request-cache-key.ts @@ -1,6 +1,7 @@ import { hashObject, sha256 } from '@agentic/platform-core' import contentType from 'fast-content-type-parse' +import { isCacheControlPubliclyCacheable } from './is-cache-control-publicly-cacheable' import { normalizeUrl } from './normalize-url' // TODO: what is a reasonable upper bound for hashing the POST body size? @@ -16,11 +17,8 @@ export async function getRequestCacheKey( } const cacheControl = request.headers.get('cache-control') - if (cacheControl) { - const directives = new Set(cacheControl.split(',').map((s) => s.trim())) - if (directives.has('no-store') || directives.has('no-cache')) { - return - } + if (!isCacheControlPubliclyCacheable(cacheControl)) { + return } if (request.method === 'POST' || request.method === 'PUT') { diff --git a/apps/gateway/src/lib/is-cache-control-publicly-cacheable.ts b/apps/gateway/src/lib/is-cache-control-publicly-cacheable.ts new file mode 100644 index 00000000..149ebda5 --- /dev/null +++ b/apps/gateway/src/lib/is-cache-control-publicly-cacheable.ts @@ -0,0 +1,19 @@ +export function isCacheControlPubliclyCacheable( + cacheControl?: string | null +): boolean { + if (!cacheControl) { + return false + } + + const directives = new Set(cacheControl.split(',').map((s) => s.trim())) + if ( + directives.has('no-store') || + directives.has('no-cache') || + directives.has('private') || + !directives.has('public') + ) { + return false + } + + return true +} diff --git a/apps/gateway/src/lib/resolve-http-edge-request.ts b/apps/gateway/src/lib/resolve-http-edge-request.ts index 907660e5..906781c4 100644 --- a/apps/gateway/src/lib/resolve-http-edge-request.ts +++ b/apps/gateway/src/lib/resolve-http-edge-request.ts @@ -19,6 +19,7 @@ export type ResolvedHttpEdgeRequest = { tool: Tool toolCallArgs: ToolCallArgs + cacheControl?: string } /** @@ -34,6 +35,7 @@ export async function resolveHttpEdgeRequest( const logger = ctx.get('logger') const ip = ctx.get('ip') + const cacheControl = ctx.req.header('cache-control') const { method } = ctx.req const requestUrl = new URL(ctx.req.url) const { pathname } = requestUrl @@ -115,6 +117,7 @@ export async function resolveHttpEdgeRequest( consumer, pricingPlan, tool, - toolCallArgs + toolCallArgs, + cacheControl } } diff --git a/apps/gateway/src/lib/resolve-origin-tool-call.ts b/apps/gateway/src/lib/resolve-origin-tool-call.ts index d50edcc7..ed49bbb5 100644 --- a/apps/gateway/src/lib/resolve-origin-tool-call.ts +++ b/apps/gateway/src/lib/resolve-origin-tool-call.ts @@ -19,6 +19,7 @@ import { createHttpRequestForOpenAPIOperation } from './create-http-request-for- import { enforceRateLimit } from './enforce-rate-limit' import { fetchCache } from './fetch-cache' import { getRequestCacheKey } from './get-request-cache-key' +import { isCacheControlPubliclyCacheable } from './is-cache-control-publicly-cacheable' import { updateOriginRequest } from './update-origin-request' export type ResolvedOriginToolCallResult = { @@ -48,6 +49,7 @@ export async function resolveOriginToolCall({ sessionId, env, ip, + cacheControl, waitUntil }: { tool: Tool @@ -58,6 +60,7 @@ export async function resolveOriginToolCall({ sessionId: string env: RawEnv ip?: string + cacheControl?: string waitUntil: (promise: Promise) => void }): Promise { // TODO: rate-limiting @@ -106,6 +109,17 @@ export async function resolveOriginToolCall({ rateLimit = toolConfig.rateLimit as RateLimit } + if (!cacheControl) { + if (toolConfig.cacheControl !== undefined) { + cacheControl = toolConfig.cacheControl + } else if (toolConfig.pure) { + cacheControl = + 'public, max-age=31560000, s-maxage=31560000, stale-while-revalidate=3600' + } else { + cacheControl = 'no-store' + } + } + const pricingPlanToolConfig = pricingPlan ? toolConfig.pricingPlanConfig?.[pricingPlan.slug] : undefined @@ -162,7 +176,7 @@ export async function resolveOriginToolCall({ deployment }) - updateOriginRequest(originRequest, { consumer, deployment }) + updateOriginRequest(originRequest, { consumer, deployment, cacheControl }) const cacheKey = await getRequestCacheKey(originRequest) @@ -173,15 +187,17 @@ export async function resolveOriginToolCall({ waitUntil }) - // non-cached version - // const originResponse = await fetch(originRequest) - return { toolCallArgs, originRequest, originResponse } } else if (originAdapter.type === 'mcp') { + const { projectIdentifier } = parseDeploymentIdentifier( + deployment.identifier, + { errorStatusCode: 500 } + ) + const id = env.DO_MCP_CLIENT.idFromName(sessionId) const originMcpClient = env.DO_MCP_CLIENT.get(id) @@ -191,11 +207,6 @@ export async function resolveOriginToolCall({ version: originAdapter.serverInfo.version }) - const { projectIdentifier } = parseDeploymentIdentifier( - deployment.identifier, - { errorStatusCode: 500 } - ) - const originMcpRequestMetadata = { agenticProxySecret: deployment._secret, sessionId, @@ -216,8 +227,35 @@ export async function resolveOriginToolCall({ projectIdentifier } as AgenticMcpRequestMetadata + let cacheKey: Request | undefined + + if (cacheControl && isCacheControlPubliclyCacheable(cacheControl)) { + const fakeOriginRequest = new Request(deployment.originUrl, { + method: 'POST', + headers: { + 'content-type': 'application/json', + 'cache-control': cacheControl + }, + body: JSON.stringify({ + name: tool.name, + args: toolCallArgs, + metadata: originMcpRequestMetadata! + }) + }) + + cacheKey = await getRequestCacheKey(fakeOriginRequest) + if (cacheKey) { + const response = await caches.default.match(cacheKey) + if (response) { + return { + toolCallArgs, + toolCallResponse: (await response.json()) as McpToolCallResponse + } + } + } + } + // TODO: add timeout support to the origin tool call? - // TODO: add response caching for origin MCP tool calls const toolCallResponseString = await originMcpClient.callTool({ name: tool.name, args: toolCallArgs, @@ -227,6 +265,16 @@ export async function resolveOriginToolCall({ toolCallResponseString ) as McpToolCallResponse + if (cacheKey && cacheControl) { + const fakeHttpResponse = new Response(toolCallResponseString, { + headers: { + 'content-type': 'application/json', + 'cache-control': cacheControl + } + }) + waitUntil(caches.default.put(cacheKey, fakeHttpResponse)) + } + return { toolCallArgs, toolCallResponse diff --git a/apps/gateway/src/lib/update-origin-request.ts b/apps/gateway/src/lib/update-origin-request.ts index 9b44525c..953b7031 100644 --- a/apps/gateway/src/lib/update-origin-request.ts +++ b/apps/gateway/src/lib/update-origin-request.ts @@ -9,10 +9,12 @@ export function updateOriginRequest( originRequest: Request, { deployment, - consumer + consumer, + cacheControl }: { deployment: AdminDeployment consumer?: AdminConsumer + cacheControl?: string } ) { // originRequest.headers.delete('authorization') @@ -88,5 +90,9 @@ export function updateOriginRequest( // https://support.cloudflare.com/hc/en-us/articles/360029779472-Troubleshooting-Cloudflare-1XXX-errors#error1000 // originRequest.headers.set('x-forwarded-for', ip) + if (cacheControl) { + originRequest.headers.set('cache-control', cacheControl) + } + originRequest.headers.set('x-agentic-proxy-secret', deployment._secret) } diff --git a/packages/types/src/openapi.d.ts b/packages/types/src/openapi.d.ts index 16a0014e..16aa49b7 100644 --- a/packages/types/src/openapi.d.ts +++ b/packages/types/src/openapi.d.ts @@ -495,7 +495,8 @@ export interface components { /** @default true */ enabled: boolean; /** @default false */ - immutable: boolean; + pure: boolean; + cacheControl?: string; /** @default true */ reportUsage: boolean; rateLimit?: components["schemas"]["RateLimit"] | null; diff --git a/packages/types/src/tools.ts b/packages/types/src/tools.ts index eb7cfc10..7054767d 100644 --- a/packages/types/src/tools.ts +++ b/packages/types/src/tools.ts @@ -114,7 +114,20 @@ export const toolConfigSchema = z * * @default false */ - immutable: z.boolean().optional().default(false), + pure: z.boolean().optional().default(false), + + /** + * A `Cache-Control` header value to use for caching this tool's responses. + * + * If `pure` is `true`, this defaults to: `public, max-age=31560000, s-maxage=31560000, stale-while-revalidate=3600` (cache publicly for up to 1 year). + * + * If `pure` is `false`, this defaults to the origin server's + * `cache-control` header value. If the origin server does not set a + * `cache-control` header, it defaults to `no-store`. + * + * @default undefined + */ + cacheControl: z.string().optional(), /** * Whether calls to this tool should be reported as usage for the default diff --git a/readme.md b/readme.md index 0e484412..7fde9e49 100644 --- a/readme.md +++ b/readme.md @@ -15,8 +15,7 @@ - => OpenAPI: `GET/POST/ETC originUrl/toolName` operation with transformed tool params - RAW: `METHOD gateway.agentic.so/deploymentIdentifier/` - => Raw HTTP: `METHOD originUrl/` simple HTTP proxy request - -**do I just ditch public REST interface and focus on MCP?** + - TODO: remove / disable `raw` support for now ## TODO @@ -36,8 +35,6 @@ - how to handle binary bodies and responses? - add support for `immutable` in `toolConfigs` - **Public MCP server interface** - - _McpAgent.serve_ - - how do I use consumer auth tokens with this flow? - how does oauth work with this flow? - **Origin MCP servers** - how to guarantee that the request is coming from agentic?