wehub-resource-sync
95 lines
2.6 KiB
TypeScript
95 lines
2.6 KiB
TypeScript
import { afterAll, afterEach, beforeAll, describe, expect, it } from "bun:test";
|
|
import * as path from "node:path";
|
|
import { Agent } from "@oh-my-pi/pi-agent-core";
|
|
import type { AssistantMessage, UserMessage } from "@oh-my-pi/pi-ai";
|
|
import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry";
|
|
import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings";
|
|
import { AgentSession } from "@oh-my-pi/pi-coding-agent/session/agent-session";
|
|
import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage";
|
|
import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager";
|
|
import { TempDir } from "@oh-my-pi/pi-utils";
|
|
|
|
describe("AgentSession session stats", () => {
|
|
let tempDir: TempDir;
|
|
let authStorage: AuthStorage;
|
|
let modelRegistry: ModelRegistry;
|
|
let session: AgentSession | undefined;
|
|
|
|
beforeAll(async () => {
|
|
tempDir = TempDir.createSync("@pi-session-stats-");
|
|
authStorage = await AuthStorage.create(path.join(tempDir.path(), "auth.db"));
|
|
modelRegistry = new ModelRegistry(authStorage);
|
|
});
|
|
|
|
afterAll(() => {
|
|
authStorage.close();
|
|
tempDir.removeSync();
|
|
});
|
|
|
|
afterEach(async () => {
|
|
await session?.dispose();
|
|
session = undefined;
|
|
});
|
|
|
|
it("includes context usage available from the active session", () => {
|
|
const model = modelRegistry.getAll().find(candidate => candidate.contextWindow && candidate.contextWindow > 0);
|
|
if (!model?.contextWindow) {
|
|
throw new Error("Expected bundled model with a context window");
|
|
}
|
|
|
|
const userMessage: UserMessage = {
|
|
role: "user",
|
|
content: "Hello",
|
|
timestamp: Date.now(),
|
|
};
|
|
const assistantMessage: AssistantMessage = {
|
|
role: "assistant",
|
|
content: [{ type: "text", text: "ok" }],
|
|
api: model.api,
|
|
provider: model.provider,
|
|
model: model.id,
|
|
usage: {
|
|
input: 10,
|
|
output: 2,
|
|
cacheRead: 0,
|
|
cacheWrite: 0,
|
|
totalTokens: 12,
|
|
cost: {
|
|
input: 0,
|
|
output: 0,
|
|
cacheRead: 0,
|
|
cacheWrite: 0,
|
|
total: 0,
|
|
},
|
|
},
|
|
stopReason: "stop",
|
|
timestamp: Date.now(),
|
|
};
|
|
|
|
const agent = new Agent({
|
|
initialState: {
|
|
model,
|
|
systemPrompt: ["Test"],
|
|
tools: [],
|
|
messages: [userMessage, assistantMessage],
|
|
},
|
|
});
|
|
session = new AgentSession({
|
|
agent,
|
|
sessionManager: SessionManager.inMemory(),
|
|
settings: Settings.isolated({ "compaction.enabled": false }),
|
|
modelRegistry,
|
|
});
|
|
|
|
const directUsage = session.getContextUsage();
|
|
const stats = session.getSessionStats();
|
|
|
|
expect(directUsage).toEqual({
|
|
tokens: 10,
|
|
contextWindow: model.contextWindow,
|
|
percent: (10 / model.contextWindow) * 100,
|
|
});
|
|
expect(stats.contextUsage).toEqual(directUsage);
|
|
});
|
|
});
|