@copilotkit/aimock 1.46.0 → 1.47.0

package/dist/ws-gemini-live.js

@@ -5,7 +5,7 @@ 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";
+import { markWsTurnStopped, watchWsClose, wsLifecycle } from "./response/lifecycle.js";
 
 //#region src/ws-gemini-live.ts
 /**
@@ -201,6 +201,7 @@ function preflightBlocks(response, journalEntry) {
 }
 function handleWebSocketGeminiLive(ws, fixtures, journal, defaults) {
 	const { logger } = defaults;
+	watchWsClose(ws);
 	const session = {
 		setupDone: false,
 		model: defaults.model,
@@ -397,40 +398,45 @@ async function processMessage(raw, ws, fixtures, journal, defaults, session) {
 	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();
+		const lifecycle = wsLifecycle(ws, fixture);
+		const interruptTurn = () => {
+			markWsTurnStopped(journalEntry, ws, lifecycle);
 			ws.destroy();
 		};
 		try {
-			if (ws.isClosed) return;
+			if (ws.isClosed) {
+				interruptTurn();
+				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();
+			if (chunkDelay > 0) await delay(chunkDelay, lifecycle.signal);
+			if (lifecycle.signal.aborted || ws.isClosed) {
+				interruptTurn();
 				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();
+			lifecycle.tick();
+			if (lifecycle.signal.aborted || ws.isClosed) {
+				interruptTurn();
 				return;
 			}
-			if (!malformed && !ws.isClosed) {
+			if (!malformed) {
 				ws.send(JSON.stringify({ serverContent: { turnComplete: true } }));
+				if (ws.isClosed) {
+					interruptTurn();
+					return;
+				}
 				session.conversationHistory.push({
 					role: "assistant",
 					content: ""
 				});
 			}
 		} finally {
-			interruption?.cleanup();
+			lifecycle.dispose();
 		}
 		return;
 	}
@@ -467,25 +473,32 @@ async function processMessage(raw, ws, fixtures, journal, defaults, session) {
 			...tc,
 			resolvedId: tc.id ?? generateToolCallId()
 		}));
+		const stoppedByClose = () => {
+			if (!ws.isClosed) return false;
+			markWsTurnStopped(journalEntry, ws);
+			return true;
+		};
+		if (stoppedByClose()) return;
 		if (resolvedToolCalls.length === 0) ws.send(JSON.stringify({ serverContent: {
 			modelTurn: { parts },
 			turnComplete: true
 		} }));
 		else {
 			ws.send(JSON.stringify({ serverContent: { modelTurn: { parts } } }));
-			if (!ws.isClosed) {
-				const functionCalls = resolvedToolCalls.map((tc) => {
-					const argsObj = liveToolArguments(tc);
-					return {
-						name: tc.name,
-						args: argsObj,
-						id: tc.resolvedId
-					};
-				});
-				ws.send(JSON.stringify({ toolCall: { functionCalls } }));
-			}
-			if (!ws.isClosed) ws.send(JSON.stringify({ serverContent: { turnComplete: true } }));
+			if (stoppedByClose()) return;
+			const functionCalls = resolvedToolCalls.map((tc) => {
+				const argsObj = liveToolArguments(tc);
+				return {
+					name: tc.name,
+					args: argsObj,
+					id: tc.resolvedId
+				};
+			});
+			ws.send(JSON.stringify({ toolCall: { functionCalls } }));
+			if (stoppedByClose()) return;
+			ws.send(JSON.stringify({ serverContent: { turnComplete: true } }));
 		}
+		if (stoppedByClose()) return;
 		session.conversationHistory.push({
 			role: "assistant",
 			content: audioResp.content ?? "[audio]",
@@ -504,7 +517,7 @@ async function processMessage(raw, ws, fixtures, journal, defaults, session) {
 		const journalEntry = addResponseEntry(200);
 		if (response.blocks && response.blocks.length > 0) {
 			const resolvedBlocks = preflightBlocks(response, journalEntry);
-			const interruption = createInterruptionSignal(fixture);
+			const lifecycle = wsLifecycle(ws, fixture);
 			const replaySpeed = fixture.replaySpeed ?? defaults.replaySpeed;
 			const { recordedTimings } = fixture;
 			let chunkIndex = 0;
@@ -514,8 +527,8 @@ async function processMessage(raw, ws, fixtures, journal, defaults, session) {
 			outer: for (const block of resolvedBlocks) if (block.type === "toolCall") {
 				if (ws.isClosed) break;
 				const tcDelay = calculateDelay(chunkIndex, void 0, latency, recordedTimings, replaySpeed);
-				if (tcDelay > 0) await delay(tcDelay, interruption?.signal);
-				if (interruption?.signal.aborted) {
+				if (tcDelay > 0) await delay(tcDelay, lifecycle.signal);
+				if (lifecycle.signal.aborted) {
 					interrupted = true;
 					break;
 				}
@@ -541,8 +554,8 @@ async function processMessage(raw, ws, fixtures, journal, defaults, session) {
 						arguments: block.arguments
 					}
 				});
-				interruption?.tick();
-				if (interruption?.signal.aborted) {
+				lifecycle.tick();
+				if (lifecycle.signal.aborted) {
 					interrupted = true;
 					break;
 				}
@@ -553,8 +566,8 @@ async function processMessage(raw, ws, fixtures, journal, defaults, session) {
 				for (let i = 0; i < text.length; i += chunkSize) {
 					if (ws.isClosed) break outer;
 					const chunkDelay = calculateDelay(chunkIndex, void 0, latency, recordedTimings, replaySpeed);
-					if (chunkDelay > 0) await delay(chunkDelay, interruption?.signal);
-					if (interruption?.signal.aborted) {
+					if (chunkDelay > 0) await delay(chunkDelay, lifecycle.signal);
+					if (lifecycle.signal.aborted) {
 						interrupted = true;
 						break outer;
 					}
@@ -566,22 +579,25 @@ async function processMessage(raw, ws, fixtures, journal, defaults, session) {
 						break outer;
 					}
 					chunkIndex++;
-					interruption?.tick();
-					if (interruption?.signal.aborted) {
+					lifecycle.tick();
+					if (lifecycle.signal.aborted) {
 						interrupted = true;
 						break outer;
 					}
 				}
 			}
-			if (interrupted) {
+			if (interrupted || ws.isClosed) {
 				ws.destroy();
-				journalEntry.response.interrupted = true;
-				journalEntry.response.interruptReason = interruption?.reason();
-				interruption?.cleanup();
+				markWsTurnStopped(journalEntry, ws, lifecycle);
+				lifecycle.dispose();
+				return;
+			}
+			lifecycle.dispose();
+			ws.send(JSON.stringify({ serverContent: { turnComplete: true } }));
+			if (ws.isClosed) {
+				markWsTurnStopped(journalEntry, ws, lifecycle);
 				return;
 			}
-			interruption?.cleanup();
-			if (!ws.isClosed) ws.send(JSON.stringify({ serverContent: { turnComplete: true } }));
 			session.conversationHistory.push({
 				role: "assistant",
 				content: historyContent || null,
@@ -593,7 +609,7 @@ async function processMessage(raw, ws, fixtures, journal, defaults, session) {
 		const content = response.content ?? "";
 		const chunkList = [];
 		for (let i = 0; i < content.length; i += chunkSize) chunkList.push(content.slice(i, i + chunkSize));
-		const interruption = createInterruptionSignal(fixture);
+		const lifecycle = wsLifecycle(ws, fixture);
 		const replaySpeed = fixture.replaySpeed ?? defaults.replaySpeed;
 		const { recordedTimings } = fixture;
 		let interrupted = false;
@@ -602,8 +618,8 @@ async function processMessage(raw, ws, fixtures, journal, defaults, session) {
 		} else for (let i = 0; i < chunkList.length; i++) {
 			if (ws.isClosed) break;
 			const chunkDelay = calculateDelay(i, void 0, latency, recordedTimings, replaySpeed);
-			if (chunkDelay > 0) await delay(chunkDelay, interruption?.signal);
-			if (interruption?.signal.aborted) {
+			if (chunkDelay > 0) await delay(chunkDelay, lifecycle.signal);
+			if (lifecycle.signal.aborted) {
 				interrupted = true;
 				break;
 			}
@@ -614,17 +630,16 @@ async function processMessage(raw, ws, fixtures, journal, defaults, session) {
 				defaults.logger.debug("[gemini-live] send failed during text streaming, closing", err);
 				break;
 			}
-			interruption?.tick();
-			if (interruption?.signal.aborted) {
+			lifecycle.tick();
+			if (lifecycle.signal.aborted) {
 				interrupted = true;
 				break;
 			}
 		}
-		if (interrupted) {
+		if (interrupted || ws.isClosed) {
 			ws.destroy();
-			journalEntry.response.interrupted = true;
-			journalEntry.response.interruptReason = interruption?.reason();
-			interruption?.cleanup();
+			markWsTurnStopped(journalEntry, ws, lifecycle);
+			lifecycle.dispose();
 			return;
 		}
 		const resolvedToolCalls = functionToolCalls.map((tc) => ({
@@ -633,12 +648,11 @@ async function processMessage(raw, ws, fixtures, journal, defaults, session) {
 		}));
 		if (!ws.isClosed) {
 			const tcDelay = calculateDelay(chunkList.length, void 0, latency, recordedTimings, replaySpeed);
-			if (tcDelay > 0) await delay(tcDelay, interruption?.signal);
-			if (interruption?.signal.aborted) {
+			if (tcDelay > 0) await delay(tcDelay, lifecycle.signal);
+			if (lifecycle.signal.aborted || ws.isClosed) {
 				ws.destroy();
-				journalEntry.response.interrupted = true;
-				journalEntry.response.interruptReason = interruption?.reason();
-				interruption?.cleanup();
+				markWsTurnStopped(journalEntry, ws, lifecycle);
+				lifecycle.dispose();
 				return;
 			}
 			const functionCalls = resolvedToolCalls.map((tc) => {
@@ -650,17 +664,20 @@ async function processMessage(raw, ws, fixtures, journal, defaults, session) {
 				};
 			});
 			ws.send(JSON.stringify({ toolCall: { functionCalls } }));
-			interruption?.tick();
+			lifecycle.tick();
 		}
-		if (interruption?.signal.aborted) {
+		if (lifecycle.signal.aborted || ws.isClosed) {
 			ws.destroy();
-			journalEntry.response.interrupted = true;
-			journalEntry.response.interruptReason = interruption?.reason();
-			interruption?.cleanup();
+			markWsTurnStopped(journalEntry, ws, lifecycle);
+			lifecycle.dispose();
+			return;
+		}
+		lifecycle.dispose();
+		ws.send(JSON.stringify({ serverContent: { turnComplete: true } }));
+		if (ws.isClosed) {
+			markWsTurnStopped(journalEntry, ws, lifecycle);
 			return;
 		}
-		interruption?.cleanup();
-		if (!ws.isClosed) ws.send(JSON.stringify({ serverContent: { turnComplete: true } }));
 		session.conversationHistory.push({
 			role: "assistant",
 			content: content || null,
@@ -679,22 +696,30 @@ async function processMessage(raw, ws, fixtures, journal, defaults, session) {
 		const journalEntry = addResponseEntry(200);
 		const content = response.content;
 		if (content.length === 0) {
-			if (ws.isClosed) return;
+			if (ws.isClosed) {
+				markWsTurnStopped(journalEntry, ws);
+				return;
+			}
 			ws.send(JSON.stringify({ serverContent: { modelTurn: { parts: [{ text: "" }] } } }));
+			if (ws.isClosed) {
+				markWsTurnStopped(journalEntry, ws);
+				return;
+			}
 			ws.send(JSON.stringify({ serverContent: { turnComplete: true } }));
+			if (ws.isClosed) markWsTurnStopped(journalEntry, ws);
 			return;
 		}
 		const chunks = [];
 		for (let i = 0; i < content.length; i += chunkSize) chunks.push(content.slice(i, i + chunkSize));
-		const interruption = createInterruptionSignal(fixture);
+		const lifecycle = wsLifecycle(ws, fixture);
 		const replaySpeed = fixture.replaySpeed ?? defaults.replaySpeed;
 		const { recordedTimings } = fixture;
 		let interrupted = false;
 		for (let i = 0; i < chunks.length; i++) {
 			if (ws.isClosed) break;
 			const chunkDelay = calculateDelay(i, void 0, latency, recordedTimings, replaySpeed);
-			if (chunkDelay > 0) await delay(chunkDelay, interruption?.signal);
-			if (interruption?.signal.aborted) {
+			if (chunkDelay > 0) await delay(chunkDelay, lifecycle.signal);
+			if (lifecycle.signal.aborted) {
 				interrupted = true;
 				break;
 			}
@@ -705,21 +730,24 @@ async function processMessage(raw, ws, fixtures, journal, defaults, session) {
 				defaults.logger.debug("[gemini-live] send failed during text streaming, closing", err);
 				break;
 			}
-			interruption?.tick();
-			if (interruption?.signal.aborted) {
+			lifecycle.tick();
+			if (lifecycle.signal.aborted) {
 				interrupted = true;
 				break;
 			}
 		}
-		if (interrupted) {
+		if (interrupted || ws.isClosed) {
 			ws.destroy();
-			journalEntry.response.interrupted = true;
-			journalEntry.response.interruptReason = interruption?.reason();
-			interruption?.cleanup();
+			markWsTurnStopped(journalEntry, ws, lifecycle);
+			lifecycle.dispose();
+			return;
+		}
+		lifecycle.dispose();
+		ws.send(JSON.stringify({ serverContent: { turnComplete: true } }));
+		if (ws.isClosed) {
+			markWsTurnStopped(journalEntry, ws, lifecycle);
 			return;
 		}
-		interruption?.cleanup();
-		if (!ws.isClosed) ws.send(JSON.stringify({ serverContent: { turnComplete: true } }));
 		session.conversationHistory.push({
 			role: "assistant",
 			content
@@ -729,24 +757,20 @@ async function processMessage(raw, ws, fixtures, journal, defaults, session) {
 	if (isToolCallResponse(response)) {
 		const journalEntry = addResponseEntry(200);
 		const functionToolCalls = preflightToolArguments(response, journalEntry);
-		const interruption = createInterruptionSignal(fixture);
+		const lifecycle = wsLifecycle(ws, fixture);
 		const replaySpeed = fixture.replaySpeed ?? defaults.replaySpeed;
 		const { recordedTimings } = fixture;
 		if (ws.isClosed) {
-			interruption?.cleanup();
+			markWsTurnStopped(journalEntry, ws, lifecycle);
+			lifecycle.dispose();
 			return;
 		}
 		const tcDelay = calculateDelay(0, void 0, latency, recordedTimings, replaySpeed);
-		if (tcDelay > 0) await delay(tcDelay, interruption?.signal);
-		if (interruption?.signal.aborted) {
+		if (tcDelay > 0) await delay(tcDelay, lifecycle.signal);
+		if (lifecycle.signal.aborted || ws.isClosed) {
 			ws.destroy();
-			journalEntry.response.interrupted = true;
-			journalEntry.response.interruptReason = interruption?.reason();
-			interruption?.cleanup();
-			return;
-		}
-		if (ws.isClosed) {
-			interruption?.cleanup();
+			markWsTurnStopped(journalEntry, ws, lifecycle);
+			lifecycle.dispose();
 			return;
 		}
 		const resolvedToolCalls = functionToolCalls.map((tc) => ({
@@ -762,16 +786,19 @@ async function processMessage(raw, ws, fixtures, journal, defaults, session) {
 			};
 		});
 		ws.send(JSON.stringify({ toolCall: { functionCalls } }));
-		interruption?.tick();
-		if (interruption?.signal.aborted) {
+		lifecycle.tick();
+		if (lifecycle.signal.aborted || ws.isClosed) {
 			ws.destroy();
-			journalEntry.response.interrupted = true;
-			journalEntry.response.interruptReason = interruption?.reason();
-			interruption?.cleanup();
+			markWsTurnStopped(journalEntry, ws, lifecycle);
+			lifecycle.dispose();
+			return;
+		}
+		lifecycle.dispose();
+		ws.send(JSON.stringify({ serverContent: { turnComplete: true } }));
+		if (ws.isClosed) {
+			markWsTurnStopped(journalEntry, ws, lifecycle);
 			return;
 		}
-		interruption?.cleanup();
-		if (!ws.isClosed) ws.send(JSON.stringify({ serverContent: { turnComplete: true } }));
 		session.conversationHistory.push({
 			role: "assistant",
 			content: null,