diff --git a/extensions/amazon-bedrock/stream.runtime.test.ts b/extensions/amazon-bedrock/stream.runtime.test.ts index 69d03f400dcd..08e4452506c8 100644 --- a/extensions/amazon-bedrock/stream.runtime.test.ts +++ b/extensions/amazon-bedrock/stream.runtime.test.ts @@ -67,6 +67,104 @@ afterEach(() => { vi.restoreAllMocks(); }); +describe("Bedrock stream client lifecycle", () => { + const context = { + messages: [{ role: "user", content: "Hello", timestamp: 0 }], + } as never; + + function expectDestroyedClient( + send: ReturnType, + destroy: ReturnType, + ) { + expect(send).toHaveBeenCalledOnce(); + expect(destroy).toHaveBeenCalledOnce(); + expect(destroy.mock.contexts[0]).toBe(send.mock.contexts[0]); + expect(destroy.mock.invocationCallOrder[0]).toBeGreaterThan( + send.mock.invocationCallOrder[0] ?? 0, + ); + } + + it("destroys the client after a successful stream", async () => { + let markStreamBlocked!: () => void; + const streamBlocked = new Promise((resolve) => { + markStreamBlocked = resolve; + }); + let releaseStream!: () => void; + const streamReleased = new Promise((resolve) => { + releaseStream = resolve; + }); + async function* successfulStream() { + yield { messageStart: { role: ConversationRole.ASSISTANT } }; + markStreamBlocked(); + await streamReleased; + yield { messageStop: { stopReason: BedrockStopReason.END_TURN } }; + } + const send = vi.spyOn(BedrockRuntimeClient.prototype, "send").mockResolvedValue({ + $metadata: { httpStatusCode: 200 }, + stream: successfulStream(), + } as never); + const destroy = vi.spyOn(BedrockRuntimeClient.prototype, "destroy"); + + const resultPromise = streamBedrockForTest(bedrockModel({}), context).result(); + await streamBlocked; + expect(destroy).not.toHaveBeenCalled(); + + releaseStream(); + const result = await resultPromise; + + expect(result.stopReason).toBe("stop"); + expectDestroyedClient(send, destroy); + }); + + it("destroys the client after a provider error", async () => { + const send = vi + .spyOn(BedrockRuntimeClient.prototype, "send") + .mockRejectedValue(new Error("synthetic provider failure")); + const destroy = vi.spyOn(BedrockRuntimeClient.prototype, "destroy"); + + const result = await streamBedrockForTest(bedrockModel({}), context).result(); + + expect(result.stopReason).toBe("error"); + expect(result.errorMessage).toBe("synthetic provider failure"); + expectDestroyedClient(send, destroy); + }); + + it("destroys the client when response stream iteration fails", async () => { + async function* failingStream() { + yield { messageStart: { role: ConversationRole.ASSISTANT } }; + throw new Error("synthetic iterator failure"); + } + const send = vi.spyOn(BedrockRuntimeClient.prototype, "send").mockResolvedValue({ + $metadata: { httpStatusCode: 200 }, + stream: failingStream(), + } as never); + const destroy = vi.spyOn(BedrockRuntimeClient.prototype, "destroy"); + + const result = await streamBedrockForTest(bedrockModel({}), context).result(); + + expect(result.stopReason).toBe("error"); + expect(result.errorMessage).toBe("synthetic iterator failure"); + expectDestroyedClient(send, destroy); + }); + + it("destroys the client after an aborted request", async () => { + const controller = new AbortController(); + controller.abort(); + const send = vi + .spyOn(BedrockRuntimeClient.prototype, "send") + .mockRejectedValue(new Error("synthetic abort")); + const destroy = vi.spyOn(BedrockRuntimeClient.prototype, "destroy"); + + const result = await streamBedrockForTest(bedrockModel({}), context, { + signal: controller.signal, + }).result(); + + expect(result.stopReason).toBe("aborted"); + expect(result.errorMessage).toBe("synthetic abort"); + expectDestroyedClient(send, destroy); + }); +}); + describe("Bedrock inbound image base64", () => { const model = () => bedrockModel({ input: ["text", "image"] }); const userImage = (data: string) => diff --git a/extensions/amazon-bedrock/stream.runtime.ts b/extensions/amazon-bedrock/stream.runtime.ts index 1e1e9cc15b9d..688e9e6f1b81 100644 --- a/extensions/amazon-bedrock/stream.runtime.ts +++ b/extensions/amazon-bedrock/stream.runtime.ts @@ -239,8 +239,9 @@ const streamBedrock: StreamFunction<"bedrock-converse-stream", BedrockOptions> = config.authSchemePreference = ["httpBearerAuth"]; } + let client: BedrockRuntimeClient | undefined; try { - const client = new BedrockRuntimeClient(config); + client = new BedrockRuntimeClient(config); const cacheRetention = resolveCacheRetention(options.cacheRetention); const additionalModelRequestFields = buildAdditionalModelRequestFields(model, options); const thinking = (additionalModelRequestFields as Record | undefined) @@ -373,6 +374,9 @@ const streamBedrock: StreamFunction<"bedrock-converse-stream", BedrockOptions> = output.errorMessage = formatBedrockError(error); stream.push({ type: "error", reason: output.stopReason, error: output }); stream.end(); + } finally { + // The SDK client owns pooled HTTP resources; release them only after its async stream settles. + client?.destroy(); } })();