refactor: unify stream next proxies

Author: jerry609

web/src/app/api/_utils/backend-proxy.test.ts

@@ -10,7 +10,7 @@ vi.mock("./auth-headers", () => ({
   withBackendAuth: withBackendAuthMock,
 }))
 
-import { apiBaseUrl, proxyJson, proxyText } from "./backend-proxy"
+import { apiBaseUrl, proxyJson, proxyStream, proxyText } from "./backend-proxy"
 
 describe("apiBaseUrl", () => {
   it("uses the shared backend base URL helper", () => {
@@ -172,4 +172,84 @@ describe("backend proxy helpers", () => {
     )
     expect(res.status).toBe(200)
   })
+
+  it("passes through SSE responses without buffering", async () => {
+    withBackendAuthMock.mockResolvedValue({
+      "Content-Type": "application/json",
+      authorization: "Bearer stream-token",
+    })
+
+    const fetchMock = vi.fn(async () =>
+      new Response("data: ping\n\n", {
+        status: 200,
+        headers: { "Content-Type": "text/event-stream" },
+      }),
+    )
+    global.fetch = fetchMock as typeof fetch
+
+    const req = new Request("https://localhost/api/studio/chat", {
+      method: "POST",
+      headers: { "Content-Type": "application/json" },
+      body: JSON.stringify({ prompt: "hello" }),
+    })
+    const res = await proxyStream(req, "https://backend/api/studio/chat", "POST", {
+      auth: true,
+      responseContentType: "text/event-stream",
+    })
+
+    expect(withBackendAuthMock).toHaveBeenCalledTimes(1)
+    const calls = fetchMock.mock.calls as unknown[][]
+    const init = calls[0]?.[1] as
+      | (RequestInit & { dispatcher?: unknown })
+      | undefined
+    expect(init).toBeDefined()
+    expect(init).toMatchObject({
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+        authorization: "Bearer stream-token",
+      },
+      body: JSON.stringify({ prompt: "hello" }),
+    })
+    expect(init?.signal).toBeUndefined()
+    expect(init?.dispatcher).toBeDefined()
+    expect(res.headers.get("content-type")).toContain("text/event-stream")
+    expect(await res.text()).toContain("data: ping")
+  })
+
+  it("returns JSON when a stream route falls back to a non-stream response", async () => {
+    withBackendAuthMock.mockResolvedValue({
+      Accept: "text/event-stream, application/json",
+      "Content-Type": "application/json",
+      authorization: "Bearer stream-token",
+    })
+
+    const fetchMock = vi.fn(async () =>
+      new Response(JSON.stringify({ done: true }), {
+        status: 200,
+        headers: { "Content-Type": "application/json" },
+      }),
+    )
+    global.fetch = fetchMock as typeof fetch
+
+    const req = new Request("https://localhost/api/research/paperscool/daily", {
+      method: "POST",
+      headers: { "Content-Type": "application/json" },
+      body: JSON.stringify({ query: "llm" }),
+    })
+    const res = await proxyStream(
+      req,
+      "https://backend/api/research/paperscool/daily",
+      "POST",
+      {
+        accept: "text/event-stream, application/json",
+        auth: true,
+        passthroughNonStreamResponse: true,
+        responseContentType: "text/event-stream",
+      },
+    )
+
+    expect(res.headers.get("content-type")).toContain("application/json")
+    await expect(res.json()).resolves.toEqual({ done: true })
+  })
 })

web/src/app/api/_utils/backend-proxy.ts

@@ -1,3 +1,5 @@
+import { Agent } from "undici"
+
 import { backendBaseUrl, withBackendAuth } from "./auth-headers"
 
 export type ProxyMethod = "DELETE" | "GET" | "HEAD" | "PATCH" | "POST" | "PUT"
@@ -22,7 +24,17 @@ type TextProxyOptions = ProxyOptions & {
   responseHeaders?: HeadersInit
 }
 
+type StreamProxyOptions = ProxyOptions & {
+  dispatcher?: Agent
+  passthroughNonStreamResponse?: boolean
+  responseContentType?: string
+}
+
 const DEFAULT_TIMEOUT_MS = 120_000
+const SSE_DISPATCHER = new Agent({
+  bodyTimeout: 0,
+  headersTimeout: 0,
+})
 
 export function apiBaseUrl(): string {
   return backendBaseUrl()
@@ -65,11 +77,58 @@ export async function proxyText(
   }
 }
 
+export async function proxyStream(
+  req: Request,
+  upstreamUrl: string,
+  method: ProxyMethod,
+  options: StreamProxyOptions = {},
+): Promise<Response> {
+  const requestOptions = {
+    responseContentType: "text/event-stream",
+    timeoutMs: 0,
+    ...options,
+  }
+
+  try {
+    const upstream = await fetchUpstream(req, upstreamUrl, method, requestOptions, {
+      dispatcher: requestOptions.dispatcher ?? SSE_DISPATCHER,
+    })
+    const upstreamContentType = upstream.headers.get("content-type") || ""
+
+    if (
+      requestOptions.passthroughNonStreamResponse &&
+      !upstreamContentType.includes("text/event-stream")
+    ) {
+      const text = await upstream.text()
+      return buildTextResponse(text, upstream, {
+        responseContentType: requestOptions.responseContentType ?? "application/json",
+        responseHeaders: undefined,
+      })
+    }
+
+    const headers = new Headers()
+    headers.set(
+      "Content-Type",
+      upstreamContentType || requestOptions.responseContentType || "text/event-stream",
+    )
+    headers.set("Cache-Control", "no-cache")
+    headers.set("Connection", "keep-alive")
+
+    return new Response(upstream.body, {
+      status: upstream.status,
+      headers,
+    })
+  } catch (error) {
+    return handleProxyError(error, upstreamUrl, requestOptions.onError)
+  }
+}
+
 async function fetchUpstream(
   req: Request,
   upstreamUrl: string,
   method: ProxyMethod,
   options: ProxyOptions,
+  init: RequestInit & { dispatcher?: Agent } = {},
 ): Promise<Response> {
   const controller = options.timeoutMs === 0 ? null : new AbortController()
   const timeoutMs = options.timeoutMs ?? DEFAULT_TIMEOUT_MS
@@ -78,12 +137,13 @@ async function fetchUpstream(
 
   try {
     return await fetch(upstreamUrl, {
+      ...init,
       method,
       headers: await resolveHeaders(req, body, options),
       body,
       cache: options.cache,
       signal: controller?.signal,
-    })
+    } as RequestInit & { dispatcher?: Agent })
   } finally {
     if (timeout) {
       clearTimeout(timeout)

web/src/app/api/gen-code/route.test.ts

@@ -0,0 +1,32 @@
+import { describe, expect, it, vi } from "vitest"
+
+const { apiBaseUrlMock, proxyStreamMock } = vi.hoisted(() => ({
+  apiBaseUrlMock: vi.fn(() => "http://backend.example.com"),
+  proxyStreamMock: vi.fn(),
+}))
+
+vi.mock("@/app/api/_utils/backend-proxy", () => ({
+  apiBaseUrl: apiBaseUrlMock,
+  proxyStream: proxyStreamMock,
+}))
+
+import { POST } from "./route"
+
+describe("gen code route", () => {
+  it("proxies code generation through the shared stream helper", async () => {
+    const req = new Request("http://localhost/api/gen-code", { method: "POST" })
+    proxyStreamMock.mockResolvedValueOnce(new Response("data: ping\n\n"))
+
+    await POST(req)
+
+    expect(proxyStreamMock).toHaveBeenCalledWith(
+      req,
+      "http://backend.example.com/api/gen-code",
+      "POST",
+      {
+        auth: true,
+        responseContentType: "text/event-stream",
+      },
+    )
+  })
+})

web/src/app/api/gen-code/route.ts

@@ -1,28 +1,10 @@
 export const runtime = "nodejs"
 
-function apiBaseUrl() {
-  return process.env.PAPERBOT_API_BASE_URL || "http://127.0.0.1:8000"
-}
-
-import { withBackendAuth } from "../_utils/auth-headers"
+import { apiBaseUrl, proxyStream } from "@/app/api/_utils/backend-proxy"
 
 export async function POST(req: Request) {
-  const body = await req.text()
-  const upstream = await fetch(`${apiBaseUrl()}/api/gen-code`, {
-    method: "POST",
-    headers: await withBackendAuth(req, {
-      "Content-Type": req.headers.get("content-type") || "application/json",
-    }),
-    body,
-  })
-
-  const headers = new Headers()
-  headers.set("Content-Type", upstream.headers.get("content-type") || "text/event-stream")
-  headers.set("Cache-Control", "no-cache")
-  headers.set("Connection", "keep-alive")
-
-  return new Response(upstream.body, {
-    status: upstream.status,
-    headers,
+  return proxyStream(req, `${apiBaseUrl()}/api/gen-code`, "POST", {
+    auth: true,
+    responseContentType: "text/event-stream",
   })
 }

web/src/app/api/research/paperscool/analyze/route.test.ts

@@ -0,0 +1,35 @@
+import { describe, expect, it, vi } from "vitest"
+
+const { apiBaseUrlMock, proxyStreamMock } = vi.hoisted(() => ({
+  apiBaseUrlMock: vi.fn(() => "http://backend.example.com"),
+  proxyStreamMock: vi.fn(),
+}))
+
+vi.mock("@/app/api/_utils/backend-proxy", () => ({
+  apiBaseUrl: apiBaseUrlMock,
+  proxyStream: proxyStreamMock,
+}))
+
+import { POST } from "./route"
+
+describe("paperscool analyze route", () => {
+  it("proxies analyze requests through the shared stream helper", async () => {
+    const req = new Request("http://localhost/api/research/paperscool/analyze", {
+      method: "POST",
+    })
+    proxyStreamMock.mockResolvedValueOnce(new Response("data: ping\n\n"))
+
+    await POST(req)
+
+    expect(proxyStreamMock).toHaveBeenCalledWith(
+      req,
+      "http://backend.example.com/api/research/paperscool/analyze",
+      "POST",
+      expect.objectContaining({
+        accept: "text/event-stream",
+        responseContentType: "text/event-stream",
+        onError: expect.any(Function),
+      }),
+    )
+  })
+})

web/src/app/api/research/paperscool/analyze/route.ts

@@ -1,43 +1,18 @@
 export const runtime = "nodejs"
 
-import { Agent } from "undici"
-
-import { apiBaseUrl } from "../../_base"
-
-// Analyze can stream for a long time; disable body timeout on the proxy hop.
-const sseDispatcher = new Agent({
-  bodyTimeout: 0,
-  headersTimeout: 0,
-})
+import { apiBaseUrl, proxyStream } from "@/app/api/_utils/backend-proxy"
 
 export async function POST(req: Request) {
-  const body = await req.text()
-  let upstream: Response
-  try {
-    upstream = await fetch(`${apiBaseUrl()}/api/research/paperscool/analyze`, {
-      method: "POST",
-      headers: {
-        "Content-Type": req.headers.get("content-type") || "application/json",
-        Accept: "text/event-stream",
-      },
-      body,
-      dispatcher: sseDispatcher,
-    } as RequestInit & { dispatcher: Agent })
-  } catch (error) {
-    const detail = error instanceof Error ? error.message : String(error)
-    return Response.json(
-      { detail: "Upstream API unreachable", error: detail },
-      { status: 502 },
-    )
-  }
-
-  const headers = new Headers()
-  headers.set("Content-Type", upstream.headers.get("content-type") || "text/event-stream")
-  headers.set("Cache-Control", "no-cache")
-  headers.set("Connection", "keep-alive")
-
-  return new Response(upstream.body, {
-    status: upstream.status,
-    headers,
+  return proxyStream(req, `${apiBaseUrl()}/api/research/paperscool/analyze`, "POST", {
+    accept: "text/event-stream",
+    responseContentType: "text/event-stream",
+    onError: ({ error }) =>
+      Response.json(
+        {
+          detail: "Upstream API unreachable",
+          error: error instanceof Error ? error.message : String(error),
+        },
+        { status: 502 },
+      ),
   })
 }