From e03d12a75e243ac75250f5f4416bbd63d1c0657b Mon Sep 17 00:00:00 2001
From: Martin PAUCOT <contact@martin-paucot.fr>
Date: Thu, 12 Dec 2024 16:16:23 +0100
Subject: [PATCH] feat(cache): cache event handler stream response

---
 src/runtime/internal/cache.ts        | 18 +++++++++++++++++-
 test/fixture/nitro.config.ts         |  1 +
 test/fixture/routes/stream-cached.ts | 13 +++++++++++++
 test/presets/netlify-legacy.test.ts  |  1 +
 test/presets/vercel.test.ts          |  4 ++++
 test/tests.ts                        | 28 ++++++++++++++++++++++++++++
 6 files changed, 64 insertions(+), 1 deletion(-)
 create mode 100644 test/fixture/routes/stream-cached.ts

diff --git a/src/runtime/internal/cache.ts b/src/runtime/internal/cache.ts
index 838fafe969..2739164196 100644
--- a/src/runtime/internal/cache.ts
+++ b/src/runtime/internal/cache.ts
@@ -5,6 +5,7 @@ import {
   fetchWithEvent,
   handleCacheHeaders,
   isEvent,
+  isStream,
   splitCookiesString,
 } from "h3";
 import type { EventHandlerRequest, EventHandlerResponse, H3Event } from "h3";
@@ -406,11 +407,26 @@ export function defineCachedEventHandler<
         headers["cache-control"] = cacheControl.join(", ");
       }
 
+      let cachedBody = body as any;
+      // When handler response is a stream, we cache the result of this stream
+      if (body instanceof ReadableStream) {
+        const td = new TextDecoder();
+        let buffer = "";
+        await body.pipeTo(
+          new WritableStream({
+            write(chunk) {
+              buffer += td.decode(chunk);
+            },
+          })
+        );
+        cachedBody = buffer;
+      }
+
       // Create cache entry for response
       const cacheEntry: ResponseCacheEntry<Response> = {
         code: event.node.res.statusCode,
         headers,
-        body,
+        body: cachedBody,
       };
 
       return cacheEntry;
diff --git a/test/fixture/nitro.config.ts b/test/fixture/nitro.config.ts
index 63d8eca4cd..c45a736006 100644
--- a/test/fixture/nitro.config.ts
+++ b/test/fixture/nitro.config.ts
@@ -91,6 +91,7 @@ export default defineNitroConfig({
     "/rules/_/cached/noncached": { cache: false, swr: false, isr: false },
     "/rules/_/cached/**": { swr: true },
     "/api/proxy/**": { proxy: "/api/echo" },
+    "/stream-cached": { swr: true },
   },
   prerender: {
     crawlLinks: true,
diff --git a/test/fixture/routes/stream-cached.ts b/test/fixture/routes/stream-cached.ts
new file mode 100644
index 0000000000..5d8ee68461
--- /dev/null
+++ b/test/fixture/routes/stream-cached.ts
@@ -0,0 +1,13 @@
+export default eventHandler(() => {
+  const encoder = new TextEncoder();
+  const stream = new ReadableStream({
+    start(controller) {
+      controller.enqueue(encoder.encode("nitro"));
+      controller.enqueue(encoder.encode("is"));
+      controller.enqueue(encoder.encode("awesome,"));
+      controller.enqueue(encoder.encode(Date.now().toString()));
+      controller.close();
+    },
+  });
+  return stream;
+});
diff --git a/test/presets/netlify-legacy.test.ts b/test/presets/netlify-legacy.test.ts
index 5de5771d7f..986af35c39 100644
--- a/test/presets/netlify-legacy.test.ts
+++ b/test/presets/netlify-legacy.test.ts
@@ -71,6 +71,7 @@ describe("nitro:preset:netlify-legacy", async () => {
         /rules/isr-ttl/*	/.netlify/builders/server 200
         /rules/isr/*	/.netlify/builders/server 200
         /rules/dynamic	/.netlify/functions/server 200
+        /stream-cached	/.netlify/builders/server 200
         /build/* /build/:splat 200
         /* /.netlify/functions/server 200"
       `);
diff --git a/test/presets/vercel.test.ts b/test/presets/vercel.test.ts
index 25ad76cfe2..b09e4751a1 100644
--- a/test/presets/vercel.test.ts
+++ b/test/presets/vercel.test.ts
@@ -142,6 +142,10 @@ describe("nitro:preset:vercel", async () => {
                 "dest": "/__nitro--rules-swr-ttl?url=$url",
                 "src": "(?<url>/rules/swr-ttl/.*)",
               },
+              {
+                "dest": "/stream-cached?url=$url",
+                "src": "/stream-cached",
+              },
               {
                 "dest": "/__nitro",
                 "src": "/(.*)",
diff --git a/test/tests.ts b/test/tests.ts
index 122e06c735..d142c4e6a8 100644
--- a/test/tests.ts
+++ b/test/tests.ts
@@ -729,6 +729,34 @@ export function testNitro(
         }
       }
     );
+
+    it.skipIf(
+      ctx.isIsolated ||
+        (isWindows && ctx.preset === "nitro-dev") ||
+        ctx.isLambda
+    )("should cache stream result", async () => {
+      const { data } = await callHandler({
+        url: "/stream-cached",
+      });
+
+      const [str, timestamp] = data.split(",") as string[];
+
+      expect(str).toBe(
+        ctx.isLambda ? btoa("nitroisawesome") : "nitroisawesome"
+      );
+
+      const calls = await Promise.all([
+        callHandler({ url: "/stream-cached" }),
+        callHandler({ url: "/stream-cached" }),
+        callHandler({ url: "/stream-cached" }),
+      ]);
+
+      for (const call of calls) {
+        expect(call.data).toBe(
+          ctx.isLambda ? btoa(`${str},${timestamp}`) : `${str},${timestamp}`
+        );
+      }
+    });
   });
 
   describe("scanned files", () => {