diff --git a/docs/record-replay.html b/docs/record-replay.html index 6e0d7e3..bdb4fc4 100644 --- a/docs/record-replay.html +++ b/docs/record-replay.html @@ -355,6 +355,55 @@

CI Pipeline Workflow

run: docker stop aimock +

Request Transform

+

+ Prompts often contain dynamic data — timestamps, UUIDs, session IDs — that + changes between runs. This causes fixture mismatches on replay because the recorded key no + longer matches the live request. The requestTransform option normalizes + requests before both matching and recording, stripping out the volatile parts. +

+ +
+
+ Strip timestamps before matching ts +
+
import { LLMock } from "@copilotkit/aimock";
+
+const mock = new LLMock({
+  requestTransform: (req) => ({
+    ...req,
+    messages: req.messages.map((m) => ({
+      ...m,
+      content:
+        typeof m.content === "string"
+          ? m.content.replace(/\d{4}-\d{2}-\d{2}T[\d:.+Z-]+/g, "")
+          : m.content,
+    })),
+  }),
+});
+
+// Fixture uses the cleaned key (no timestamp)
+mock.onMessage("tell me the weather ", { content: "Sunny" });
+
+// Request with a timestamp still matches after transform
+await mock.start();
+
+ +

+ When requestTransform is set, string matching for + userMessage and inputText switches from substring + (includes) to exact equality (===). This prevents shortened keys + from accidentally matching unrelated prompts. Without a transform, the existing + includes behavior is preserved for backward compatibility. +

+ +

+ The transform is applied in both directions: recording saves the + transformed match key (no timestamps in the fixture file), and + matching transforms the incoming request before comparison. This means + recorded fixtures and live requests always use the same normalized key. +

+

Building Fixture Sets

A practical workflow for building and maintaining fixture sets:

    diff --git a/src/__tests__/request-transform.test.ts b/src/__tests__/request-transform.test.ts new file mode 100644 index 0000000..e5c8189 --- /dev/null +++ b/src/__tests__/request-transform.test.ts @@ -0,0 +1,342 @@ +import { describe, it, expect, afterEach } from "vitest"; +import http from "node:http"; +import { matchFixture } from "../router.js"; +import { LLMock } from "../llmock.js"; +import type { ChatCompletionRequest, Fixture } from "../types.js"; + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +function makeReq(overrides: Partial = {}): ChatCompletionRequest { + return { + model: "gpt-4o", + messages: [{ role: "user", content: "hello" }], + ...overrides, + }; +} + +function makeFixture( + match: Fixture["match"], + response: Fixture["response"] = { content: "ok" }, +): Fixture { + return { match, response }; +} + +async function httpPost(url: string, body: object): Promise<{ status: number; body: string }> { + return new Promise((resolve, reject) => { + const req = http.request( + url, + { + method: "POST", + headers: { "Content-Type": "application/json" }, + }, + (res) => { + const chunks: Buffer[] = []; + res.on("data", (c) => chunks.push(c)); + res.on("end", () => + resolve({ + status: res.statusCode!, + body: Buffer.concat(chunks).toString(), + }), + ); + }, + ); + req.on("error", reject); + req.write(JSON.stringify(body)); + req.end(); + }); +} + +/** Strip ISO timestamps from text content. */ +const stripTimestamps = (req: ChatCompletionRequest): ChatCompletionRequest => ({ + ...req, + messages: req.messages.map((m) => ({ + ...m, + content: + typeof m.content === "string" + ? m.content.replace(/\d{4}-\d{2}-\d{2}T[\d:.+Z-]+/g, "") + : m.content, + })), +}); + +// --------------------------------------------------------------------------- +// Unit tests — matchFixture with requestTransform +// --------------------------------------------------------------------------- + +describe("matchFixture — requestTransform", () => { + it("matches after transform strips dynamic data", () => { + const fixture = makeFixture({ userMessage: "tell me the weather" }); + const req = makeReq({ + messages: [{ role: "user", content: "tell me the weather 2026-04-02T10:30:00.000Z" }], + }); + + // Without transform — exact match would fail, but includes works + expect(matchFixture([fixture], req)).toBe(fixture); + + // With transform — also matches (exact match against stripped text) + const transformedFixture = makeFixture({ userMessage: "tell me the weather " }); + expect(matchFixture([transformedFixture], req, undefined, stripTimestamps)).toBe( + transformedFixture, + ); + }); + + it("uses exact equality (===) when transform is provided", () => { + // Fixture matches a substring — without transform, includes would match + const fixture = makeFixture({ userMessage: "hello" }); + const req = makeReq({ + messages: [{ role: "user", content: "hello world" }], + }); + + // Without transform — includes matches + expect(matchFixture([fixture], req)).toBe(fixture); + + // With transform (identity) — exact match fails because "hello world" !== "hello" + const identity = (r: ChatCompletionRequest): ChatCompletionRequest => r; + expect(matchFixture([fixture], req, undefined, identity)).toBeNull(); + }); + + it("exact match succeeds when text matches precisely", () => { + const fixture = makeFixture({ userMessage: "hello world" }); + const req = makeReq({ + messages: [{ role: "user", content: "hello world" }], + }); + + const identity = (r: ChatCompletionRequest): ChatCompletionRequest => r; + expect(matchFixture([fixture], req, undefined, identity)).toBe(fixture); + }); + + it("preserves includes behavior when no transform is provided", () => { + const fixture = makeFixture({ userMessage: "hello" }); + const req = makeReq({ + messages: [{ role: "user", content: "say hello to me" }], + }); + + // No transform — includes matching + expect(matchFixture([fixture], req)).toBe(fixture); + }); + + it("applies transform to inputText (embedding) matching with exact equality", () => { + const fixture = makeFixture({ inputText: "embed this text" }); + const req = makeReq({ embeddingInput: "embed this text plus extra" }); + + // Without transform — includes matches + expect(matchFixture([fixture], req)).toBe(fixture); + + // With identity transform — exact match fails + const identity = (r: ChatCompletionRequest): ChatCompletionRequest => r; + expect(matchFixture([fixture], req, undefined, identity)).toBeNull(); + + // With identity transform — exact match succeeds + const exactFixture = makeFixture({ inputText: "embed this text plus extra" }); + expect(matchFixture([exactFixture], req, undefined, identity)).toBe(exactFixture); + }); + + it("regex matching still works with transform", () => { + const fixture = makeFixture({ userMessage: /weather/i }); + const req = makeReq({ + messages: [{ role: "user", content: "tell me the weather 2026-04-02T10:30:00.000Z" }], + }); + + // Regex always uses .test(), not exact match + expect(matchFixture([fixture], req, undefined, stripTimestamps)).toBe(fixture); + }); + + it("predicate receives original (untransformed) request", () => { + let receivedContent: string | null = null; + const fixture = makeFixture({ + predicate: (r) => { + const msg = r.messages.find((m) => m.role === "user"); + receivedContent = typeof msg?.content === "string" ? msg.content : null; + return true; + }, + }); + + const originalContent = "hello 2026-04-02T10:30:00.000Z"; + const req = makeReq({ + messages: [{ role: "user", content: originalContent }], + }); + + matchFixture([fixture], req, undefined, stripTimestamps); + // Predicate should see the original request, not the transformed one + expect(receivedContent).toBe(originalContent); + }); + + it("transform applies to model matching", () => { + const fixture = makeFixture({ model: "cleaned-model" }); + const req = makeReq({ model: "original-model" }); + + const modelTransform = (r: ChatCompletionRequest): ChatCompletionRequest => ({ + ...r, + model: "cleaned-model", + }); + + expect(matchFixture([fixture], req, undefined, modelTransform)).toBe(fixture); + }); + + it("identity transform does not break tool call matching", () => { + const fixture = makeFixture({ toolName: "get_weather" }); + const req = makeReq({ + tools: [ + { + type: "function", + function: { name: "get_weather", description: "Get weather" }, + }, + ], + }); + + const identity = (r: ChatCompletionRequest): ChatCompletionRequest => r; + expect(matchFixture([fixture], req, undefined, identity)).toBe(fixture); + }); + + it("identity transform does not break toolCallId matching", () => { + const fixture = makeFixture({ toolCallId: "call_123" }); + const req = makeReq({ + messages: [ + { role: "user", content: "hi" }, + { role: "tool", content: "result", tool_call_id: "call_123" }, + ], + }); + + const identity = (r: ChatCompletionRequest): ChatCompletionRequest => r; + expect(matchFixture([fixture], req, undefined, identity)).toBe(fixture); + }); + + it("sequenceIndex still works with transform", () => { + const fixture = makeFixture({ userMessage: "cleaned", sequenceIndex: 1 }); + const req = makeReq({ + messages: [{ role: "user", content: "cleaned" }], + }); + + const identity = (r: ChatCompletionRequest): ChatCompletionRequest => r; + const counts = new Map(); + + // First call (count 0) — sequenceIndex 1 should not match + expect(matchFixture([fixture], req, counts, identity)).toBeNull(); + + // Simulate count increment + counts.set(fixture, 1); + expect(matchFixture([fixture], req, counts, identity)).toBe(fixture); + }); +}); + +// --------------------------------------------------------------------------- +// Integration tests — LLMock server with requestTransform +// --------------------------------------------------------------------------- + +let mock: LLMock | null = null; + +afterEach(async () => { + if (mock) { + await mock.stop(); + mock = null; + } +}); + +describe("LLMock server — requestTransform", () => { + it("matches fixture after transform strips timestamps from request", async () => { + mock = new LLMock({ + requestTransform: stripTimestamps, + }); + + // Fixture expects the cleaned message (no timestamp) + mock.onMessage("tell me the weather ", { content: "It will be sunny" }); + + const url = await mock.start(); + + const res = await httpPost(`${url}/v1/chat/completions`, { + model: "gpt-4", + messages: [ + { + role: "user", + content: "tell me the weather 2026-04-02T10:30:00.000Z", + }, + ], + }); + + expect(res.status).toBe(200); + const body = JSON.parse(res.body); + expect(body.choices[0].message.content).toBe("It will be sunny"); + }); + + it("uses exact equality with transform — prevents false positive substring matches", async () => { + mock = new LLMock({ + requestTransform: (req) => req, // identity + }); + + // "hello" is a substring of "hello world" — but with transform, + // exact match is used, so this should NOT match + mock.onMessage("hello", { content: "should not match" }); + mock.onMessage("hello world", { content: "correct match" }); + + const url = await mock.start(); + + const res = await httpPost(`${url}/v1/chat/completions`, { + model: "gpt-4", + messages: [{ role: "user", content: "hello world" }], + }); + + expect(res.status).toBe(200); + const body = JSON.parse(res.body); + expect(body.choices[0].message.content).toBe("correct match"); + }); + + it("works without requestTransform — backward compatible includes matching", async () => { + mock = new LLMock(); + + mock.onMessage("hello", { content: "matched via includes" }); + + const url = await mock.start(); + + const res = await httpPost(`${url}/v1/chat/completions`, { + model: "gpt-4", + messages: [{ role: "user", content: "say hello to everyone" }], + }); + + expect(res.status).toBe(200); + const body = JSON.parse(res.body); + expect(body.choices[0].message.content).toBe("matched via includes"); + }); + + it("transform works with streaming responses", async () => { + mock = new LLMock({ + requestTransform: stripTimestamps, + }); + + mock.onMessage("weather ", { content: "sunny" }); + + const url = await mock.start(); + + const res = await httpPost(`${url}/v1/chat/completions`, { + model: "gpt-4", + stream: true, + messages: [{ role: "user", content: "weather 2026-01-01T00:00:00Z" }], + }); + + expect(res.status).toBe(200); + // Streaming responses have SSE format — just verify it returned 200 + expect(res.body).toContain("sunny"); + }); + + it("transform works with embedding requests", async () => { + mock = new LLMock({ + requestTransform: (req) => ({ + ...req, + embeddingInput: req.embeddingInput?.replace(/\d{4}-\d{2}-\d{2}T[\d:.+Z-]+/g, ""), + }), + }); + + mock.onEmbedding("embed this ", { embedding: [0.1, 0.2, 0.3] }); + + const url = await mock.start(); + + const res = await httpPost(`${url}/v1/embeddings`, { + model: "text-embedding-3-small", + input: "embed this 2026-04-02T10:30:00Z", + }); + + expect(res.status).toBe(200); + const body = JSON.parse(res.body); + expect(body.data[0].embedding).toEqual([0.1, 0.2, 0.3]); + }); +}); diff --git a/src/bedrock-converse.ts b/src/bedrock-converse.ts index 933e0af..3f744dc 100644 --- a/src/bedrock-converse.ts +++ b/src/bedrock-converse.ts @@ -263,7 +263,12 @@ export async function handleConverse( const completionReq = converseToCompletionRequest(converseReq, modelId); - const fixture = matchFixture(fixtures, completionReq, journal.fixtureMatchCounts); + const fixture = matchFixture( + fixtures, + completionReq, + journal.fixtureMatchCounts, + defaults.requestTransform, + ); if (fixture) { journal.incrementFixtureMatchCount(fixture, fixtures); @@ -466,7 +471,12 @@ export async function handleConverseStream( const completionReq = converseToCompletionRequest(converseReq, modelId); - const fixture = matchFixture(fixtures, completionReq, journal.fixtureMatchCounts); + const fixture = matchFixture( + fixtures, + completionReq, + journal.fixtureMatchCounts, + defaults.requestTransform, + ); if (fixture) { journal.incrementFixtureMatchCount(fixture, fixtures); diff --git a/src/bedrock.ts b/src/bedrock.ts index d45f64e..b545a70 100644 --- a/src/bedrock.ts +++ b/src/bedrock.ts @@ -309,7 +309,12 @@ export async function handleBedrock( // Convert to ChatCompletionRequest for fixture matching const completionReq = bedrockToCompletionRequest(bedrockReq, modelId); - const fixture = matchFixture(fixtures, completionReq, journal.fixtureMatchCounts); + const fixture = matchFixture( + fixtures, + completionReq, + journal.fixtureMatchCounts, + defaults.requestTransform, + ); if (fixture) { journal.incrementFixtureMatchCount(fixture, fixtures); @@ -626,7 +631,12 @@ export async function handleBedrockStream( const completionReq = bedrockToCompletionRequest(bedrockReq, modelId); - const fixture = matchFixture(fixtures, completionReq, journal.fixtureMatchCounts); + const fixture = matchFixture( + fixtures, + completionReq, + journal.fixtureMatchCounts, + defaults.requestTransform, + ); if (fixture) { journal.incrementFixtureMatchCount(fixture, fixtures); diff --git a/src/cohere.ts b/src/cohere.ts index 5bc00fa..bdf9748 100644 --- a/src/cohere.ts +++ b/src/cohere.ts @@ -465,7 +465,12 @@ export async function handleCohere( // Convert to ChatCompletionRequest for fixture matching const completionReq = cohereToCompletionRequest(cohereReq); - const fixture = matchFixture(fixtures, completionReq, journal.fixtureMatchCounts); + const fixture = matchFixture( + fixtures, + completionReq, + journal.fixtureMatchCounts, + defaults.requestTransform, + ); if (fixture) { journal.incrementFixtureMatchCount(fixture, fixtures); diff --git a/src/embeddings.ts b/src/embeddings.ts index 95dc678..970d140 100644 --- a/src/embeddings.ts +++ b/src/embeddings.ts @@ -86,7 +86,12 @@ export async function handleEmbeddings( embeddingInput: combinedInput, }; - const fixture = matchFixture(fixtures, syntheticReq, journal.fixtureMatchCounts); + const fixture = matchFixture( + fixtures, + syntheticReq, + journal.fixtureMatchCounts, + defaults.requestTransform, + ); if (fixture) { journal.incrementFixtureMatchCount(fixture, fixtures); diff --git a/src/gemini.ts b/src/gemini.ts index 4229839..5e5493c 100644 --- a/src/gemini.ts +++ b/src/gemini.ts @@ -415,7 +415,12 @@ export async function handleGemini( // Convert to ChatCompletionRequest for fixture matching const completionReq = geminiToCompletionRequest(geminiReq, model, streaming); - const fixture = matchFixture(fixtures, completionReq, journal.fixtureMatchCounts); + const fixture = matchFixture( + fixtures, + completionReq, + journal.fixtureMatchCounts, + defaults.requestTransform, + ); const path = req.url ?? `/v1beta/models/${model}:generateContent`; if (fixture) { diff --git a/src/messages.ts b/src/messages.ts index ee17efc..7cd5547 100644 --- a/src/messages.ts +++ b/src/messages.ts @@ -500,7 +500,12 @@ export async function handleMessages( // Convert to ChatCompletionRequest for fixture matching const completionReq = claudeToCompletionRequest(claudeReq); - const fixture = matchFixture(fixtures, completionReq, journal.fixtureMatchCounts); + const fixture = matchFixture( + fixtures, + completionReq, + journal.fixtureMatchCounts, + defaults.requestTransform, + ); if (fixture) { journal.incrementFixtureMatchCount(fixture, fixtures); diff --git a/src/ollama.ts b/src/ollama.ts index 20ed12f..eba0111 100644 --- a/src/ollama.ts +++ b/src/ollama.ts @@ -342,7 +342,12 @@ export async function handleOllama( // Convert to ChatCompletionRequest for fixture matching const completionReq = ollamaToCompletionRequest(ollamaReq); - const fixture = matchFixture(fixtures, completionReq, journal.fixtureMatchCounts); + const fixture = matchFixture( + fixtures, + completionReq, + journal.fixtureMatchCounts, + defaults.requestTransform, + ); if (fixture) { journal.incrementFixtureMatchCount(fixture, fixtures); @@ -585,7 +590,12 @@ export async function handleOllamaGenerate( // Convert to ChatCompletionRequest for fixture matching const completionReq = ollamaGenerateToCompletionRequest(generateReq); - const fixture = matchFixture(fixtures, completionReq, journal.fixtureMatchCounts); + const fixture = matchFixture( + fixtures, + completionReq, + journal.fixtureMatchCounts, + defaults.requestTransform, + ); if (fixture) { journal.incrementFixtureMatchCount(fixture, fixtures); diff --git a/src/recorder.ts b/src/recorder.ts index 59ea6f8..20f5a9b 100644 --- a/src/recorder.ts +++ b/src/recorder.ts @@ -51,7 +51,11 @@ export async function proxyAndRecord( providerKey: RecordProviderKey, pathname: string, fixtures: Fixture[], - defaults: { record?: RecordConfig; logger: Logger }, + defaults: { + record?: RecordConfig; + logger: Logger; + requestTransform?: (req: ChatCompletionRequest) => ChatCompletionRequest; + }, rawBody?: string, ): Promise { const record = defaults.record; @@ -170,8 +174,9 @@ export async function proxyAndRecord( fixtureResponse = buildFixtureResponse(parsedResponse, upstreamStatus, encodingFormat); } - // Build the match criteria from the original request - const fixtureMatch = buildFixtureMatch(request); + // Build the match criteria from the (optionally transformed) request + const matchRequest = defaults.requestTransform ? defaults.requestTransform(request) : request; + const fixtureMatch = buildFixtureMatch(matchRequest); // Build and save the fixture const fixture: Fixture = { match: fixtureMatch, response: fixtureResponse }; diff --git a/src/responses.ts b/src/responses.ts index f31f5dc..d7fa30d 100644 --- a/src/responses.ts +++ b/src/responses.ts @@ -675,7 +675,12 @@ export async function handleResponses( // Convert to ChatCompletionRequest for fixture matching const completionReq = responsesToCompletionRequest(responsesReq); - const fixture = matchFixture(fixtures, completionReq, journal.fixtureMatchCounts); + const fixture = matchFixture( + fixtures, + completionReq, + journal.fixtureMatchCounts, + defaults.requestTransform, + ); if (fixture) { journal.incrementFixtureMatchCount(fixture, fixtures); diff --git a/src/router.ts b/src/router.ts index c1fdd88..efc79c1 100644 --- a/src/router.ts +++ b/src/router.ts @@ -27,22 +27,31 @@ export function matchFixture( fixtures: Fixture[], req: ChatCompletionRequest, matchCounts?: Map, + requestTransform?: (req: ChatCompletionRequest) => ChatCompletionRequest, ): Fixture | null { + // Apply transform once before matching — used for stripping dynamic data + const effective = requestTransform ? requestTransform(req) : req; + const useExactMatch = !!requestTransform; + for (const fixture of fixtures) { const { match } = fixture; - // predicate — if present, must return true + // predicate — if present, must return true (receives original request) if (match.predicate !== undefined) { if (!match.predicate(req)) continue; } // userMessage — match against the last user message content if (match.userMessage !== undefined) { - const msg = getLastMessageByRole(req.messages, "user"); + const msg = getLastMessageByRole(effective.messages, "user"); const text = msg ? getTextContent(msg.content) : null; if (!text) continue; if (typeof match.userMessage === "string") { - if (!text.includes(match.userMessage)) continue; + if (useExactMatch) { + if (text !== match.userMessage) continue; + } else { + if (!text.includes(match.userMessage)) continue; + } } else { if (!match.userMessage.test(text)) continue; } @@ -50,23 +59,27 @@ export function matchFixture( // toolCallId — match against the last tool message's tool_call_id if (match.toolCallId !== undefined) { - const msg = getLastMessageByRole(req.messages, "tool"); + const msg = getLastMessageByRole(effective.messages, "tool"); if (!msg || msg.tool_call_id !== match.toolCallId) continue; } // toolName — match against any tool definition by function.name if (match.toolName !== undefined) { - const tools = req.tools ?? []; + const tools = effective.tools ?? []; const found = tools.some((t) => t.function.name === match.toolName); if (!found) continue; } // inputText — match against the embedding input text (used by embeddings endpoint) if (match.inputText !== undefined) { - const embeddingInput = req.embeddingInput; + const embeddingInput = effective.embeddingInput; if (!embeddingInput) continue; if (typeof match.inputText === "string") { - if (!embeddingInput.includes(match.inputText)) continue; + if (useExactMatch) { + if (embeddingInput !== match.inputText) continue; + } else { + if (!embeddingInput.includes(match.inputText)) continue; + } } else { if (!match.inputText.test(embeddingInput)) continue; } @@ -74,16 +87,16 @@ export function matchFixture( // responseFormat — exact string match against request response_format.type if (match.responseFormat !== undefined) { - const reqType = req.response_format?.type; + const reqType = effective.response_format?.type; if (reqType !== match.responseFormat) continue; } // model — exact string or regexp if (match.model !== undefined) { if (typeof match.model === "string") { - if (req.model !== match.model) continue; + if (effective.model !== match.model) continue; } else { - if (!match.model.test(req.model)) continue; + if (!match.model.test(effective.model)) continue; } } diff --git a/src/server.ts b/src/server.ts index 5e5f08c..02120b0 100644 --- a/src/server.ts +++ b/src/server.ts @@ -340,7 +340,12 @@ async function handleCompletions( } // Match fixture - const fixture = matchFixture(fixtures, body, journal.fixtureMatchCounts); + const fixture = matchFixture( + fixtures, + body, + journal.fixtureMatchCounts, + defaults.requestTransform, + ); if (fixture) { journal.incrementFixtureMatchCount(fixture, fixtures); @@ -562,6 +567,9 @@ export async function createServer( get strict() { return serverOptions.strict; }, + get requestTransform() { + return serverOptions.requestTransform; + }, }; // Validate chaos config rates diff --git a/src/types.ts b/src/types.ts index 5f76d7d..50e9a52 100644 --- a/src/types.ts +++ b/src/types.ts @@ -265,6 +265,16 @@ export interface MockServerOptions { strict?: boolean; /** Record-and-replay: proxy unmatched requests to upstream and save fixtures. */ record?: RecordConfig; + /** + * Normalize requests before matching and recording. Useful for stripping + * dynamic data (timestamps, UUIDs, session IDs) that would cause fixture + * mismatches on replay. + * + * When set, string matching for `userMessage` and `inputText` uses exact + * equality (`===`) instead of substring (`includes`) to prevent false + * positives from shortened keys. + */ + requestTransform?: (req: ChatCompletionRequest) => ChatCompletionRequest; } // Handler defaults — the common shape passed from server.ts to every handler @@ -279,4 +289,5 @@ export interface HandlerDefaults { registry?: MetricsRegistry; record?: RecordConfig; strict?: boolean; + requestTransform?: (req: ChatCompletionRequest) => ChatCompletionRequest; } diff --git a/src/ws-gemini-live.ts b/src/ws-gemini-live.ts index 15f70bf..11a9c21 100644 --- a/src/ws-gemini-live.ts +++ b/src/ws-gemini-live.ts @@ -171,7 +171,14 @@ export function handleWebSocketGeminiLive( ws: WebSocketConnection, fixtures: Fixture[], journal: Journal, - defaults: { latency: number; chunkSize: number; model: string; logger: Logger; strict?: boolean }, + defaults: { + latency: number; + chunkSize: number; + model: string; + logger: Logger; + strict?: boolean; + requestTransform?: (req: ChatCompletionRequest) => ChatCompletionRequest; + }, ): void { const { logger } = defaults; const session: SessionState = { @@ -206,7 +213,14 @@ async function processMessage( ws: WebSocketConnection, fixtures: Fixture[], journal: Journal, - defaults: { latency: number; chunkSize: number; model: string; logger: Logger; strict?: boolean }, + defaults: { + latency: number; + chunkSize: number; + model: string; + logger: Logger; + strict?: boolean; + requestTransform?: (req: ChatCompletionRequest) => ChatCompletionRequest; + }, session: SessionState, ): Promise { let parsed: GeminiLiveMessage; @@ -295,7 +309,12 @@ async function processMessage( tools: session.tools.length > 0 ? session.tools : undefined, }; - const fixture = matchFixture(fixtures, completionReq, journal.fixtureMatchCounts); + const fixture = matchFixture( + fixtures, + completionReq, + journal.fixtureMatchCounts, + defaults.requestTransform, + ); const path = WS_PATH; if (fixture) { diff --git a/src/ws-realtime.ts b/src/ws-realtime.ts index 6c9955d..9deb16e 100644 --- a/src/ws-realtime.ts +++ b/src/ws-realtime.ts @@ -130,7 +130,14 @@ export function handleWebSocketRealtime( ws: WebSocketConnection, fixtures: Fixture[], journal: Journal, - defaults: { latency: number; chunkSize: number; model: string; logger: Logger; strict?: boolean }, + defaults: { + latency: number; + chunkSize: number; + model: string; + logger: Logger; + strict?: boolean; + requestTransform?: (req: ChatCompletionRequest) => ChatCompletionRequest; + }, ): void { const { logger } = defaults; const sessionId = generateId("sess"); @@ -176,7 +183,14 @@ async function processMessage( ws: WebSocketConnection, fixtures: Fixture[], journal: Journal, - defaults: { latency: number; chunkSize: number; model: string; logger: Logger; strict?: boolean }, + defaults: { + latency: number; + chunkSize: number; + model: string; + logger: Logger; + strict?: boolean; + requestTransform?: (req: ChatCompletionRequest) => ChatCompletionRequest; + }, session: SessionConfig, conversationItems: RealtimeItem[], ): Promise { @@ -246,7 +260,14 @@ async function handleResponseCreate( ws: WebSocketConnection, fixtures: Fixture[], journal: Journal, - defaults: { latency: number; chunkSize: number; model: string; logger: Logger; strict?: boolean }, + defaults: { + latency: number; + chunkSize: number; + model: string; + logger: Logger; + strict?: boolean; + requestTransform?: (req: ChatCompletionRequest) => ChatCompletionRequest; + }, session: SessionConfig, conversationItems: RealtimeItem[], ): Promise { @@ -258,7 +279,12 @@ async function handleResponseCreate( messages, }; - const fixture = matchFixture(fixtures, completionReq, journal.fixtureMatchCounts); + const fixture = matchFixture( + fixtures, + completionReq, + journal.fixtureMatchCounts, + defaults.requestTransform, + ); const responseId = generateId("resp"); if (fixture) { diff --git a/src/ws-responses.ts b/src/ws-responses.ts index 6b92c20..d8c8c01 100644 --- a/src/ws-responses.ts +++ b/src/ws-responses.ts @@ -6,7 +6,7 @@ * handler, but as individual WebSocket text frames. */ -import type { Fixture } from "./types.js"; +import type { ChatCompletionRequest, Fixture } from "./types.js"; import { matchFixture } from "./router.js"; import { responsesToCompletionRequest, @@ -57,7 +57,14 @@ export function handleWebSocketResponses( ws: WebSocketConnection, fixtures: Fixture[], journal: Journal, - defaults: { latency: number; chunkSize: number; model: string; logger: Logger; strict?: boolean }, + defaults: { + latency: number; + chunkSize: number; + model: string; + logger: Logger; + strict?: boolean; + requestTransform?: (req: ChatCompletionRequest) => ChatCompletionRequest; + }, ): void { const { logger } = defaults; // Serialize message processing to prevent event interleaving @@ -82,7 +89,14 @@ async function processMessage( ws: WebSocketConnection, fixtures: Fixture[], journal: Journal, - defaults: { latency: number; chunkSize: number; model: string; logger: Logger; strict?: boolean }, + defaults: { + latency: number; + chunkSize: number; + model: string; + logger: Logger; + strict?: boolean; + requestTransform?: (req: ChatCompletionRequest) => ChatCompletionRequest; + }, ): Promise { let parsed: unknown; try { @@ -136,7 +150,12 @@ async function processMessage( }; const completionReq = responsesToCompletionRequest(responsesReq); - const fixture = matchFixture(fixtures, completionReq, journal.fixtureMatchCounts); + const fixture = matchFixture( + fixtures, + completionReq, + journal.fixtureMatchCounts, + defaults.requestTransform, + ); if (fixture) { journal.incrementFixtureMatchCount(fixture, fixtures);