From b774480a1510e2c2725c6d7aeddeeb197864fa15 Mon Sep 17 00:00:00 2001
From: Travis Fischer <fisch0920@gmail.com>
Date: Sun, 1 Jun 2025 01:33:39 +0700
Subject: [PATCH] =?UTF-8?q?=F0=9F=9A=85?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/gateway/src/lib/fetch-cache.ts           | 17 +++++++++++++----
 .../gateway/src/lib/resolve-origin-request.ts |  4 +++-
 apps/gateway/src/worker.ts                    | 19 +++++++++++++++----
 3 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/apps/gateway/src/lib/fetch-cache.ts b/apps/gateway/src/lib/fetch-cache.ts
index 2dbc74fb..6fcfd400 100644
--- a/apps/gateway/src/lib/fetch-cache.ts
+++ b/apps/gateway/src/lib/fetch-cache.ts
@@ -1,8 +1,17 @@
+import type { Context } from './types'
+
 const cache = caches.default
 
-export async function fetchCache(opts) {
-  const { event, cacheKey, fetch: fetchResponse } = opts
-
+export async function fetchCache(
+  ctx: Context,
+  {
+    cacheKey,
+    fetchResponse
+  }: {
+    cacheKey?: string
+    fetchResponse: () => Promise<Response>
+  }
+) {
   let response
 
   if (cacheKey) {
@@ -16,7 +25,7 @@ export async function fetchCache(opts) {
     if (cacheKey) {
       if (response.headers.has('Cache-Control')) {
         // cache will respect response headers
-        event.waitUntil(
+        ctx.waitUntil(
           cache.put(cacheKey, response.clone()).catch((err) => {
             console.warn('cache put error', cacheKey, err)
           })
diff --git a/apps/gateway/src/lib/resolve-origin-request.ts b/apps/gateway/src/lib/resolve-origin-request.ts
index 44bb3bad..b0c94e18 100644
--- a/apps/gateway/src/lib/resolve-origin-request.ts
+++ b/apps/gateway/src/lib/resolve-origin-request.ts
@@ -26,6 +26,8 @@ export async function resolveOriginRequest(
   const { search, pathname } = requestUrl
   const method = req.method.toLowerCase()
   const requestPathParts = pathname.split('/')
+
+  // TODO: the isMCPRequest logic needs to be completely redone.
   const isMCPRequest = requestPathParts[0] === 'mcp'
   const requestPath = isMCPRequest
     ? requestPathParts.slice(1).join('/')
@@ -70,7 +72,7 @@ export async function resolveOriginRequest(
       `Auth token "${token}" is not authorized for project "${deployment.projectId}"`
     )
 
-    // TODO: Ensure that consumer.plan is compatible with the target deployment
+    // TODO: Ensure that consumer.plan is compatible with the target deployment?
     // TODO: This could definitely cause issues when changing pricing plans.
 
     pricingPlan = deployment.pricingPlans.find(
diff --git a/apps/gateway/src/worker.ts b/apps/gateway/src/worker.ts
index e638bffd..a878ae54 100644
--- a/apps/gateway/src/worker.ts
+++ b/apps/gateway/src/worker.ts
@@ -1,5 +1,5 @@
 import { AgenticApiClient } from '@agentic/platform-api-client'
-import { parseZodSchema } from '@agentic/platform-core'
+import { assert, parseZodSchema } from '@agentic/platform-core'
 
 import type { Context } from './lib/types'
 import { type AgenticEnv, envSchema } from './lib/env'
@@ -67,21 +67,32 @@ export default {
 
         switch (resolvedOriginRequest.deployment.originAdapter.type) {
           case 'openapi':
-            originResponse = await fetch(resolvedOriginRequest.originRequest!)
+            assert(
+              resolvedOriginRequest.originRequest,
+              500,
+              'Origin request is required'
+            )
+            originResponse = await fetch(resolvedOriginRequest.originRequest)
             break
 
           case 'raw':
-            originResponse = await fetch(resolvedOriginRequest.originRequest!)
+            assert(
+              resolvedOriginRequest.originRequest,
+              500,
+              'Origin request is required'
+            )
+            originResponse = await fetch(resolvedOriginRequest.originRequest)
             break
 
           case 'mcp':
             throw new Error('MCP not yet supported')
         }
 
+        assert(originResponse, 500, 'Origin response is required')
         const res = new Response(originResponse.body, originResponse)
-        recordTimespans()
 
         // Record the time it took for both the origin and gateway proxy to respond
+        recordTimespans()
         res.headers.set('x-response-time', `${originTimespan!}ms`)
         res.headers.set('x-proxy-response-time', `${gatewayTimespan!}ms`)