can1357
2f151fea9a
- Added `close()` method to SessionManager and AuthStorage for proper resource cleanup and finalization of prepared statements. - Added `initiatorOverride` option support in OpenAI and Anthropic providers for message attribution control. - Fixed resource leaks in RpcClient timeout handling by centralizing timeout creation with unref() and adding explicit clearTimeout() calls. - Fixed AgentSession disposal to call SessionManager's `close()` method for guaranteed resource cleanup instead of fallback flush. - Updated all test suites to properly dispose AuthStorage instances in cleanup hooks to prevent resource leaks between tests.
92 lines
2.7 KiB
TypeScript
92 lines
2.7 KiB
TypeScript
import { afterEach, beforeEach, describe, expect, test, vi } from "bun:test";
|
|
import * as fs from "node:fs";
|
|
import * as os from "node:os";
|
|
import * as path from "node:path";
|
|
import type { OAuthCredentials, UsageProvider } from "@oh-my-pi/pi-ai";
|
|
import * as ai from "@oh-my-pi/pi-ai";
|
|
import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage";
|
|
import { Snowflake } from "@oh-my-pi/pi-utils";
|
|
|
|
describe("AuthStorage account rotation", () => {
|
|
let tempDir: string;
|
|
let authStorage: AuthStorage;
|
|
let usageExhausted = false;
|
|
|
|
const usageProvider: UsageProvider = {
|
|
id: "openai-codex",
|
|
async fetchUsage(params) {
|
|
const accountId = params.credential.accountId ?? "unknown";
|
|
return {
|
|
provider: "openai-codex",
|
|
fetchedAt: Date.now(),
|
|
limits: [
|
|
{
|
|
id: `requests-${accountId}`,
|
|
label: "Requests",
|
|
scope: { provider: "openai-codex", accountId },
|
|
amount: { unit: "requests", used: usageExhausted ? 100 : 10, limit: 100 },
|
|
status: usageExhausted ? "exhausted" : "ok",
|
|
},
|
|
],
|
|
};
|
|
},
|
|
};
|
|
|
|
beforeEach(async () => {
|
|
tempDir = path.join(os.tmpdir(), `pi-test-auth-rotation-${Snowflake.next()}`);
|
|
fs.mkdirSync(tempDir, { recursive: true });
|
|
usageExhausted = false;
|
|
|
|
authStorage = await AuthStorage.create(path.join(tempDir, "testauth.db"), {
|
|
usageProviderResolver: provider => (provider === "openai-codex" ? usageProvider : undefined),
|
|
});
|
|
|
|
vi.spyOn(ai, "getOAuthApiKey").mockImplementation(async (_provider, credentials) => {
|
|
const credential = credentials["openai-codex"] as OAuthCredentials | undefined;
|
|
if (!credential) return null;
|
|
return {
|
|
apiKey: `api-${credential.accountId ?? "unknown"}`,
|
|
newCredentials: credential,
|
|
};
|
|
});
|
|
});
|
|
|
|
afterEach(() => {
|
|
vi.restoreAllMocks();
|
|
authStorage.close();
|
|
if (tempDir && fs.existsSync(tempDir)) {
|
|
fs.rmSync(tempDir, { recursive: true });
|
|
}
|
|
});
|
|
|
|
test("returns a fallback key when every OAuth account is usage-limited", async () => {
|
|
await authStorage.set("openai-codex", [
|
|
{
|
|
type: "oauth",
|
|
access: "access-1",
|
|
refresh: "refresh-1",
|
|
expires: Date.now() + 60_000,
|
|
accountId: "acct-1",
|
|
},
|
|
{
|
|
type: "oauth",
|
|
access: "access-2",
|
|
refresh: "refresh-2",
|
|
expires: Date.now() + 60_000,
|
|
accountId: "acct-2",
|
|
},
|
|
]);
|
|
|
|
const sessionId = "issue-55-session";
|
|
const firstKey = await authStorage.getApiKey("openai-codex", sessionId);
|
|
expect(firstKey).toMatch(/^api-acct-/);
|
|
|
|
usageExhausted = true;
|
|
const switched = await authStorage.markUsageLimitReached("openai-codex", sessionId);
|
|
expect(switched).toBe(true);
|
|
|
|
const exhaustedFallbackKey = await authStorage.getApiKey("openai-codex", sessionId);
|
|
expect(exhaustedFallbackKey).toMatch(/^api-acct-/);
|
|
});
|
|
});
|