Integrate reply-from for workspace proxy

2025-11-19 02:27:07 +00:00
parent 146eae5220
commit 7aa94e7a88
3 changed files with 51 additions and 74 deletions
--- a/package-lock.json
+++ b/package-lock.json
@@ -727,6 +727,15 @@
      "integrity": "sha512-NM8/P9n3XjXhIZn1lLhkFaACTOURQXjWhV4BA/RnOv8xvgqtqpAX9IO4mRQxSx1Rlo4tqzeqb0sOlruaOy3dug==",
      "license": "MIT"
    },
+    "node_modules/@fastify/busboy": {
+      "version": "2.1.1",
+      "resolved": "https://registry.npmjs.org/@fastify/busboy/-/busboy-2.1.1.tgz",
+      "integrity": "sha512-vBZP4NlzfOlerQTnba4aqZoMhE/a9HY7HRqoOPaETQcSQuWEIyZMHGfVu6w9wGtGK5fED5qRs2DteVCjOH60sA==",
+      "license": "MIT",
+      "engines": {
+        "node": ">=14"
+      }
+    },
    "node_modules/@fastify/cors": {
      "version": "8.5.0",
      "resolved": "https://registry.npmjs.org/@fastify/cors/-/cors-8.5.0.tgz",
@@ -761,6 +770,33 @@
        "fast-deep-equal": "^3.1.3"
      }
    },
+    "node_modules/@fastify/reply-from": {
+      "version": "9.8.0",
+      "resolved": "https://registry.npmjs.org/@fastify/reply-from/-/reply-from-9.8.0.tgz",
+      "integrity": "sha512-bPNVaFhEeNI0Lyl6404YZaPFokudCplidE3QoOcr78yOy6H9sYw97p5KPYvY/NJNUHfFtvxOaSAHnK+YSiv/Mg==",
+      "license": "MIT",
+      "dependencies": {
+        "@fastify/error": "^3.0.0",
+        "end-of-stream": "^1.4.4",
+        "fast-content-type-parse": "^1.1.0",
+        "fast-querystring": "^1.0.0",
+        "fastify-plugin": "^4.0.0",
+        "toad-cache": "^3.7.0",
+        "undici": "^5.19.1"
+      }
+    },
+    "node_modules/@fastify/reply-from/node_modules/undici": {
+      "version": "5.29.0",
+      "resolved": "https://registry.npmjs.org/undici/-/undici-5.29.0.tgz",
+      "integrity": "sha512-raqeBD6NQK4SkWhQzeYKd1KmIG6dllBOTt55Rmkt4HtI9mwdWtJljnrXjAFUBLTSN67HWrOIZ3EPF4kjUw80Bg==",
+      "license": "MIT",
+      "dependencies": {
+        "@fastify/busboy": "^2.0.0"
+      },
+      "engines": {
+        "node": ">=14.0"
+      }
+    },
    "node_modules/@fastify/send": {
      "version": "2.1.0",
      "resolved": "https://registry.npmjs.org/@fastify/send/-/send-2.1.0.tgz",
@@ -3704,7 +3740,6 @@
      "version": "1.4.5",
      "resolved": "https://registry.npmjs.org/end-of-stream/-/end-of-stream-1.4.5.tgz",
      "integrity": "sha512-ooEGc6HP26xXq/N+GCGOT0JKCLDGrq2bQUZrQ7gyrJiZANJ/8YDTxTpQBXGMn+WbIQXNVpyWymm7KYVICQnyOg==",
-      "dev": true,
      "license": "MIT",
      "dependencies": {
        "once": "^1.4.0"
@@ -5975,7 +6010,6 @@
      "version": "1.4.0",
      "resolved": "https://registry.npmjs.org/once/-/once-1.4.0.tgz",
      "integrity": "sha512-lNaJgI+2Q5URQBkccEKHTQOPaXdUxnZZElQTZY0MFUAuaEqe1E+Nyvgdz/aIyNi6Z9MzO5dv1H8n58/GELp3+w==",
-      "dev": true,
      "license": "ISC",
      "dependencies": {
        "wrappy": "1"
@@ -8230,7 +8264,6 @@
      "version": "1.0.2",
      "resolved": "https://registry.npmjs.org/wrappy/-/wrappy-1.0.2.tgz",
      "integrity": "sha512-l4Sp/DRseor9wL6EvV2+TuQn63dMkPjZ/sp9XkghTEbV9KlPS1xUsZ3u7/IQO4wxtcFB4bgpQPRcR3QCvezPcQ==",
-      "dev": true,
      "license": "ISC"
    },
    "node_modules/xmlbuilder": {
@@ -8371,6 +8404,7 @@
      "version": "0.1.0",
      "dependencies": {
        "@fastify/cors": "^8.5.0",
+        "@fastify/reply-from": "^9.8.0",
        "@fastify/static": "^7.0.4",
        "commander": "^12.1.0",
        "fastify": "^4.28.1",
--- a/packages/cli/package.json
+++ b/packages/cli/package.json
@@ -16,6 +16,7 @@
  },
  "dependencies": {
    "@fastify/cors": "^8.5.0",
+    "@fastify/reply-from": "^9.8.0",
    "@fastify/static": "^7.0.4",
    "commander": "^12.1.0",
    "fastify": "^4.28.1",
--- a/packages/cli/src/server/http-server.ts
+++ b/packages/cli/src/server/http-server.ts
@@ -1,10 +1,9 @@
 import Fastify, { type FastifyInstance, type FastifyReply, type FastifyRequest } from "fastify"
 import cors from "@fastify/cors"
 import fastifyStatic from "@fastify/static"
+import replyFrom, { type FastifyReplyFromOptions } from "@fastify/reply-from"
 import fs from "fs"
 import path from "path"
-import { Readable } from "node:stream"
-import type { ReadableStream as NodeReadableStream } from "node:stream/web"
 import { fetch } from "undici"
 import type { Logger } from "../logger"
 import { WorkspaceManager } from "../workspaces/manager"
@@ -59,6 +58,10 @@ export function createHttpServer(deps: HttpServerDeps) {
    credentials: true,
  })

+  app.register(replyFrom, {
+    contentTypesToEncode: [],
+  })
+
  registerWorkspaceRoutes(app, { workspaceManager: deps.workspaceManager })
  registerConfigRoutes(app, { configStore: deps.configStore, binaryRegistry: deps.binaryRegistry })
  registerFilesystemRoutes(app, { fileSystemBrowser: deps.fileSystemBrowser })
@@ -91,7 +94,7 @@ interface InstanceProxyDeps {
 function registerInstanceProxyRoutes(app: FastifyInstance, deps: InstanceProxyDeps) {
  app.register(async (instance) => {
    instance.removeAllContentTypeParsers()
-    instance.addContentTypeParser("*", { parseAs: "buffer" }, (req, body, done) => done(null, body))
+    instance.addContentTypeParser("*", (req, body, done) => done(null, body))

    const proxyBaseHandler = async (request: FastifyRequest<{ Params: { id: string } }>, reply: FastifyReply) => {
      await proxyWorkspaceRequest({
@@ -122,7 +125,6 @@ function registerInstanceProxyRoutes(app: FastifyInstance, deps: InstanceProxyDe
 }

 const INSTANCE_PROXY_HOST = "127.0.0.1"
-const METHODS_WITHOUT_BODY = new Set(["GET", "HEAD", "OPTIONS"])

 async function proxyWorkspaceRequest(args: {
  request: FastifyRequest
@@ -151,74 +153,14 @@ async function proxyWorkspaceRequest(args: {
  const search = queryIndex >= 0 ? (request.raw.url ?? "").slice(queryIndex) : ""
  const targetUrl = `http://${INSTANCE_PROXY_HOST}:${port}${normalizedSuffix}${search}`

-  try {
-    const abortController = new AbortController()
-    const bodyPayload = METHODS_WITHOUT_BODY.has(request.method.toUpperCase())
-      ? undefined
-      : (request.body as Buffer | undefined)
-    const headers = buildProxyHeaders(request.headers)
-    if (bodyPayload && bodyPayload.byteLength > 0) {
-      headers["content-length"] = String(bodyPayload.byteLength)
-    } else {
-      delete headers["content-length"]
-    }
-    const response = await fetch(targetUrl, {
-      method: request.method,
-      headers,
-      body: bodyPayload,
-      signal: abortController.signal,
-    })
-
-    const headersToForward: Record<string, string> = {}
-    response.headers.forEach((value, key) => {
-      if (key.toLowerCase() === "content-length") {
-        return
+  return reply.from(targetUrl, {
+    onError: (proxyReply, { error }) => {
+      logger.error({ err: error, workspaceId, targetUrl }, "Failed to proxy workspace request")
+      if (!proxyReply.sent) {
+        proxyReply.code(502).send({ error: "Workspace instance proxy failed" })
      }
-      headersToForward[key] = value
-    })
-
-    const contentType = (response.headers.get("content-type") ?? "").toLowerCase()
-    const isEventStream = contentType.includes("text/event-stream")
-
-    if (isEventStream && response.body) {
-      reply.hijack()
-      Object.entries(headersToForward).forEach(([key, value]) => reply.raw.setHeader(key, value))
-      reply.raw.setHeader("Cache-Control", "no-cache")
-      reply.raw.setHeader("Connection", "keep-alive")
-      reply.raw.setHeader("Content-Type", "text/event-stream")
-      reply.raw.writeHead(response.status)
-      const stream = Readable.fromWeb(response.body as NodeReadableStream)
-      const cleanup = () => {
-        stream.destroy()
-        abortController.abort()
-      }
-      request.raw.on("close", cleanup)
-      request.raw.on("error", cleanup)
-      stream.on("error", cleanup)
-      stream.pipe(reply.raw)
-      return
-    }
-
-    Object.entries(headersToForward).forEach(([key, value]) => reply.header(key, value))
-
-    reply.code(response.status)
-
-    if (request.method === "HEAD") {
-      reply.send()
-      abortController.abort()
-      return
-    }
-
-    const bodyBuffer = Buffer.from(await response.arrayBuffer())
-    reply.header("content-length", String(bodyBuffer.byteLength))
-    reply.send(bodyBuffer)
-    abortController.abort()
-  } catch (error) {
-    logger.error({ err: error, workspaceId, targetUrl }, "Failed to proxy workspace request")
-    if (!reply.sent) {
-      reply.code(502).send({ error: "Workspace instance proxy failed" })
-    }
-  }
+    },
+  })
 }

 function normalizeInstanceSuffix(pathSuffix: string | undefined) {