fix(auth): clear per-model cooldown on success
When a model succeeds, also clear its per-model cooldown key so the system doesn't think it's still rate-limited. - Add optional `model` param to markAuthProfileUsed - Pass modelId when marking profile used in agent runner - Add tests for per-model cooldown clearing behavior
This commit is contained in:
parent
715728c989
commit
33f9bcc3ce
@ -1,8 +1,13 @@
|
||||
import fs from "node:fs/promises";
|
||||
import os from "node:os";
|
||||
import path from "node:path";
|
||||
import { describe, expect, it } from "vitest";
|
||||
import {
|
||||
calculateAuthProfileCooldownMs,
|
||||
cooldownKey,
|
||||
isProfileInCooldown,
|
||||
markAuthProfileUsed,
|
||||
saveAuthProfileStore,
|
||||
} from "./auth-profiles.js";
|
||||
import type { AuthProfileStore } from "./auth-profiles.js";
|
||||
import { AUTH_STORE_VERSION } from "./auth-profiles/constants.js";
|
||||
@ -105,3 +110,70 @@ describe("isProfileInCooldown with per-model support", () => {
|
||||
expect(isProfileInCooldown(store, "openai:default", "gpt-4")).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe("markAuthProfileUsed with per-model support", () => {
|
||||
it("clears per-model cooldown when model is provided", async () => {
|
||||
const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "moltbot-auth-"));
|
||||
const cooldownTime = Date.now() + 60_000;
|
||||
const store: AuthProfileStore = {
|
||||
version: AUTH_STORE_VERSION,
|
||||
profiles: {
|
||||
"openai:default": { type: "api_key", provider: "openai", key: "test" },
|
||||
},
|
||||
usageStats: {
|
||||
"openai:default": { cooldownUntil: cooldownTime },
|
||||
"openai:default:gpt-4": { cooldownUntil: cooldownTime, errorCount: 3 },
|
||||
"openai:default:gpt-3.5": { cooldownUntil: cooldownTime },
|
||||
},
|
||||
};
|
||||
saveAuthProfileStore(store, tempDir);
|
||||
|
||||
try {
|
||||
// Mark gpt-4 as used (successful)
|
||||
await markAuthProfileUsed({
|
||||
store,
|
||||
profileId: "openai:default",
|
||||
model: "gpt-4",
|
||||
agentDir: tempDir,
|
||||
});
|
||||
|
||||
// Profile-level cooldown should be cleared
|
||||
expect(store.usageStats?.["openai:default"]?.cooldownUntil).toBeUndefined();
|
||||
// Per-model cooldown for gpt-4 should be cleared
|
||||
expect(store.usageStats?.["openai:default:gpt-4"]?.cooldownUntil).toBeUndefined();
|
||||
expect(store.usageStats?.["openai:default:gpt-4"]?.errorCount).toBe(0);
|
||||
// Per-model cooldown for gpt-3.5 should remain (different model)
|
||||
expect(store.usageStats?.["openai:default:gpt-3.5"]?.cooldownUntil).toBe(cooldownTime);
|
||||
} finally {
|
||||
await fs.rm(tempDir, { recursive: true, force: true });
|
||||
}
|
||||
});
|
||||
|
||||
it("only clears profile-level cooldown when model is not provided", async () => {
|
||||
const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "moltbot-auth-"));
|
||||
const cooldownTime = Date.now() + 60_000;
|
||||
const store: AuthProfileStore = {
|
||||
version: AUTH_STORE_VERSION,
|
||||
profiles: {
|
||||
"openai:default": { type: "api_key", provider: "openai", key: "test" },
|
||||
},
|
||||
usageStats: {
|
||||
"openai:default": { cooldownUntil: cooldownTime },
|
||||
"openai:default:gpt-4": { cooldownUntil: cooldownTime },
|
||||
},
|
||||
};
|
||||
saveAuthProfileStore(store, tempDir);
|
||||
|
||||
try {
|
||||
// Mark profile as used without specifying model
|
||||
await markAuthProfileUsed({ store, profileId: "openai:default", agentDir: tempDir });
|
||||
|
||||
// Profile-level cooldown should be cleared
|
||||
expect(store.usageStats?.["openai:default"]?.cooldownUntil).toBeUndefined();
|
||||
// Per-model cooldown should remain (no model specified)
|
||||
expect(store.usageStats?.["openai:default:gpt-4"]?.cooldownUntil).toBe(cooldownTime);
|
||||
} finally {
|
||||
await fs.rm(tempDir, { recursive: true, force: true });
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
@ -62,17 +62,24 @@ export function isProfileInCooldown(
|
||||
* Mark a profile as successfully used. Resets error count and updates lastUsed.
|
||||
* Uses store lock to avoid overwriting concurrent usage updates.
|
||||
*/
|
||||
/**
|
||||
* Mark a profile as successfully used. Resets error count and updates lastUsed.
|
||||
* Uses store lock to avoid overwriting concurrent usage updates.
|
||||
* When model is provided, also clears the per-model cooldown.
|
||||
*/
|
||||
export async function markAuthProfileUsed(params: {
|
||||
store: AuthProfileStore;
|
||||
profileId: string;
|
||||
model?: string;
|
||||
agentDir?: string;
|
||||
}): Promise<void> {
|
||||
const { store, profileId, agentDir } = params;
|
||||
const { store, profileId, model, agentDir } = params;
|
||||
const updated = await updateAuthProfileStoreWithLock({
|
||||
agentDir,
|
||||
updater: (freshStore) => {
|
||||
if (!freshStore.profiles[profileId]) return false;
|
||||
freshStore.usageStats = freshStore.usageStats ?? {};
|
||||
// Clear profile-level cooldown
|
||||
freshStore.usageStats[profileId] = {
|
||||
...freshStore.usageStats[profileId],
|
||||
lastUsed: Date.now(),
|
||||
@ -82,6 +89,20 @@ export async function markAuthProfileUsed(params: {
|
||||
disabledReason: undefined,
|
||||
failureCounts: undefined,
|
||||
};
|
||||
// Also clear per-model cooldown if model provided
|
||||
if (model) {
|
||||
const modelKey = cooldownKey(profileId, model);
|
||||
if (freshStore.usageStats[modelKey]) {
|
||||
freshStore.usageStats[modelKey] = {
|
||||
...freshStore.usageStats[modelKey],
|
||||
errorCount: 0,
|
||||
cooldownUntil: undefined,
|
||||
disabledUntil: undefined,
|
||||
disabledReason: undefined,
|
||||
failureCounts: undefined,
|
||||
};
|
||||
}
|
||||
}
|
||||
return true;
|
||||
},
|
||||
});
|
||||
@ -92,6 +113,7 @@ export async function markAuthProfileUsed(params: {
|
||||
if (!store.profiles[profileId]) return;
|
||||
|
||||
store.usageStats = store.usageStats ?? {};
|
||||
// Clear profile-level cooldown
|
||||
store.usageStats[profileId] = {
|
||||
...store.usageStats[profileId],
|
||||
lastUsed: Date.now(),
|
||||
@ -101,6 +123,20 @@ export async function markAuthProfileUsed(params: {
|
||||
disabledReason: undefined,
|
||||
failureCounts: undefined,
|
||||
};
|
||||
// Also clear per-model cooldown if model provided
|
||||
if (model) {
|
||||
const modelKey = cooldownKey(profileId, model);
|
||||
if (store.usageStats[modelKey]) {
|
||||
store.usageStats[modelKey] = {
|
||||
...store.usageStats[modelKey],
|
||||
errorCount: 0,
|
||||
cooldownUntil: undefined,
|
||||
disabledUntil: undefined,
|
||||
disabledReason: undefined,
|
||||
failureCounts: undefined,
|
||||
};
|
||||
}
|
||||
}
|
||||
saveAuthProfileStore(store, agentDir);
|
||||
}
|
||||
|
||||
|
||||
@ -646,6 +646,7 @@ export async function runEmbeddedPiAgent(
|
||||
await markAuthProfileUsed({
|
||||
store: authStore,
|
||||
profileId: lastProfileId,
|
||||
model: modelId,
|
||||
agentDir: params.agentDir,
|
||||
});
|
||||
}
|
||||
|
||||
Loading…
Reference in New Issue
Block a user