8000 fix(migrate): skip imported answers without usage in the usage backfill by geminixiang · Pull Request #165 · geminixiang/mikan · GitHub
[go: up one dir, main page]

Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,13 @@ change in a minor release.

## [Unreleased]

### Fixed

- `mikan migrate` step `0013-session-usage` no longer stops with
`Cannot read properties of undefined (reading 'input')` on an imported answer that recorded no
usage, and the Admin portal's daily usage chart no longer drops a session holding one; such
answers count as no spend.

## [1.3.0-beta.0]

### Added
Expand Down
10 changes: 7 additions & 3 deletions src/sessions/session-store.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ import { dirname } from "node:path";
import type { AgentMessage } from "@earendil-works/pi-agent-core";
import type { JsonValue } from "@earendil-works/chord";
import { BACKGROUND_CONTEXT } from "@earendil-works/chord/context";
import type { Message, Models, Usage } from "@earendil-works/pi-ai";
import type { AssistantMessage, Message, Models, Usage } from "@earendil-works/pi-ai";
import { calculateContextTokens } from "@earendil-works/pi-ai/utils/estimate";
import {
AssistantEntry,
Expand Down Expand Up @@ -687,7 +687,7 @@ async function importedUsage(
record.conversationId !== conversationId ||
record.byTaskId !== undefined ||
!AssistantEntry.is(record) ||
message?.role !== "assistant"
!recordsUsage(message)
) {
continue;
}
Expand All @@ -701,6 +701,10 @@ async function importedUsage(
return byModel;
}

function recordsUsage(message: AgentMessage | undefined): message is AssistantMessage {
return message?.role === "assistant" && isRecord(message.usage);
}

async function recordImportedUsage(
harness: Harness,
conversationId: ConversationId,
Expand Down Expand Up @@ -1035,7 +1039,7 @@ export class SessionStore implements SessionInspection {
const id = session.conversation.id;
return (await readRecords(session.conversation)).flatMap((record) => {
const message = record.model?.[0];
return record.conversationId === id && message?.role === "assistant"
return record.conversationId === id && recordsUsage(message)
? [{ timestamp: message.timestamp, usage: message.usage }]
: [];
});
Expand Down
24 changes: 21 additions & 3 deletions src/test/migrate-session-usage.test.ts
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import type { AssistantMessage } from "@earendil-works/pi-ai";
import { mkdtempSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
Expand Down Expand Up @@ -42,10 +43,15 @@ function context(dryRun = false): MigrationContext {
};
}

async function importedAnswer(key: string, input: number, cause?: string): Promise<void> {
async function importedAnswer(
key: string,
input: number,
cause?: string,
withUsage = true,
): Promise<void> {
const session = await SessionStore.open(office, key);
if (cause) await session.appendCustomEntry(RUN_CAUSE_CUSTOM_TYPE, { messageId: cause });
await session.appendMessage({
const answer: AssistantMessage = {
role: "assistant",
content: [{ type: "text", text: "imported answer" }],
api: "openai-responses",
Expand All @@ -61,7 +67,9 @@ async function importedAnswer(key: string, input: number, cause?: string): Promi
},
stopReason: "stop",
timestamp: 1,
});
};
if (!withUsage) Reflect.deleteProperty(answer, "usage");
await session.appendMessage(answer);
await session.close();
}

Expand All @@ -86,6 +94,16 @@ test("records the spend of imported answers once, and only in the session that h
expect(await spentTokens("C1:1000.1")).toBe(0);
});

test("skips imported answers that recorded no usage", async () => {
await importedAnswer("C1", 10, undefined, false);
await importedAnswer("C1", 10);

await sessionUsageMigration.run(context());

expect(await spentTokens("C1")).toBe(12);
expect(await SessionStore.ownResponseUsage(office, "C1")).toHaveLength(1);
});

test("does nothing for an office without session storage", async () => {
await sessionUsageMigration.run(context());
expect(reports).toEqual([]);
Expand Down
Loading
0