@copilotkit/aimock 1.46.0 → 1.47.0

package/dist/bedrock.js

@@ -1,16 +1,19 @@
 import { renderError } from "./response/errors.js";
 import { resolveRequestDecisions, strictOverrideField } from "./response/policy.js";
-import { extractOverrides, flattenHeaders, generateMessageId, generateToolUseId, getContext, getTestId, isCombinedFixtureResponse, isErrorResponse, isJsonObject, isTextResponse, isToolCallResponse, requireEmittedFunctionToolCalls, resolveFixtureBlockOutcome, resolveReasoningForModel, resolveResponse, servedToolArgs, strictNoMatchLogLine, strictNoMatchMessage, validateChatMessages, validateToolsField } from "./helpers.js";
+import { extractOverrides, flattenHeaders, generateMessageId, generateToolUseId, getContext, getTestId, isCombinedFixtureResponse, isErrorResponse, isJsonObject, isRawResponse, isTextResponse, isToolCallResponse, requireEmittedFunctionToolCalls, resolveFixtureBlockOutcome, resolveReasoningForModel, resolveResponse, servedToolArgs, strictNoMatchLogLine, strictNoMatchMessage, validateChatMessages, validateToolsField } from "./helpers.js";
 import { writeErrorResponse } from "./sse-writer.js";
 import { applyChaosAsync } from "./chaos.js";
 import { planMisbehavior, recordMisbehaviorOutcome } from "./misbehavior.js";
 import { matchFixtureDiagnostic, recordMatchOptions } from "./router.js";
 import { render, resolve } from "./response/usage/codecs/anthropic.js";
 import { proxyAndRecord } from "./recorder.js";
-import { createInterruptionSignal } from "./interruption.js";
+import { misbehaviorPayloadLength, servedPayloadLength } from "./response/timing.js";
+import { deliverJson, deliverStream } from "./response/deliver.js";
+import { proxyInterruption, withHangUpDetection } from "./response/proxy.js";
+import { encodeEventStreamMessage, eventStreamCodec } from "./aws-event-stream.js";
+import { serveRawResponse } from "./response/raw.js";
 import { mapAnthropicFinish } from "./response/finish.js";
 import { prepareClaudeMisbehavior } from "./messages.js";
-import { writeEventStream } from "./aws-event-stream.js";
 
 //#region src/bedrock.ts
 function bedrockStopReason(overrideFinishReason, defaultReason) {
@@ -266,7 +269,10 @@ async function handleBedrock(req, res, raw, modelId, fixtures, journal, defaults
 	completionReq._endpointType = "chat";
 	completionReq._context = getContext(req);
 	const testId = getTestId(req);
-	const { fixture, skippedBySequenceOrTurn } = matchFixtureDiagnostic(fixtures, completionReq, journal.getFixtureMatchCountsForTest(testId), defaults.requestTransform, recordMatchOptions(!!defaults.record, defaults.logger));
+	const { fixture, skippedBySequenceOrTurn } = matchFixtureDiagnostic(fixtures, completionReq, journal.getFixtureMatchCountsForTest(testId), defaults.requestTransform, {
+		...recordMatchOptions(!!defaults.record, defaults.logger),
+		wire: "bedrock-invoke"
+	});
 	if (fixture) logger.debug(`Fixture matched: ${JSON.stringify(fixture.match).slice(0, 120)}`);
 	else logger.debug(`No fixture matched for request`);
 	if (fixture) journal.incrementFixtureMatchCount(fixture, fixtures, testId);
@@ -299,7 +305,7 @@ async function handleBedrock(req, res, raw, modelId, fixtures, journal, defaults
 			return;
 		}
 		if (defaults.record) {
-			const outcome = await proxyAndRecord(req, res, completionReq, "bedrock", urlPath, fixtures, defaults, raw);
+			const { outcome, clientHungUp } = await withHangUpDetection(req, res, () => proxyAndRecord(req, res, completionReq, "bedrock", urlPath, fixtures, defaults, raw));
 			if (outcome === "handled_by_hook") return;
 			if (outcome !== "not_configured") {
 				journal.add({
@@ -310,7 +316,8 @@ async function handleBedrock(req, res, raw, modelId, fixtures, journal, defaults
 					response: {
 						status: res.statusCode ?? 200,
 						fixture: null,
-						source: "proxy"
+						source: "proxy",
+						...proxyInterruption(res, clientHungUp)
 					}
 				});
 				return;
@@ -376,32 +383,60 @@ async function handleBedrock(req, res, raw, modelId, fixtures, journal, defaults
 		writeErrorResponse(res, misbehavior.status, renderedError.body);
 		return;
 	}
+	const chunkSize = Math.max(1, fixture.chunkSize ?? defaults.chunkSize);
+	const jsonOpts = (payloadLength, duplicateId = false) => ({
+		surface: "other",
+		payloadLength,
+		chunkSize,
+		recordedTimings: fixture.recordedTimings,
+		duplicateId,
+		replaySpeed: fixture.replaySpeed ?? defaults.replaySpeed
+	});
 	if (misbehavior.kind === "applied") {
 		const prepared = prepareClaudeMisbehavior(misbehavior, {
 			request: completionReq,
 			stream: false,
-			chunkSize: Math.max(1, fixture.chunkSize ?? defaults.chunkSize),
+			chunkSize,
 			logger,
-			strict: resolveRequestDecisions(defaults.strict, req.headers).reasoningSuppression
+			strict: resolveRequestDecisions(defaults.strict, req.headers).reasoningSuppression,
+			citations: false
 		});
 		if (prepared.stream) throw new TypeError("InvokeModel requires non-stream Claude output");
+		const entry = journal.add({
+			method: req.method ?? "POST",
+			path: urlPath,
+			headers: flattenHeaders(req.headers),
+			body: completionReq,
+			response: {
+				status: 200,
+				fixture
+			}
+		});
 		recordMisbehaviorOutcome({
-			entry: journal.add({
-				method: req.method ?? "POST",
-				path: urlPath,
-				headers: flattenHeaders(req.headers),
-				body: completionReq,
-				response: {
-					status: 200,
-					fixture
-				}
-			}),
+			entry,
 			summary: prepared.plan.summary,
 			defaults,
 			testId
 		});
-		res.writeHead(200, { "Content-Type": "application/json" });
-		res.end(JSON.stringify(prepared.body));
+		const preparedBody = prepared.body;
+		await deliverJson(res, entry, fixture, () => preparedBody, jsonOpts(misbehaviorPayloadLength(misbehavior, misbehavior.reasoning === void 0 ? prepared.reasoning : void 0), !!misbehavior.duplicateId));
+		return;
+	}
+	if (isRawResponse(response)) {
+		await serveRawResponse({
+			wire: "bedrock-invoke",
+			stream: false,
+			req,
+			res,
+			fixture,
+			response,
+			body: completionReq,
+			journal,
+			defaults,
+			testId,
+			misbehaviorSummary: misbehavior.summary,
+			errorWire: "bedrock"
+		});
 		return;
 	}
 	if (isErrorResponse(response)) {
@@ -423,7 +458,7 @@ async function handleBedrock(req, res, raw, modelId, fixtures, journal, defaults
 		if (response.webSearches?.length) logger.warn("webSearches in fixture response are not supported for Bedrock API — ignoring");
 		const overrides = extractOverrides(response);
 		const effReasoning = resolveReasoningForModel(response.reasoning, completionReq.model, resolveRequestDecisions(defaults.strict, req.headers).reasoningSuppression, logger);
-		addResponseEntry({
+		const journalEntry = addResponseEntry({
 			method: req.method ?? "POST",
 			path: urlPath,
 			headers: flattenHeaders(req.headers),
@@ -434,28 +469,27 @@ async function handleBedrock(req, res, raw, modelId, fixtures, journal, defaults
 			}
 		});
 		const functionToolCalls = requireEmittedFunctionToolCalls(response, "Bedrock InvokeModel");
-		if (response.blocks && response.blocks.length > 0) {
-			const blocksBody = buildBedrockBlocksResponse(response.blocks, completionReq.model, logger, effReasoning, overrides);
-			res.writeHead(200, { "Content-Type": "application/json" });
-			res.end(JSON.stringify(blocksBody));
+		const blocks = response.blocks;
+		if (blocks && blocks.length > 0) {
+			await deliverJson(res, journalEntry, fixture, () => buildBedrockBlocksResponse(blocks, completionReq.model, logger, effReasoning, overrides), jsonOpts(servedPayloadLength(response, effReasoning)));
 			return;
 		}
-		const textBody = buildBedrockTextResponse(response.content ?? "", completionReq.model, effReasoning, overrides);
-		const toolBody = buildBedrockToolCallResponse(functionToolCalls, completionReq.model, logger, void 0, overrides);
-		const merged = {
-			...textBody,
-			content: [...textBody.content, ...toolBody.content],
-			stop_reason: bedrockStopReason(overrides?.finishReason, "tool_use")
-		};
-		res.writeHead(200, { "Content-Type": "application/json" });
-		res.end(JSON.stringify(merged));
+		await deliverJson(res, journalEntry, fixture, () => {
+			const textBody = buildBedrockTextResponse(response.content ?? "", completionReq.model, effReasoning, overrides);
+			const toolBody = buildBedrockToolCallResponse(functionToolCalls, completionReq.model, logger, void 0, overrides);
+			return {
+				...textBody,
+				content: [...textBody.content, ...toolBody.content],
+				stop_reason: bedrockStopReason(overrides?.finishReason, "tool_use")
+			};
+		}, jsonOpts(servedPayloadLength(response, effReasoning)));
 		return;
 	}
 	if (isTextResponse(response)) {
 		if (response.webSearches?.length) logger.warn("webSearches in fixture response are not supported for Bedrock API — ignoring");
 		const overrides = extractOverrides(response);
 		const effReasoning = resolveReasoningForModel(response.reasoning, completionReq.model, resolveRequestDecisions(defaults.strict, req.headers).reasoningSuppression, logger);
-		addResponseEntry({
+		await deliverJson(res, addResponseEntry({
 			method: req.method ?? "POST",
 			path: urlPath,
 			headers: flattenHeaders(req.headers),
@@ -464,17 +498,14 @@ async function handleBedrock(req, res, raw, modelId, fixtures, journal, defaults
 				status: 200,
 				fixture
 			}
-		});
-		const body = buildBedrockTextResponse(response.content, completionReq.model, effReasoning, overrides);
-		res.writeHead(200, { "Content-Type": "application/json" });
-		res.end(JSON.stringify(body));
+		}), fixture, () => buildBedrockTextResponse(response.content, completionReq.model, effReasoning, overrides), jsonOpts(servedPayloadLength(response, effReasoning)));
 		return;
 	}
 	if (isToolCallResponse(response)) {
 		if (response.webSearches?.length) logger.warn("webSearches in fixture response are not supported for Bedrock API — ignoring");
 		const overrides = extractOverrides(response);
 		const effReasoning = resolveReasoningForModel(response.reasoning, completionReq.model, resolveRequestDecisions(defaults.strict, req.headers).reasoningSuppression, logger);
-		addResponseEntry({
+		const journalEntry = addResponseEntry({
 			method: req.method ?? "POST",
 			path: urlPath,
 			headers: flattenHeaders(req.headers),
@@ -484,9 +515,8 @@ async function handleBedrock(req, res, raw, modelId, fixtures, journal, defaults
 				fixture
 			}
 		});
-		const body = buildBedrockToolCallResponse(requireEmittedFunctionToolCalls(response, "Bedrock InvokeModel"), completionReq.model, logger, effReasoning, overrides);
-		res.writeHead(200, { "Content-Type": "application/json" });
-		res.end(JSON.stringify(body));
+		const functionToolCalls = requireEmittedFunctionToolCalls(response, "Bedrock InvokeModel");
+		await deliverJson(res, journalEntry, fixture, () => buildBedrockToolCallResponse(functionToolCalls, completionReq.model, logger, effReasoning, overrides), jsonOpts(servedPayloadLength(response, effReasoning)));
 		return;
 	}
 	addResponseEntry({
@@ -506,12 +536,7 @@ async function handleBedrock(req, res, raw, modelId, fixtures, journal, defaults
 }
 const BEDROCK_INVOKE_STREAM_EVENT_TYPE = "chunk";
 /** Serialize Claude events as InvokeModel PayloadPart, leaving generic AWS framing unchanged. */
-function writeBedrockInvokeStream(res, events, options) {
-	return writeEventStream(res, events.map((event) => ({
-		eventType: event.eventType,
-		payload: { bytes: Buffer.from(JSON.stringify(event.payload), "utf8").toString("base64") }
-	})), options);
-}
+const bedrockInvokeCodec = eventStreamCodec((event) => encodeEventStreamMessage(event.eventType, { bytes: Buffer.from(JSON.stringify(event.payload), "utf8").toString("base64") }));
 function buildBedrockInvokeMessageStart(model, overrides) {
 	return {
 		eventType: BEDROCK_INVOKE_STREAM_EVENT_TYPE,
@@ -1010,7 +1035,10 @@ async function handleBedrockStream(req, res, raw, modelId, fixtures, journal, de
 	completionReq._endpointType = "chat";
 	completionReq._context = getContext(req);
 	const testId = getTestId(req);
-	const { fixture, skippedBySequenceOrTurn } = matchFixtureDiagnostic(fixtures, completionReq, journal.getFixtureMatchCountsForTest(testId), defaults.requestTransform, recordMatchOptions(!!defaults.record, defaults.logger));
+	const { fixture, skippedBySequenceOrTurn } = matchFixtureDiagnostic(fixtures, completionReq, journal.getFixtureMatchCountsForTest(testId), defaults.requestTransform, {
+		...recordMatchOptions(!!defaults.record, defaults.logger),
+		wire: "bedrock-invoke"
+	});
 	if (fixture) logger.debug(`Fixture matched: ${JSON.stringify(fixture.match).slice(0, 120)}`);
 	else logger.debug(`No fixture matched for request`);
 	if (fixture) journal.incrementFixtureMatchCount(fixture, fixtures, testId);
@@ -1043,7 +1071,7 @@ async function handleBedrockStream(req, res, raw, modelId, fixtures, journal, de
 			return;
 		}
 		if (defaults.record) {
-			const outcome = await proxyAndRecord(req, res, completionReq, "bedrock", urlPath, fixtures, defaults, raw);
+			const { outcome, clientHungUp } = await withHangUpDetection(req, res, () => proxyAndRecord(req, res, completionReq, "bedrock", urlPath, fixtures, defaults, raw));
 			if (outcome === "handled_by_hook") return;
 			if (outcome !== "not_configured") {
 				journal.add({
@@ -1054,7 +1082,8 @@ async function handleBedrockStream(req, res, raw, modelId, fixtures, journal, de
 					response: {
 						status: res.statusCode ?? 200,
 						fixture: null,
-						source: "proxy"
+						source: "proxy",
+						...proxyInterruption(res, clientHungUp)
 					}
 				});
 				return;
@@ -1120,15 +1149,21 @@ async function handleBedrockStream(req, res, raw, modelId, fixtures, journal, de
 		writeErrorResponse(res, misbehavior.status, renderedError.body);
 		return;
 	}
-	const latency = fixture.latency ?? defaults.latency;
 	const chunkSize = Math.max(1, fixture.chunkSize ?? defaults.chunkSize);
+	const streamTiming = {
+		latency: fixture.latency ?? defaults.latency,
+		streamingProfile: fixture.streamingProfile,
+		recordedTimings: fixture.recordedTimings,
+		replaySpeed: fixture.replaySpeed ?? defaults.replaySpeed
+	};
 	if (misbehavior.kind === "applied") {
 		const prepared = prepareClaudeMisbehavior(misbehavior, {
 			request: completionReq,
 			stream: true,
 			chunkSize,
 			logger,
-			strict: resolveRequestDecisions(defaults.strict, req.headers).reasoningSuppression
+			strict: resolveRequestDecisions(defaults.strict, req.headers).reasoningSuppression,
+			citations: false
 		});
 		if (!prepared.stream) throw new TypeError("InvokeModel stream requires Claude events");
 		const entry = journal.add({
@@ -1147,27 +1182,27 @@ async function handleBedrockStream(req, res, raw, modelId, fixtures, journal, de
 			defaults,
 			testId
 		});
-		const events = prepared.events.map((event) => ({
+		await deliverStream(res, entry, fixture, bedrockInvokeCodec, prepared.events.map((event) => ({
 			eventType: BEDROCK_INVOKE_STREAM_EVENT_TYPE,
 			payload: event
-		}));
-		const interruption = createInterruptionSignal(fixture);
-		try {
-			if (!await writeBedrockInvokeStream(res, events, {
-				latency,
-				streamingProfile: fixture.streamingProfile,
-				recordedTimings: fixture.recordedTimings,
-				replaySpeed: fixture.replaySpeed ?? defaults.replaySpeed,
-				signal: interruption?.signal,
-				onChunkSent: interruption?.tick
-			})) {
-				if (!res.writableEnded) res.destroy();
-				entry.response.interrupted = true;
-				entry.response.interruptReason = interruption?.reason();
-			}
-		} finally {
-			interruption?.cleanup();
-		}
+		})), streamTiming);
+		return;
+	}
+	if (isRawResponse(response)) {
+		await serveRawResponse({
+			wire: "bedrock-invoke",
+			stream: true,
+			req,
+			res,
+			fixture,
+			response,
+			body: completionReq,
+			journal,
+			defaults,
+			testId,
+			misbehaviorSummary: misbehavior.summary,
+			errorWire: "bedrock"
+		});
 		return;
 	}
 	if (isErrorResponse(response)) {
@@ -1200,28 +1235,14 @@ async function handleBedrockStream(req, res, raw, modelId, fixtures, journal, de
 			}
 		});
 		const functionToolCalls = requireEmittedFunctionToolCalls(response, "Bedrock InvokeModel");
-		const events = buildBedrockStreamContentWithToolCallsEvents(response.content ?? "", functionToolCalls, completionReq.model, chunkSize, logger, effReasoning, overrides, response.blocks);
-		const interruption = createInterruptionSignal(fixture);
-		if (!await writeBedrockInvokeStream(res, events, {
-			latency,
-			streamingProfile: fixture.streamingProfile,
-			recordedTimings: fixture.recordedTimings,
-			replaySpeed: fixture.replaySpeed ?? defaults.replaySpeed,
-			signal: interruption?.signal,
-			onChunkSent: interruption?.tick
-		})) {
-			if (!res.writableEnded) res.destroy();
-			journalEntry.response.interrupted = true;
-			journalEntry.response.interruptReason = interruption?.reason();
-		}
-		interruption?.cleanup();
+		await deliverStream(res, journalEntry, fixture, bedrockInvokeCodec, buildBedrockStreamContentWithToolCallsEvents(response.content ?? "", functionToolCalls, completionReq.model, chunkSize, logger, effReasoning, overrides, response.blocks), streamTiming);
 		return;
 	}
 	if (isTextResponse(response)) {
 		if (response.webSearches?.length) logger.warn("webSearches in fixture response are not supported for Bedrock API — ignoring");
 		const overrides = extractOverrides(response);
 		const effReasoning = resolveReasoningForModel(response.reasoning, completionReq.model, resolveRequestDecisions(defaults.strict, req.headers).reasoningSuppression, logger);
-		const journalEntry = addResponseEntry({
+		await deliverStream(res, addResponseEntry({
 			method: req.method ?? "POST",
 			path: urlPath,
 			headers: flattenHeaders(req.headers),
@@ -1230,29 +1251,14 @@ async function handleBedrockStream(req, res, raw, modelId, fixtures, journal, de
 				status: 200,
 				fixture
 			}
-		});
-		const events = buildBedrockStreamTextEvents(response.content, completionReq.model, chunkSize, effReasoning, overrides);
-		const interruption = createInterruptionSignal(fixture);
-		if (!await writeBedrockInvokeStream(res, events, {
-			latency,
-			streamingProfile: fixture.streamingProfile,
-			recordedTimings: fixture.recordedTimings,
-			replaySpeed: fixture.replaySpeed ?? defaults.replaySpeed,
-			signal: interruption?.signal,
-			onChunkSent: interruption?.tick
-		})) {
-			if (!res.writableEnded) res.destroy();
-			journalEntry.response.interrupted = true;
-			journalEntry.response.interruptReason = interruption?.reason();
-		}
-		interruption?.cleanup();
+		}), fixture, bedrockInvokeCodec, buildBedrockStreamTextEvents(response.content, completionReq.model, chunkSize, effReasoning, overrides), streamTiming);
 		return;
 	}
 	if (isToolCallResponse(response)) {
 		if (response.webSearches?.length) logger.warn("webSearches in fixture response are not supported for Bedrock API — ignoring");
 		const overrides = extractOverrides(response);
 		const effReasoning = resolveReasoningForModel(response.reasoning, completionReq.model, resolveRequestDecisions(defaults.strict, req.headers).reasoningSuppression, logger);
-		const journalEntry = addResponseEntry({
+		await deliverStream(res, addResponseEntry({
 			method: req.method ?? "POST",
 			path: urlPath,
 			headers: flattenHeaders(req.headers),
@@ -1261,22 +1267,7 @@ async function handleBedrockStream(req, res, raw, modelId, fixtures, journal, de
 				status: 200,
 				fixture
 			}
-		});
-		const events = buildBedrockStreamToolCallEvents(requireEmittedFunctionToolCalls(response, "Bedrock InvokeModel"), completionReq.model, chunkSize, logger, effReasoning, overrides);
-		const interruption = createInterruptionSignal(fixture);
-		if (!await writeBedrockInvokeStream(res, events, {
-			latency,
-			streamingProfile: fixture.streamingProfile,
-			recordedTimings: fixture.recordedTimings,
-			replaySpeed: fixture.replaySpeed ?? defaults.replaySpeed,
-			signal: interruption?.signal,
-			onChunkSent: interruption?.tick
-		})) {
-			if (!res.writableEnded) res.destroy();
-			journalEntry.response.interrupted = true;
-			journalEntry.response.interruptReason = interruption?.reason();
-		}
-		interruption?.cleanup();
+		}), fixture, bedrockInvokeCodec, buildBedrockStreamToolCallEvents(requireEmittedFunctionToolCalls(response, "Bedrock InvokeModel"), completionReq.model, chunkSize, logger, effReasoning, overrides), streamTiming);
 		return;
 	}
 	addResponseEntry({