History

Handle streaming OpenAPI responses

Rhys Sullivan committed 4 months ago

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

Showing the changes to packages/plugins/openapi (2 of 2 files). Show all

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 @@
1import { Effect, Layer, Option } from "effect";1import { Effect, Layer, Option, Stream } from "effect";
2import { HttpClient, HttpClientRequest } from "effect/unstable/http";2import { HttpClient, HttpClientRequest } from "effect/unstable/http";
3import type { ToolFileValue } from "@executor-js/sdk/core";3import type { ToolFileValue } from "@executor-js/sdk/core";
4 4
@@ -151,6 +151,28 @@ const applyHeaders = ( · 146 unmodified lines
151const normalizeContentType = (ct: string | null | undefined): string =>151const normalizeContentType = (ct: string | null | undefined): string =>
152 ct?.split(";")[0]?.trim().toLowerCase() ?? "";152 ct?.split(";")[0]?.trim().toLowerCase() ?? "";
153 153
154export const STREAM_MAX_BYTES = 1_000_000;
155export const STREAM_MAX_MS = 10_000;
156
157class StreamCapReached extends Error {
158 readonly _tag = "StreamCapReached";
159}
160
161export interface StreamingResponseCaps {
162 readonly maxBytes?: number;
163 readonly maxMs?: number;
164}
165
166const STREAMING_RESPONSE_CONTENT_TYPES = new Set([
167 "application/stream+json",
168 "application/x-ndjson",
169 "application/jsonl",
170 "text/event-stream",
171]);
172
173const isStreamingResponseContentType = (ct: string | null | undefined): boolean =>
174 STREAMING_RESPONSE_CONTENT_TYPES.has(normalizeContentType(ct));
175
154const isJsonContentType = (ct: string | null | undefined): boolean => {176const isJsonContentType = (ct: string | null | undefined): boolean => {
155 const normalized = normalizeContentType(ct);177 const normalized = normalizeContentType(ct);
156 if (!normalized) return false;178 if (!normalized) return false;
@@ -179,6 +201,115 @@ const isTextContentType = (ct: string | null | undefined): boolean => · 22 unmodified lines
179const isOctetStream = (ct: string | null | undefined): boolean =>201const isOctetStream = (ct: string | null | undefined): boolean =>
180 normalizeContentType(ct) === "application/octet-stream";202 normalizeContentType(ct) === "application/octet-stream";
181 203
204const 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
230const 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
246export 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
182const bytesToBase64 = (bytes: Uint8Array): string => {313const bytesToBase64 = (bytes: Uint8Array): string => {
183 let binary = "";314 let binary = "";
184 const chunkSize = 0x8000;315 const chunkSize = 0x8000;
@@ -861,6 +992,16 @@ export const invoke = Effect.fn("OpenApi.invoke")(function* ( · 676 unmodified lines
861 ? Option.getOrUndefined(responseBodyBinding.fileHint)992 ? Option.getOrUndefined(responseBodyBinding.fileHint)
862 : undefined;993 : undefined;
863 const ok = status >= 200 && status < 300;994 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 }
864 const responseBody: unknown =1005 const responseBody: unknown =
865 status === 2041006 status === 204
866 ? null1007 ? null
@@ -870,12 +1011,14 @@ export const invoke = Effect.fn("OpenApi.invoke")(function* ( · 3 unmodified lines
870 fileHint,1011 fileHint,
871 contentType,1012 contentType,
872 )1013 )
873 : isJsonContentType(contentType)1014 : streamingBody
874 ? yield* response.json.pipe(1015 ? streamingBody.data
875 Effect.catch(() => response.text),1016 : isJsonContentType(contentType)
876 mapBodyError,1017 ? yield* response.json.pipe(
877 )1018 Effect.catch(() => response.text),
878 : yield* response.text.pipe(mapBodyError);1019 mapBodyError,
1020 )
1021 : yield* response.text.pipe(mapBodyError);
879 1022
880 const dataBody =1023 const dataBody =
881 ok && fileHint?.kind === "byteField"1024 ok && fileHint?.kind === "byteField"
packages/plugins/openapi/src/sdk/streaming-response.test.ts+276−0
@@ -0,0 +1,276 @@
1import { describe, expect, it } from "@effect/vitest";
2import { Effect, Exit, Schema, Stream } from "effect";
3import { FetchHttpClient } from "effect/unstable/http";
4import { HttpApi, HttpApiEndpoint, HttpApiGroup, OpenApi } from "effect/unstable/httpapi";
5import { createServer, type ServerResponse } from "node:http";
6import type { AddressInfo, Socket } from "node:net";
7
8import {
9 createExecutor,
10 AuthTemplateSlug,
11 ConnectionName,
12 IntegrationSlug,
13 ToolAddress,
14} from "@executor-js/sdk";
15import { makeTestConfig, memoryCredentialsPlugin } from "@executor-js/sdk/testing";
16import { makeOpenApiHttpApiTestSourceConfig, unwrapInvocation } from "../testing";
17
18import { collectStreamingBody, STREAM_MAX_BYTES } from "./invoke";
19import { openApiPlugin } from "./plugin";
20
21const testPlugins = () =>
22 [openApiPlugin({ httpClientLayer: FetchHttpClient.layer }), memoryCredentialsPlugin()] as const;
23
24const STREAM_TOOL = "logs.getLogs";
25
26const LogsGroup = HttpApiGroup.make("logs").add(
27 HttpApiEndpoint.get("getLogs", "/logs", {
28 success: Schema.Unknown,
29 }),
30);
31
32const StreamingApi = HttpApi.make("streamingResponseTest")
33 .add(LogsGroup)
34 .annotateMerge(OpenApi.annotations({ title: "StreamingResponseTest", version: "1.0.0" }));
35
36const 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
55const 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
87const encoder = new TextEncoder();
88
89const 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
118describe("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});