Skip to content

Commit 2acdbf0

Browse files
committed
fix(memory): abort batch upload response reads
1 parent d6cefe2 commit 2acdbf0

2 files changed

Lines changed: 83 additions & 1 deletion

File tree

packages/memory-host-sdk/src/host/batch-upload.test.ts

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,19 @@ function streamingTextResponse(params: {
3131
return new Response(stream, { status: params.status, headers: params.headers });
3232
}
3333

34+
function stallingResponse(params: { status: number; onCancel: () => void }): Response {
35+
const stream = new ReadableStream<Uint8Array>({
36+
start() {
37+
// Leave the body open until the caller aborts.
38+
},
39+
cancel() {
40+
params.onCancel();
41+
},
42+
});
43+
44+
return new Response(stream, { status: params.status });
45+
}
46+
3447
describe("uploadBatchJsonlFile", () => {
3548
beforeEach(() => {
3649
vi.clearAllMocks();
@@ -108,4 +121,70 @@ describe("uploadBatchJsonlFile", () => {
108121
).rejects.toThrow("file upload failed: response body too large: 64 bytes (limit: 8 bytes)");
109122
expect(canceled).toBe(true);
110123
});
124+
125+
it("passes caller abort signals through non-ok file-upload response snippets", async () => {
126+
let canceled = false;
127+
remoteHttpMock.mockImplementationOnce(async (params) => {
128+
return await params.onResponse(
129+
stallingResponse({
130+
status: 500,
131+
onCancel: () => {
132+
canceled = true;
133+
},
134+
}),
135+
);
136+
});
137+
const controller = new AbortController();
138+
const upload = uploadBatchJsonlFile({
139+
client: {
140+
baseUrl: "https://memory.example/v1",
141+
headers: { Authorization: "Bearer test" },
142+
},
143+
requests: [{ input: "one" }],
144+
errorPrefix: "file upload failed",
145+
signal: controller.signal,
146+
});
147+
148+
await new Promise((resolve) => {
149+
setTimeout(resolve, 0);
150+
});
151+
controller.abort(new Error("upload aborted"));
152+
153+
await expect(upload).rejects.toThrow("upload aborted");
154+
expect(canceled).toBe(true);
155+
expect(remoteHttpMock.mock.calls[0]?.[0].signal).toBe(controller.signal);
156+
});
157+
158+
it("passes caller abort signals through successful file-upload JSON reads", async () => {
159+
let canceled = false;
160+
remoteHttpMock.mockImplementationOnce(async (params) => {
161+
return await params.onResponse(
162+
stallingResponse({
163+
status: 200,
164+
onCancel: () => {
165+
canceled = true;
166+
},
167+
}),
168+
);
169+
});
170+
const controller = new AbortController();
171+
const upload = uploadBatchJsonlFile({
172+
client: {
173+
baseUrl: "https://memory.example/v1",
174+
headers: { Authorization: "Bearer test" },
175+
},
176+
requests: [{ input: "one" }],
177+
errorPrefix: "file upload failed",
178+
signal: controller.signal,
179+
});
180+
181+
await new Promise((resolve) => {
182+
setTimeout(resolve, 0);
183+
});
184+
controller.abort(new Error("upload json aborted"));
185+
186+
await expect(upload).rejects.toThrow("upload json aborted");
187+
expect(canceled).toBe(true);
188+
expect(remoteHttpMock.mock.calls[0]?.[0].signal).toBe(controller.signal);
189+
});
111190
});

packages/memory-host-sdk/src/host/batch-upload.ts

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ export async function uploadBatchJsonlFile(params: {
1616
requests: unknown[];
1717
errorPrefix: string;
1818
maxResponseBytes?: number;
19+
signal?: AbortSignal;
1920
}): Promise<string> {
2021
const baseUrl = normalizeBatchBaseUrl(params.client);
2122
const jsonl = params.requests.map((request) => JSON.stringify(request)).join("\n");
@@ -31,19 +32,21 @@ export async function uploadBatchJsonlFile(params: {
3132
url: `${baseUrl}/files`,
3233
ssrfPolicy: params.client.ssrfPolicy,
3334
fetchImpl: params.client.fetchImpl,
35+
signal: params.signal,
3436
init: {
3537
method: "POST",
3638
headers: buildBatchHeaders(params.client, { json: false }),
3739
body: form,
3840
},
3941
onResponse: async (fileRes) => {
4042
if (!fileRes.ok) {
41-
const text = await readResponseTextSnippet(fileRes);
43+
const text = await readResponseTextSnippet(fileRes, { signal: params.signal });
4244
throw new Error(`${params.errorPrefix}: ${fileRes.status} ${text}`);
4345
}
4446
return (await readResponseJsonWithLimit(fileRes, {
4547
errorPrefix: params.errorPrefix,
4648
maxBytes: params.maxResponseBytes,
49+
signal: params.signal,
4750
})) as { id?: string };
4851
},
4952
});

0 commit comments

Comments
 (0)