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:
Bruno Guidolim 2026-01-28 18:45:26 +01:00
parent 715728c989
commit 33f9bcc3ce
3 changed files with 110 additions and 1 deletions

View File

@ -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 });
}
});
});

View File

@ -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);
}

View File

@ -646,6 +646,7 @@ export async function runEmbeddedPiAgent(
await markAuthProfileUsed({
store: authStore,
profileId: lastProfileId,
model: modelId,
agentDir: params.agentDir,
});
}