57 lines
1.6 KiB
TypeScript
57 lines
1.6 KiB
TypeScript
import { fauxAssistantMessage } from "@earendil-works/pi-ai";
|
|
import { afterEach, describe, expect, it } from "vitest";
|
|
import { createHarness, type Harness } from "../harness.ts";
|
|
|
|
describe("regression #3982: message_end cost override", () => {
|
|
const harnesses: Harness[] = [];
|
|
|
|
afterEach(() => {
|
|
while (harnesses.length > 0) {
|
|
harnesses.pop()?.cleanup();
|
|
}
|
|
});
|
|
|
|
it("allows extensions to replace finalized assistant usage cost", async () => {
|
|
const harness = await createHarness({
|
|
extensionFactories: [
|
|
(pi) => {
|
|
pi.on("message_end", (event) => {
|
|
if (event.message.role !== "assistant") return;
|
|
|
|
return {
|
|
message: {
|
|
...event.message,
|
|
usage: {
|
|
...event.message.usage,
|
|
cost: {
|
|
...event.message.usage.cost,
|
|
total: 0.123,
|
|
},
|
|
},
|
|
},
|
|
};
|
|
});
|
|
},
|
|
],
|
|
});
|
|
harnesses.push(harness);
|
|
harness.setResponses([fauxAssistantMessage("hello")]);
|
|
|
|
await harness.session.prompt("hi");
|
|
|
|
const assistantMessage = harness.session.messages.find((message) => message.role === "assistant");
|
|
expect(assistantMessage?.role).toBe("assistant");
|
|
if (assistantMessage?.role !== "assistant") {
|
|
throw new Error("missing assistant message");
|
|
}
|
|
expect(assistantMessage.usage.cost.total).toBe(0.123);
|
|
|
|
const messageEnd = harness.eventsOfType("message_end").find((event) => event.message.role === "assistant");
|
|
expect(messageEnd?.message.role).toBe("assistant");
|
|
if (messageEnd?.message.role !== "assistant") {
|
|
throw new Error("missing assistant message_end event");
|
|
}
|
|
expect(messageEnd.message.usage.cost.total).toBe(0.123);
|
|
});
|
|
});
|