History

Handle streaming OpenAPI responses

Rhys Sullivan committed 4 months ago

Commit d0104a31 parent c46730b2 files · +426 −7Browse files

Unified · Split

2 files
  1. packages/plugins/openapi/src/sdk/invoke.ts +150 −7
  2. packages/plugins/openapi/src/sdk/streaming-response.test.ts +276 −0
packages/plugins/openapi/src/sdk/invoke.ts+150−7
@@ -1,4 +1,4 @@
1−import { Effect, Layer, Option } from "effect";
1+import { Effect, Layer, Option, Stream } from "effect";
22 import { HttpClient, HttpClientRequest } from "effect/unstable/http";
33 import type { ToolFileValue } from "@executor-js/sdk/core";
44
@@ -151,6 +151,28 @@ const applyHeaders = ( · 146 unmodified lines
151151 const normalizeContentType = (ct: string | null | undefined): string =>
152152 ct?.split(";")[0]?.trim().toLowerCase() ?? "";
153153
154+export const STREAM_MAX_BYTES = 1_000_000;
155+export const STREAM_MAX_MS = 10_000;
156+
157+class StreamCapReached extends Error {
158+ readonly _tag = "StreamCapReached";
159+}
160+
161+export interface StreamingResponseCaps {
162+ readonly maxBytes?: number;
163+ readonly maxMs?: number;
164+}
165+
166+const STREAMING_RESPONSE_CONTENT_TYPES = new Set([
167+ "application/stream+json",
168+ "application/x-ndjson",
169+ "application/jsonl",
170+ "text/event-stream",
171+]);
172+
173+const isStreamingResponseContentType = (ct: string | null | undefined): boolean =>
174+ STREAMING_RESPONSE_CONTENT_TYPES.has(normalizeContentType(ct));
175+
154176 const isJsonContentType = (ct: string | null | undefined): boolean => {
155177 const normalized = normalizeContentType(ct);
156178 if (!normalized) return false;
@@ -179,6 +201,115 @@ const isTextContentType = (ct: string | null | undefined): boolean => · 22 unmodified lines
179201 const isOctetStream = (ct: string | null | undefined): boolean =>
180202 normalizeContentType(ct) === "application/octet-stream";
181203
204+const parseStreamingJsonLines = (text: string, truncated: boolean): unknown => {
205+ const lines = text.split(/\r?\n/);
206+ let lastNonEmptyIndex = -1;
207+ for (let index = lines.length - 1; index >= 0; index--) {
208+ const line = lines[index];
209+ if (line !== undefined && line.trim() !== "") {
210+ lastNonEmptyIndex = index;
211+ break;
212+ }
213+ }
214+ const rows: unknown[] = [];
215+
216+ for (let index = 0; index < lines.length; index++) {
217+ const line = lines[index];
218+ if (line === undefined || line.trim() === "") continue;
219+ try {
220+ rows.push(JSON.parse(line));
221+ } catch {
222+ if (truncated && index === lastNonEmptyIndex) continue;
223+ return text;
224+ }
225+ }
226+
227+ return rows;
228+};
229+
230+const decodeStreamingResponseBody = (
231+ contentType: string | null,
232+ text: string,
233+ truncated: boolean,
234+): unknown => {
235+ switch (normalizeContentType(contentType)) {
236+ case "application/stream+json":
237+ case "application/x-ndjson":
238+ case "application/jsonl":
239+ return parseStreamingJsonLines(text, truncated);
240+ case "text/event-stream":
241+ default:
242+ return text;
243+ }
244+};
245+
246+export const collectStreamingBody = (
247+ stream: Stream.Stream<Uint8Array, unknown, never>,
248+ contentType: string | null,
249+ caps: StreamingResponseCaps = {},
250+) =>
251+ Effect.gen(function* () {
252+ const maxBytes = caps.maxBytes ?? STREAM_MAX_BYTES;
253+ const maxMs = caps.maxMs ?? STREAM_MAX_MS;
254+ const chunks: Uint8Array[] = [];
255+ let bytes = 0;
256+ let truncated = false;
257+ const startedAt = Date.now();
258+
259+ const completed = yield* stream.pipe(
260+ Stream.timeout(maxMs),
261+ Stream.runForEach((chunk) =>
262+ Effect.gen(function* () {
263+ const remaining = maxBytes - bytes;
264+ if (remaining <= 0) {
265+ truncated = true;
266+ return yield* Effect.fail(new StreamCapReached());
267+ }
268+
269+ if (chunk.byteLength > remaining) {
270+ chunks.push(chunk.subarray(0, remaining));
271+ bytes += remaining;
272+ truncated = true;
273+ return yield* Effect.fail(new StreamCapReached());
274+ }
275+
276+ chunks.push(chunk);
277+ bytes += chunk.byteLength;
278+ if (bytes >= maxBytes) {
279+ truncated = true;
280+ return yield* Effect.fail(new StreamCapReached());
281+ }
282+ }),
283+ ),
284+ Effect.timeoutOption(maxMs + 1_000),
285+ Effect.catch((error: unknown) =>
286+ error instanceof StreamCapReached
287+ ? Effect.succeed(Option.some(undefined))
288+ : Effect.fail(error),
289+ ),
290+ );
291+ if (Option.isNone(completed)) truncated = true;
292+
293+ const durationMs = Date.now() - startedAt;
294+ if (durationMs >= maxMs) truncated = true;
295+ const collected = new Uint8Array(bytes);
296+ let offset = 0;
297+ for (const chunk of chunks) {
298+ collected.set(chunk, offset);
299+ offset += chunk.byteLength;
300+ }
301+ const text = new TextDecoder("utf-8").decode(collected);
302+
303+ return {
304+ data: decodeStreamingResponseBody(contentType, text, truncated),
305+ headers: {
306+ "x-executor-stream": truncated ? "truncated" : "complete",
307+ "x-executor-stream-bytes": String(bytes),
308+ "x-executor-stream-duration-ms": String(durationMs),
309+ },
310+ };
311+ });
312+
182313 const bytesToBase64 = (bytes: Uint8Array): string => {
183314 let binary = "";
184315 const chunkSize = 0x8000;
@@ -861,6 +992,16 @@ export const invoke = Effect.fn("OpenApi.invoke")(function* ( · 676 unmodified lines
861992 ? Option.getOrUndefined(responseBodyBinding.fileHint)
862993 : undefined;
863994 const ok = status >= 200 && status < 300;
995+ const streamingBody =
996+ ok &&
997+ status !== 204 &&
998+ fileHint?.kind !== "binaryResponse" &&
999+ isStreamingResponseContentType(contentType)
1000+ ? yield* collectStreamingBody(response.stream, contentType).pipe(mapBodyError)
1001+ : null;
1002+ if (streamingBody) {
1003+ Object.assign(responseHeaders, streamingBody.headers);
1004+ }
8641005 const responseBody: unknown =
8651006 status === 204
8661007 ? null
@@ -870,12 +1011,14 @@ export const invoke = Effect.fn("OpenApi.invoke")(function* ( · 3 unmodified lines
8701011 fileHint,
8711012 contentType,
8721013 )
873− : isJsonContentType(contentType)
874− ? yield* response.json.pipe(
875− Effect.catch(() => response.text),
876− mapBodyError,
877− )
878− : yield* response.text.pipe(mapBodyError);
1014+ : streamingBody
1015+ ? streamingBody.data
1016+ : isJsonContentType(contentType)
1017+ ? yield* response.json.pipe(
1018+ Effect.catch(() => response.text),
1019+ mapBodyError,
1020+ )
1021+ : yield* response.text.pipe(mapBodyError);
8791022
8801023 const dataBody =
8811024 ok && fileHint?.kind === "byteField"
packages/plugins/openapi/src/sdk/streaming-response.test.ts+276−0
@@ -0,0 +1,276 @@
1+import { describe, expect, it } from "@effect/vitest";
2+import { Effect, Exit, Schema, Stream } from "effect";
3+import { FetchHttpClient } from "effect/unstable/http";
4+import { HttpApi, HttpApiEndpoint, HttpApiGroup, OpenApi } from "effect/unstable/httpapi";
5+import { createServer, type ServerResponse } from "node:http";
6+import type { AddressInfo, Socket } from "node:net";
7+
8+import {
9+ createExecutor,
10+ AuthTemplateSlug,
11+ ConnectionName,
12+ IntegrationSlug,
13+ ToolAddress,
14+} from "@executor-js/sdk";
15+import { makeTestConfig, memoryCredentialsPlugin } from "@executor-js/sdk/testing";
16+import { makeOpenApiHttpApiTestSourceConfig, unwrapInvocation } from "../testing";
17+
18+import { collectStreamingBody, STREAM_MAX_BYTES } from "./invoke";
19+import { openApiPlugin } from "./plugin";
20+
21+const testPlugins = () =>
22+ [openApiPlugin({ httpClientLayer: FetchHttpClient.layer }), memoryCredentialsPlugin()] as const;
23+
24+const STREAM_TOOL = "logs.getLogs";
25+
26+const LogsGroup = HttpApiGroup.make("logs").add(
27+ HttpApiEndpoint.get("getLogs", "/logs", {
28+ success: Schema.Unknown,
29+ }),
30+);
31+
32+const StreamingApi = HttpApi.make("streamingResponseTest")
33+ .add(LogsGroup)
34+ .annotateMerge(OpenApi.annotations({ title: "StreamingResponseTest", version: "1.0.0" }));
35+
36+const buildExecutor = (baseUrl: string) =>
37+ Effect.gen(function* () {
38+ const executor = yield* createExecutor(makeTestConfig({ plugins: testPlugins() }));
39+ yield* executor.openapi.addSpec(
40+ makeOpenApiHttpApiTestSourceConfig(StreamingApi, { slug: "streaming", baseUrl }),
41+ );
42+ yield* executor.connections.create({
43+ owner: "org",
44+ name: ConnectionName.make("main"),
45+ integration: IntegrationSlug.make("streaming"),
46+ template: AuthTemplateSlug.make("apiKey"),
47+ value: "token",
48+ });
49+ return {
50+ executor,
51+ address: ToolAddress.make(`tools.streaming.org.main.${STREAM_TOOL}`),
52+ };
53+ });
54+
55+const startRawServer = (handler: (res: ServerResponse) => void) =>
56+ Effect.acquireRelease(
57+ Effect.callback<{ baseUrl: string; close: () => void }>((resume) => {
58+ const sockets = new Set<Socket>();
59+ const server = createServer((req, res) => {
60+ if (req.url?.split("?")[0] !== "/logs") {
61+ res.writeHead(200, { "content-type": "application/json" });
62+ res.end(JSON.stringify({ ok: true }));
63+ return;
64+ }
65+ handler(res);
66+ });
67+ server.on("connection", (socket) => {
68+ sockets.add(socket);
69+ socket.on("close", () => sockets.delete(socket));
70+ });
71+ server.listen(0, "127.0.0.1", () => {
72+ const port = (server.address() as AddressInfo).port;
73+ resume(
74+ Effect.succeed({
75+ baseUrl: `http://127.0.0.1:${port}`,
76+ close: () => {
77+ for (const socket of sockets) socket.destroy();
78+ server.close();
79+ },
80+ }),
81+ );
82+ });
83+ }),
84+ (server) => Effect.sync(() => server.close()),
85+ );
86+
87+const encoder = new TextEncoder();
88+
89+const delayedReadableStream = (
90+ chunks: readonly Uint8Array[],
91+ delayMs: number,
92+ options: { close: boolean },
93+) =>
94+ Stream.fromReadableStream({
95+ evaluate: () => {
96+ const timers: ReturnType<typeof setTimeout>[] = [];
97+ return new ReadableStream<Uint8Array>({
98+ start: (controller) => {
99+ chunks.forEach((chunk, index) => {
100+ timers.push(
101+ setTimeout(() => {
102+ controller.enqueue(chunk);
103+ if (options.close && index === chunks.length - 1) {
104+ controller.close();
105+ }
106+ }, index * delayMs),
107+ );
108+ });
109+ },
110+ cancel: () => {
111+ for (const timer of timers) clearTimeout(timer);
112+ },
113+ });
114+ },
115+ onError: (error) => error,
116+ });
117+
118+describe("OpenAPI streaming responses", () => {
119+ it.effect("returns from an oversized NDJSON stream with truncation metadata", () =>
120+ Effect.gen(function* () {
121+ const rows = Array.from(
122+ { length: 6 },
123+ (_, index) => `${JSON.stringify({ index, phase: "immediate" })}\n`,
124+ ).join("");
125+ const rowBytes = encoder.encode(rows);
126+ const partialBytes = encoder.encode(`{"partial":"${"x".repeat(STREAM_MAX_BYTES)}"`);
127+
128+ const result = yield* collectStreamingBody(
129+ Stream.make(rowBytes, partialBytes),
130+ "application/stream+json",
131+ );
132+
133+ expect(result.headers["x-executor-stream"]).toBe("truncated");
134+ expect(result.headers["x-executor-stream-bytes"]).toBe(String(STREAM_MAX_BYTES));
135+ expect(Array.isArray(result.data)).toBe(true);
136+ const data = result.data as unknown[];
137+ expect(data.slice(0, 6)).toEqual(
138+ Array.from({ length: 6 }, (_, index) => ({ index, phase: "immediate" })),
139+ );
140+ }),
141+ );
142+
143+ it.effect("collects every chunk from a delayed ReadableStream", () =>
144+ Effect.gen(function* () {
145+ const chunks = [
146+ encoder.encode(`${JSON.stringify({ id: 1 })}\n`),
147+ encoder.encode(`${JSON.stringify({ id: 2 })}\n`),
148+ encoder.encode(`${JSON.stringify({ id: 3 })}\n`),
149+ ];
150+ const expectedBytes = chunks.reduce((sum, chunk) => sum + chunk.byteLength, 0);
151+
152+ const result = yield* collectStreamingBody(
153+ delayedReadableStream(chunks, 75, { close: true }),
154+ "application/x-ndjson",
155+ { maxBytes: STREAM_MAX_BYTES, maxMs: 1_000 },
156+ );
157+
158+ console.log(
159+ JSON.stringify({
160+ case: "delayed-readable-complete",
161+ bytes: result.headers["x-executor-stream-bytes"],
162+ rows: Array.isArray(result.data) ? result.data.length : 0,
163+ }),
164+ );
165+ expect(result.headers["x-executor-stream"]).toBe("complete");
166+ expect(result.headers["x-executor-stream-bytes"]).toBe(String(expectedBytes));
167+ expect(result.data).toEqual([{ id: 1 }, { id: 2 }, { id: 3 }]);
168+ }),
169+ );
170+
171+ it("timeout truncation preserves all chunks received before interruption", async () => {
172+ await Effect.runPromise(
173+ Effect.gen(function* () {
174+ const chunks = [
175+ encoder.encode(`${JSON.stringify({ id: 1 })}\n`),
176+ encoder.encode(`${JSON.stringify({ id: 2 })}\n`),
177+ ];
178+ const expectedBytes = chunks.reduce((sum, chunk) => sum + chunk.byteLength, 0);
179+ const startedAt = Date.now();
180+
181+ const result = yield* collectStreamingBody(
182+ delayedReadableStream(chunks, 100, { close: false }),
183+ "application/stream+json",
184+ { maxBytes: STREAM_MAX_BYTES, maxMs: 500 },
185+ );
186+ const elapsedMs = Date.now() - startedAt;
187+
188+ console.log(
189+ JSON.stringify({
190+ case: "delayed-readable-timeout",
191+ elapsedMs,
192+ bytes: result.headers["x-executor-stream-bytes"],
193+ rows: Array.isArray(result.data) ? result.data.length : 0,
194+ }),
195+ );
196+ expect(elapsedMs).toBeGreaterThanOrEqual(450);
197+ expect(elapsedMs).toBeLessThan(1_500);
198+ expect(result.headers["x-executor-stream"]).toBe("truncated");
199+ expect(result.headers["x-executor-stream-bytes"]).toBe(String(expectedBytes));
200+ expect(result.data).toEqual([{ id: 1 }, { id: 2 }]);
201+ }),
202+ );
203+ });
204+
205+ it.effect("drops a trailing partial JSON line when the byte cap truncates", () =>
206+ Effect.gen(function* () {
207+ let body = "";
208+ for (let index = 0; Buffer.byteLength(body) < STREAM_MAX_BYTES - 2_000; index++) {
209+ body += `${JSON.stringify({ index, value: "x".repeat(512) })}\n`;
210+ }
211+ const encoded = encoder.encode(`${body}{"partial":"${"y".repeat(STREAM_MAX_BYTES)}"`);
212+
213+ const result = yield* collectStreamingBody(Stream.make(encoded), "application/x-ndjson");
214+
215+ expect(result.headers["x-executor-stream"]).toBe("truncated");
216+ expect(Array.isArray(result.data)).toBe(true);
217+ const data = result.data as unknown[];
218+ expect(data.length).toBeGreaterThan(0);
219+ expect(data.every((row: unknown) => typeof row === "object" && row !== null)).toBe(true);
220+ expect(data.some((row: unknown) => "partial" in (row as Record<string, unknown>))).toBe(
221+ false,
222+ );
223+ }),
224+ );
225+
226+ it.effect("parses a completed short NDJSON stream", () =>
227+ Effect.gen(function* () {
228+ const server = yield* startRawServer((res) => {
229+ res.writeHead(200, { "content-type": "application/jsonl" });
230+ res.end(`${JSON.stringify({ id: 1 })}\n${JSON.stringify({ id: 2 })}\n`);
231+ });
232+ const { executor, address } = yield* buildExecutor(server.baseUrl);
233+
234+ const result = unwrapInvocation<unknown[]>(yield* executor.execute(address, {}));
235+
236+ expect(result.data).toEqual([{ id: 1 }, { id: 2 }]);
237+ expect(result.headers?.["x-executor-stream"]).toBe("complete");
238+ }),
239+ );
240+
241+ it.effect("leaves ordinary JSON responses on the buffered path", () =>
242+ Effect.gen(function* () {
243+ const server = yield* startRawServer((res) => {
244+ res.writeHead(200, { "content-type": "application/json" });
245+ res.end(JSON.stringify({ ok: true }));
246+ });
247+ const { executor, address } = yield* buildExecutor(server.baseUrl);
248+
249+ const result = unwrapInvocation<Record<string, unknown>>(
250+ yield* executor.execute(address, {}),
251+ );
252+
253+ expect(result.data).toEqual({ ok: true });
254+ expect(result.headers?.["x-executor-stream"]).toBeUndefined();
255+ }),
256+ );
257+
258+ it.effect("leaves non-2xx JSON error responses on the buffered path", () =>
259+ Effect.gen(function* () {
260+ const server = yield* startRawServer((res) => {
261+ res.writeHead(422, { "content-type": "application/json" });
262+ res.end('{"error":{"field":"name","reason":"too_short"}}');
263+ });
264+ const { executor, address } = yield* buildExecutor(server.baseUrl);
265+
266+ const exit = yield* executor.execute(address, {}).pipe(Effect.exit);
267+ const text = Exit.match(exit, {
268+ onFailure: (cause) => JSON.stringify(cause),
269+ onSuccess: (value) => JSON.stringify(value),
270+ });
271+
272+ expect(text).toMatch(/422|too_short|name|response|error/i);
273+ expect(text).not.toMatch(/x-executor-stream/);
274+ }),
275+ );
276+});