Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
4.6 kB · 115 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116import { existsSync, mkdirSync, rmSync } from "node:fs";import { tmpdir } from "node:os";import { join } from "node:path";import { getModel, usageProviders } from "@earendil-works/pi-ai";import { afterEach, beforeEach, describe, expect, it } from "vitest";import { AuthStorage } from "../src/core/auth-storage.js";import { DefaultResourceLoader } from "../src/core/resource-loader.js";import type { ExtensionFactory } from "../src/core/sdk.js";import { createAgentSession } from "../src/core/sdk.js";import { SessionManager } from "../src/core/session-manager.js";import { SettingsManager } from "../src/core/settings-manager.js";import { fetchUsageReports } from "../src/modes/interactive/usage-report.js";
/** * A provider registered by an extension can report its own quota. Without this seam the feed * exists but /usage never calls it, so the provider looks like it has no limits at all. */describe("extension provider usage feeds", () => { let tempDir: string; let agentDir: string;
beforeEach(() => { tempDir = join(tmpdir(), `pi-usage-feed-test-${Date.now()}-${Math.random().toString(36).slice(2)}`); agentDir = join(tempDir, "agent"); mkdirSync(agentDir, { recursive: true }); });
afterEach(() => { if (tempDir && existsSync(tempDir)) rmSync(tempDir, { recursive: true, force: true }); });
async function createSession(extensionFactories: ExtensionFactory[], apiKey?: { provider: string; key: string }) { const settingsManager = SettingsManager.create(tempDir, agentDir); const authStorage = apiKey ? AuthStorage.inMemory({ [apiKey.provider]: [{ type: "api_key", key: apiKey.key }] }) : AuthStorage.inMemory(); const resourceLoader = new DefaultResourceLoader({ cwd: tempDir, agentDir, settingsManager, extensionFactories }); await resourceLoader.reload(); const { session } = await createAgentSession({ cwd: tempDir, agentDir, model: getModel("anthropic", "claude-sonnet-4-5")!, settingsManager, sessionManager: SessionManager.inMemory(), authStorage, resourceLoader, }); return session; }
const extensionWithFeed: ExtensionFactory = (pi) => { pi.registerProvider("hyper", { name: "Charm Hyper", apiKey: "unused", usage: { credentialKinds: ["apiKey"], advance: (state, responses) => { if (responses.length === 0) { return { requests: [{ kind: "http", method: "GET", url: "https://hyper.charm.land/v1/teams" }] }; } return { report: { limits: [{ label: "Credits", usedPercent: 5.2, windowLabel: "24h window", status: "ok" }], balance: "$11.83", planType: "shred", email: state.credential.email, }, }; }, }, }); };
it("registers the feed and reports it through /usage's fetch path", async () => { const session = await createSession([extensionWithFeed], { provider: "hyper", key: "hyper-key" }); const feed = usageProviders().find((candidate) => candidate.provider === "hyper"); expect(feed?.providerName).toBe("Charm Hyper");
const result = await fetchUsageReports(session.modelRegistry.authStorage, { fetchImpl: async () => new Response(JSON.stringify({ teams: [{ id: "team-1" }] }), { status: 200 }), }); const report = result.reports.find((candidate) => candidate.provider === "hyper"); expect(report?.limits[0]).toMatchObject({ label: "Credits", usedPercent: 5.2, windowLabel: "24h window" }); expect(report?.balance).toBe("$11.83"); expect(report?.planType).toBe("shred");
session.dispose(); });
it("leaves a provider with no feed skipped, with the reason", async () => { const session = await createSession([(pi) => pi.registerProvider("plain-provider", { apiKey: "unused" })], { provider: "plain-provider", key: "plain-key", }); const result = await fetchUsageReports(session.modelRegistry.authStorage, { fetchImpl: async () => new Response("{}"), }); expect(result.skipped).toContainEqual({ provider: "plain-provider", reason: "no usage feed" }); session.dispose(); });
it("drops the feed when the provider is unregistered", async () => { const session = await createSession([]); session.modelRegistry.registerProvider("hyper-two", { name: "Hyper Two", apiKey: "unused", usage: { advance: () => ({ report: { limits: [] } }) }, }); expect(usageProviders().some((candidate) => candidate.provider === "hyper-two")).toBe(true);
session.modelRegistry.unregisterProvider("hyper-two"); expect(usageProviders().some((candidate) => candidate.provider === "hyper-two")).toBe(false); session.dispose(); });});