test: share sanitize session usage helpers

This commit is contained in:
Peter Steinberger 2026-03-13 23:59:46 +00:00
parent 91b9c47dad
commit 903cb0679d
1 changed files with 62 additions and 74 deletions

View File

@ -177,6 +177,14 @@ describe("sanitizeSessionHistory", () => {
AgentMessage & { usage?: unknown; content?: unknown }
>;
const getSingleAssistantUsage = async (messages: AgentMessage[]) => {
vi.mocked(mockedHelpers.isGoogleModelApi).mockReturnValue(false);
const result = await sanitizeOpenAIHistory(messages);
return result.find((message) => message.role === "assistant") as
| (AgentMessage & { usage?: unknown })
| undefined;
};
beforeEach(async () => {
testTimestamp = 1;
const harness = await loadSanitizeSessionHistoryWithCleanMocks();
@ -358,43 +366,33 @@ describe("sanitizeSessionHistory", () => {
});
it("adds a zeroed assistant usage snapshot when usage is missing", async () => {
vi.mocked(mockedHelpers.isGoogleModelApi).mockReturnValue(false);
const messages = castAgentMessages([
{ role: "user", content: "question" },
{
role: "assistant",
content: [{ type: "text", text: "answer without usage" }],
},
]);
const result = await sanitizeOpenAIHistory(messages);
const assistant = result.find((message) => message.role === "assistant") as
| (AgentMessage & { usage?: unknown })
| undefined;
const assistant = await getSingleAssistantUsage(
castAgentMessages([
{ role: "user", content: "question" },
{
role: "assistant",
content: [{ type: "text", text: "answer without usage" }],
},
]),
);
expect(assistant?.usage).toEqual(makeZeroUsageSnapshot());
});
it("normalizes mixed partial assistant usage fields to numeric totals", async () => {
vi.mocked(mockedHelpers.isGoogleModelApi).mockReturnValue(false);
const messages = castAgentMessages([
{ role: "user", content: "question" },
{
role: "assistant",
content: [{ type: "text", text: "answer with partial usage" }],
usage: {
output: 3,
cache_read_input_tokens: 9,
const assistant = await getSingleAssistantUsage(
castAgentMessages([
{ role: "user", content: "question" },
{
role: "assistant",
content: [{ type: "text", text: "answer with partial usage" }],
usage: {
output: 3,
cache_read_input_tokens: 9,
},
},
},
]);
const result = await sanitizeOpenAIHistory(messages);
const assistant = result.find((message) => message.role === "assistant") as
| (AgentMessage & { usage?: unknown })
| undefined;
]),
);
expect(assistant?.usage).toEqual({
input: 0,
@ -406,31 +404,26 @@ describe("sanitizeSessionHistory", () => {
});
it("preserves existing usage cost while normalizing token fields", async () => {
vi.mocked(mockedHelpers.isGoogleModelApi).mockReturnValue(false);
const messages = castAgentMessages([
{ role: "user", content: "question" },
{
role: "assistant",
content: [{ type: "text", text: "answer with partial usage and cost" }],
usage: {
output: 3,
cache_read_input_tokens: 9,
cost: {
input: 1.25,
output: 2.5,
cacheRead: 0.25,
cacheWrite: 0,
total: 4,
const assistant = await getSingleAssistantUsage(
castAgentMessages([
{ role: "user", content: "question" },
{
role: "assistant",
content: [{ type: "text", text: "answer with partial usage and cost" }],
usage: {
output: 3,
cache_read_input_tokens: 9,
cost: {
input: 1.25,
output: 2.5,
cacheRead: 0.25,
cacheWrite: 0,
total: 4,
},
},
},
},
]);
const result = await sanitizeOpenAIHistory(messages);
const assistant = result.find((message) => message.role === "assistant") as
| (AgentMessage & { usage?: unknown })
| undefined;
]),
);
expect(assistant?.usage).toEqual({
...makeZeroUsageSnapshot(),
@ -450,27 +443,22 @@ describe("sanitizeSessionHistory", () => {
});
it("preserves unknown cost when token fields already match", async () => {
vi.mocked(mockedHelpers.isGoogleModelApi).mockReturnValue(false);
const messages = castAgentMessages([
{ role: "user", content: "question" },
{
role: "assistant",
content: [{ type: "text", text: "answer with complete numeric usage but no cost" }],
usage: {
input: 1,
output: 2,
cacheRead: 3,
cacheWrite: 4,
totalTokens: 10,
const assistant = await getSingleAssistantUsage(
castAgentMessages([
{ role: "user", content: "question" },
{
role: "assistant",
content: [{ type: "text", text: "answer with complete numeric usage but no cost" }],
usage: {
input: 1,
output: 2,
cacheRead: 3,
cacheWrite: 4,
totalTokens: 10,
},
},
},
]);
const result = await sanitizeOpenAIHistory(messages);
const assistant = result.find((message) => message.role === "assistant") as
| (AgentMessage & { usage?: unknown })
| undefined;
]),
);
expect(assistant?.usage).toEqual({
input: 1,