diff --git a/extensions/lmstudio/src/models.fetch.ts b/extensions/lmstudio/src/models.fetch.ts index 2cd2dbdd9927..8b75a6c7e731 100644 --- a/extensions/lmstudio/src/models.fetch.ts +++ b/extensions/lmstudio/src/models.fetch.ts @@ -48,6 +48,12 @@ type DiscoverLmstudioModelsParams = { fetchImpl?: typeof fetch; }; +async function cancelUnreadResponseBody(response: Response): Promise { + if (!response.bodyUsed) { + await response.body?.cancel().catch(() => undefined); + } +} + async function fetchLmstudioEndpoint(params: { url: string; init?: RequestInit; @@ -57,8 +63,10 @@ async function fetchLmstudioEndpoint(params: { auditContext: string; }): Promise<{ response: Response; release: () => Promise }> { const timeoutMs = resolveTimerTimeoutMs(params.timeoutMs, 1); + let response: Response; + let release: () => Promise; if (params.ssrfPolicy) { - return await fetchWithSsrFGuard({ + const guarded = await fetchWithSsrFGuard({ url: params.url, init: params.init, timeoutMs, @@ -66,14 +74,22 @@ async function fetchLmstudioEndpoint(params: { policy: params.ssrfPolicy, auditContext: params.auditContext, }); - } - const fetchFn = params.fetchImpl ?? fetch; - return { - response: await fetchFn(params.url, { + response = guarded.response; + release = guarded.release; + } else { + const fetchFn = params.fetchImpl ?? fetch; + response = await fetchFn(params.url, { ...params.init, signal: AbortSignal.timeout(timeoutMs), - }), - release: async () => {}, + }); + release = async () => undefined; + } + return { + response, + release: async () => { + await cancelUnreadResponseBody(response); + await release(); + }, }; } diff --git a/extensions/lmstudio/src/models.test.ts b/extensions/lmstudio/src/models.test.ts index b5ec7b1f9a0f..897447aa1362 100644 --- a/extensions/lmstudio/src/models.test.ts +++ b/extensions/lmstudio/src/models.test.ts @@ -408,6 +408,38 @@ describe("lmstudio-models", () => { }); }); + it("cancels the response body after a non-ok model discovery response", async () => { + const tracked = cancelTrackedResponse("unavailable", { status: 503 }); + const fetchMock = vi.fn(async () => tracked.response); + + const result = await fetchLmstudioModels({ + baseUrl: "http://localhost:1234/v1", + fetchImpl: asFetch(fetchMock), + }); + + expect(result).toEqual({ + reachable: true, + status: 503, + models: [], + }); + expect(tracked.wasCanceled()).toBe(true); + }); + + it("cancels guarded non-ok discovery bodies before releasing the dispatcher", async () => { + const tracked = cancelTrackedResponse("unavailable", { status: 503 }); + const release = vi.fn(async () => undefined); + fetchWithSsrFGuardMock.mockResolvedValue({ response: tracked.response, release }); + + const result = await fetchLmstudioModels({ + baseUrl: "http://localhost:1234/v1", + ssrfPolicy: {}, + }); + + expect(result).toMatchObject({ reachable: true, status: 503, models: [] }); + expect(tracked.wasCanceled()).toBe(true); + expect(release).toHaveBeenCalledOnce(); + }); + it("reports malformed model list JSON with an owned error", async () => { const fetchMock = vi.fn(async () => malformedJsonResponse());