@copilotkit/aimock 1.43.0 → 1.45.0

package/dist/ws-gemini-live.js

@@ -1,6 +1,7 @@
 import { DEFAULT_TEST_ID } from "./constants.js";
-import { flattenHeaders, formatToMime, generateToolCallId, isAudioResponse, isContentWithToolCallsResponse, isErrorResponse, isTextResponse, isToolCallResponse, resolveFixtureBlocks, resolveResponse, resolveStrictMode, strictNoMatchLogLine, strictNoMatchMessage, strictOverrideField } from "./helpers.js";
+import { fixtureToolCallErrorCode, flattenHeaders, formatToMime, generateToolCallId, googleFixtureToolCallErrorDetails, isAudioResponse, isCombinedFixtureResponse, isErrorResponse, isFixtureToolCallError, isTextResponse, isToolCallResponse, journalFixtureToolCallError, rejectResponsesOnlyToolCalls, requireServedFunctionToolCalls, resolveFixtureBlockOutcome, resolveFixtureBlocks, resolveResponse, resolveStrictMode, runWithToolArgumentsScope, servedToolArgs, strictNoMatchLogLine, strictNoMatchMessage, strictOverrideField, strictToolArgumentsEnabled, toolArgsForWire } from "./helpers.js";
 import { calculateDelay, delay } from "./sse-writer.js";
+import { planMisbehavior, recordMisbehaviorOutcome } from "./misbehavior.js";
 import { matchFixtureDiagnostic } from "./router.js";
 import "./journal.js";
 import { createInterruptionSignal } from "./interruption.js";
@@ -102,10 +103,94 @@ function convertTools(geminiTools) {
 		function: {
 			name: d.name,
 			description: d.description,
-			parameters: d.parameters
+			parameters: d.parameters ?? d.parametersJsonSchema
 		}
 	}));
 }
+function liveToolArguments(tc) {
+	return servedToolArgs(tc, "object").value;
+}
+const GEMINI_LIVE_WIRE = "Gemini Live";
+/** Prepare native object arguments and IDs once for both output and the journal. */
+function prepareLiveMisbehavior(plan) {
+	if (plan.summary.fault === "tool-args-invalid-json" || plan.summary.fault === "empty-response") return {
+		...plan,
+		summary: {
+			...plan.summary,
+			servedToolCalls: []
+		}
+	};
+	const response = plan.response;
+	const combined = isCombinedFixtureResponse(response);
+	const outcome = combined && response.blocks?.length ? resolveFixtureBlockOutcome(response.blocks) : void 0;
+	const calls = outcome?.toolCalls ?? (combined || isToolCallResponse(response) ? response.toolCalls ?? [] : []);
+	const duplicate = plan.duplicateId;
+	const toolCalls = calls.map((call, index) => {
+		const args = toolArgsForWire(call);
+		return {
+			...call,
+			arguments: args.kind === "parsed" ? args.text : args.raw,
+			...index === duplicate?.destinationIndex ? {} : { id: call.id ?? generateToolCallId() }
+		};
+	});
+	if (duplicate) toolCalls[duplicate.destinationIndex].id = toolCalls[duplicate.sourceIndex].id;
+	let preparedResponse = { ...response };
+	if ("usage" in preparedResponse) delete preparedResponse.usage;
+	if (outcome) {
+		let index = 0;
+		const blocks = outcome.ordered.map((block) => block.type === "text" ? { ...block } : {
+			...block,
+			...toolCalls[index++]
+		});
+		preparedResponse = {
+			...preparedResponse,
+			content: outcome.content,
+			toolCalls,
+			blocks
+		};
+	} else if (combined || isToolCallResponse(response)) preparedResponse = {
+		...preparedResponse,
+		toolCalls
+	};
+	return {
+		...plan,
+		response: preparedResponse,
+		summary: {
+			...plan.summary,
+			servedToolCalls: toolCalls.map(({ name, arguments: args, id }) => ({
+				name,
+				arguments: args,
+				id
+			}))
+		}
+	};
+}
+/**
+* Validate the whole selected response before audio, text, or any tool is
+* sent: every entry is validated, custom tool calls are rejected, and every
+* argument string must parse. Returns the calls narrowed to function calls. On
+* failure the journal entry becomes a 500, and the outer message handler sends
+* code 13 with `google.rpc.ErrorInfo` details carrying the aimock code, drops
+* this turn from the history, and leaves the socket open.
+*/
+function preflightToolArguments(response, journalEntry) {
+	return journalFixtureToolCallError(journalEntry, () => {
+		const functionCalls = requireServedFunctionToolCalls(response, GEMINI_LIVE_WIRE);
+		if (strictToolArgumentsEnabled()) for (const tc of functionCalls) liveToolArguments(tc);
+		return functionCalls;
+	});
+}
+/** Resolve ordered blocks (rejecting custom tool calls) and preflight their arguments. */
+function preflightBlocks(response, journalEntry) {
+	return journalFixtureToolCallError(journalEntry, () => {
+		rejectResponsesOnlyToolCalls(response, GEMINI_LIVE_WIRE);
+		const resolved = resolveFixtureBlocks(response.blocks ?? []);
+		if (strictToolArgumentsEnabled()) {
+			for (const block of resolved) if (block.type === "toolCall") liveToolArguments(block);
+		}
+		return resolved;
+	});
+}
 function handleWebSocketGeminiLive(ws, fixtures, journal, defaults) {
 	const { logger } = defaults;
 	const session = {
@@ -116,19 +201,31 @@ function handleWebSocketGeminiLive(ws, fixtures, journal, defaults) {
 	};
 	let pending = Promise.resolve();
 	ws.on("message", (raw) => {
-		pending = pending.then(() => processMessage(raw, ws, fixtures, journal, defaults, session).catch((err) => {
-			const msg = err instanceof Error ? err.message : "Internal error";
-			logger.error(`WebSocket Gemini Live error: ${msg}`);
+		pending = pending.then(async () => {
+			const historyMark = session.conversationHistory.length;
 			try {
-				ws.send(JSON.stringify({ error: {
-					code: 13,
-					message: msg,
-					status: "INTERNAL"
-				} }));
-			} catch (sendErr) {
-				defaults.logger.debug(`Failed to send error to client: ${sendErr instanceof Error ? sendErr.message : "unknown"}`);
+				defaults.beforeProcessMessage?.();
+				await runWithToolArgumentsScope({
+					strict: defaults.strictToolArguments === true,
+					logger
+				}, () => processMessage(raw, ws, fixtures, journal, defaults, session));
+			} catch (err) {
+				const msg = err instanceof Error ? err.message : "Internal error";
+				logger.error(`WebSocket Gemini Live error: ${msg}`);
+				const fixtureToolCallError = isFixtureToolCallError(err);
+				if (fixtureToolCallError) session.conversationHistory.length = historyMark;
+				try {
+					ws.send(JSON.stringify({ error: {
+						code: 13,
+						message: msg,
+						status: "INTERNAL",
+						...fixtureToolCallError ? { details: googleFixtureToolCallErrorDetails(fixtureToolCallErrorCode(err)) } : {}
+					} }));
+				} catch (sendErr) {
+					defaults.logger.debug(`Failed to send error to client: ${sendErr instanceof Error ? sendErr.message : "unknown"}`);
+				}
 			}
-		}));
+		});
 	});
 }
 async function processMessage(raw, ws, fixtures, journal, defaults, session) {
@@ -240,12 +337,28 @@ async function processMessage(raw, ws, fixtures, journal, defaults, session) {
 		return;
 	}
 	session.conversationHistory.push(...newMessages);
-	const response = await resolveResponse(fixture, completionReq);
-	const latency = fixture.latency ?? defaults.latency;
-	const chunkSize = Math.max(1, fixture.chunkSize ?? defaults.chunkSize);
-	if (isErrorResponse(response)) {
-		const status = response.status ?? 500;
-		journal.add({
+	let response = await resolveResponse(fixture, completionReq);
+	const plannerDefaults = {
+		...defaults,
+		replaySpeed: defaults.replaySpeed ?? 1
+	};
+	let misbehavior = planMisbehavior({
+		wire: "gemini-live",
+		emitsToolCallIds: true,
+		fixture,
+		response,
+		request: completionReq,
+		stream: true,
+		defaults: plannerDefaults,
+		rawHeaders: defaults.upgradeHeaders ?? {},
+		url: `${path}?testId=${encodeURIComponent(testId)}`
+	});
+	if (misbehavior.kind === "applied") {
+		misbehavior = prepareLiveMisbehavior(misbehavior);
+		response = misbehavior.response;
+	}
+	const addResponseEntry = (status) => {
+		const entry = journal.add({
 			method: "WS",
 			path,
 			headers: flattenHeaders(defaults.upgradeHeaders ?? {}),
@@ -255,6 +368,68 @@ async function processMessage(raw, ws, fixtures, journal, defaults, session) {
 				fixture
 			}
 		});
+		recordMisbehaviorOutcome({
+			entry,
+			summary: misbehavior.summary,
+			defaults: plannerDefaults,
+			testId
+		});
+		return entry;
+	};
+	if (misbehavior.kind === "error") {
+		addResponseEntry(misbehavior.status);
+		if (!misbehavior.summary?.evaluations.length) defaults.logger.error(`[misbehavior] ${misbehavior.message} (testId=${testId})`);
+		ws.send(JSON.stringify({ error: {
+			code: httpToGrpc(misbehavior.status),
+			message: misbehavior.message,
+			status: misbehavior.status === 400 ? "INVALID_ARGUMENT" : "UNIMPLEMENTED"
+		} }));
+		return;
+	}
+	const latency = fixture.latency ?? defaults.latency;
+	if (misbehavior.kind === "applied" && (misbehavior.summary.fault === "tool-args-invalid-json" || misbehavior.summary.fault === "empty-response")) {
+		const journalEntry = addResponseEntry(200);
+		const interruption = createInterruptionSignal(fixture);
+		const markInterrupted = () => {
+			journalEntry.response.interrupted = true;
+			journalEntry.response.interruptReason = interruption?.reason();
+			ws.destroy();
+		};
+		try {
+			if (ws.isClosed) return;
+			const chunkDelay = calculateDelay(0, void 0, latency, fixture.recordedTimings, fixture.replaySpeed ?? defaults.replaySpeed);
+			if (chunkDelay > 0) await delay(chunkDelay, interruption?.signal);
+			if (interruption?.signal.aborted) {
+				markInterrupted();
+				return;
+			}
+			if (ws.isClosed) return;
+			const malformed = misbehavior.summary.fault === "tool-args-invalid-json";
+			ws.send(JSON.stringify({ serverContent: malformed ? {
+				turnComplete: true,
+				turnCompleteReason: "MALFORMED_FUNCTION_CALL"
+			} : { modelTurn: { parts: [{ text: "" }] } } }));
+			interruption?.tick();
+			if (interruption?.signal.aborted) {
+				markInterrupted();
+				return;
+			}
+			if (!malformed && !ws.isClosed) {
+				ws.send(JSON.stringify({ serverContent: { turnComplete: true } }));
+				session.conversationHistory.push({
+					role: "assistant",
+					content: ""
+				});
+			}
+		} finally {
+			interruption?.cleanup();
+		}
+		return;
+	}
+	const chunkSize = Math.max(1, fixture.chunkSize ?? defaults.chunkSize);
+	if (isErrorResponse(response)) {
+		const status = response.status ?? 500;
+		addResponseEntry(status);
 		ws.send(JSON.stringify({ error: {
 			code: httpToGrpc(status),
 			message: response.error.message,
@@ -263,17 +438,9 @@ async function processMessage(raw, ws, fixtures, journal, defaults, session) {
 		return;
 	}
 	if (isAudioResponse(response)) {
-		journal.add({
-			method: "WS",
-			path,
-			headers: flattenHeaders(defaults.upgradeHeaders ?? {}),
-			body: completionReq,
-			response: {
-				status: 200,
-				fixture
-			}
-		});
+		const journalEntry = addResponseEntry(200);
 		const audioResp = response;
+		const audioToolCalls = preflightToolArguments(audioResp, journalEntry);
 		let mimeType;
 		let data;
 		if (typeof audioResp.audio === "string") {
@@ -288,7 +455,7 @@ async function processMessage(raw, ws, fixtures, journal, defaults, session) {
 			data
 		} }];
 		if (audioResp.content) parts.push({ text: audioResp.content });
-		const resolvedToolCalls = (audioResp.toolCalls ?? []).map((tc) => ({
+		const resolvedToolCalls = audioToolCalls.map((tc) => ({
 			...tc,
 			resolvedId: tc.id ?? generateToolCallId()
 		}));
@@ -300,13 +467,7 @@ async function processMessage(raw, ws, fixtures, journal, defaults, session) {
 			ws.send(JSON.stringify({ serverContent: { modelTurn: { parts } } }));
 			if (!ws.isClosed) {
 				const functionCalls = resolvedToolCalls.map((tc) => {
-					let argsObj;
-					try {
-						argsObj = JSON.parse(tc.arguments || "{}");
-					} catch {
-						defaults.logger.warn(`Malformed JSON in fixture tool call arguments for "${tc.name}": ${tc.arguments}`);
-						argsObj = {};
-					}
+					const argsObj = liveToolArguments(tc);
 					return {
 						name: tc.name,
 						args: argsObj,
@@ -331,19 +492,10 @@ async function processMessage(raw, ws, fixtures, journal, defaults, session) {
 		});
 		return;
 	}
-	if (isContentWithToolCallsResponse(response)) {
-		const journalEntry = journal.add({
-			method: "WS",
-			path,
-			headers: flattenHeaders(defaults.upgradeHeaders ?? {}),
-			body: completionReq,
-			response: {
-				status: 200,
-				fixture
-			}
-		});
+	if (isCombinedFixtureResponse(response)) {
+		const journalEntry = addResponseEntry(200);
 		if (response.blocks && response.blocks.length > 0) {
-			const resolvedBlocks = resolveFixtureBlocks(response.blocks);
+			const resolvedBlocks = preflightBlocks(response, journalEntry);
 			const interruption = createInterruptionSignal(fixture);
 			const replaySpeed = fixture.replaySpeed ?? defaults.replaySpeed;
 			const { recordedTimings } = fixture;
@@ -361,13 +513,7 @@ async function processMessage(raw, ws, fixtures, journal, defaults, session) {
 				}
 				if (ws.isClosed) break;
 				const resolvedId = block.id ?? generateToolCallId();
-				let argsObj;
-				try {
-					argsObj = JSON.parse(block.arguments || "{}");
-				} catch {
-					defaults.logger.warn(`Malformed JSON in fixture tool call arguments for "${block.name}": ${block.arguments}`);
-					argsObj = {};
-				}
+				const argsObj = liveToolArguments(block);
 				try {
 					ws.send(JSON.stringify({ toolCall: { functionCalls: [{
 						name: block.name,
@@ -435,6 +581,7 @@ async function processMessage(raw, ws, fixtures, journal, defaults, session) {
 			});
 			return;
 		}
+		const functionToolCalls = preflightToolArguments(response, journalEntry);
 		const content = response.content ?? "";
 		const chunkList = [];
 		for (let i = 0; i < content.length; i += chunkSize) chunkList.push(content.slice(i, i + chunkSize));
@@ -472,7 +619,7 @@ async function processMessage(raw, ws, fixtures, journal, defaults, session) {
 			interruption?.cleanup();
 			return;
 		}
-		const resolvedToolCalls = (response.toolCalls ?? []).map((tc) => ({
+		const resolvedToolCalls = functionToolCalls.map((tc) => ({
 			...tc,
 			resolvedId: tc.id ?? generateToolCallId()
 		}));
@@ -487,13 +634,7 @@ async function processMessage(raw, ws, fixtures, journal, defaults, session) {
 				return;
 			}
 			const functionCalls = resolvedToolCalls.map((tc) => {
-				let argsObj;
-				try {
-					argsObj = JSON.parse(tc.arguments || "{}");
-				} catch {
-					defaults.logger.warn(`Malformed JSON in fixture tool call arguments for "${tc.name}": ${tc.arguments}`);
-					argsObj = {};
-				}
+				const argsObj = liveToolArguments(tc);
 				return {
 					name: tc.name,
 					args: argsObj,
@@ -527,16 +668,7 @@ async function processMessage(raw, ws, fixtures, journal, defaults, session) {
 		return;
 	}
 	if (isTextResponse(response)) {
-		const journalEntry = journal.add({
-			method: "WS",
-			path,
-			headers: flattenHeaders(defaults.upgradeHeaders ?? {}),
-			body: completionReq,
-			response: {
-				status: 200,
-				fixture
-			}
-		});
+		const journalEntry = addResponseEntry(200);
 		const content = response.content;
 		if (content.length === 0) {
 			if (ws.isClosed) return;
@@ -587,16 +719,8 @@ async function processMessage(raw, ws, fixtures, journal, defaults, session) {
 		return;
 	}
 	if (isToolCallResponse(response)) {
-		const journalEntry = journal.add({
-			method: "WS",
-			path,
-			headers: flattenHeaders(defaults.upgradeHeaders ?? {}),
-			body: completionReq,
-			response: {
-				status: 200,
-				fixture
-			}
-		});
+		const journalEntry = addResponseEntry(200);
+		const functionToolCalls = preflightToolArguments(response, journalEntry);
 		const interruption = createInterruptionSignal(fixture);
 		const replaySpeed = fixture.replaySpeed ?? defaults.replaySpeed;
 		const { recordedTimings } = fixture;
@@ -617,18 +741,12 @@ async function processMessage(raw, ws, fixtures, journal, defaults, session) {
 			interruption?.cleanup();
 			return;
 		}
-		const resolvedToolCalls = (response.toolCalls ?? []).map((tc) => ({
+		const resolvedToolCalls = functionToolCalls.map((tc) => ({
 			...tc,
 			resolvedId: tc.id ?? generateToolCallId()
 		}));
 		const functionCalls = resolvedToolCalls.map((tc) => {
-			let argsObj;
-			try {
-				argsObj = JSON.parse(tc.arguments || "{}");
-			} catch {
-				defaults.logger.warn(`Malformed JSON in fixture tool call arguments for "${tc.name}": ${tc.arguments}`);
-				argsObj = {};
-			}
+			const argsObj = liveToolArguments(tc);
 			return {
 				name: tc.name,
 				args: argsObj,
@@ -660,16 +778,7 @@ async function processMessage(raw, ws, fixtures, journal, defaults, session) {
 		});
 		return;
 	}
-	journal.add({
-		method: "WS",
-		path,
-		headers: flattenHeaders(defaults.upgradeHeaders ?? {}),
-		body: completionReq,
-		response: {
-			status: 500,
-			fixture
-		}
-	});
+	addResponseEntry(500);
 	ws.send(JSON.stringify({ error: {
 		code: 13,
 		message: "Fixture response did not match any known type",