diff --git a/src/agents/auth-profiles.auth-profile-cooldowns.test.ts b/src/agents/auth-profiles.auth-profile-cooldowns.test.ts index 8eab8d3d7..7ed3df38b 100644 --- a/src/agents/auth-profiles.auth-profile-cooldowns.test.ts +++ b/src/agents/auth-profiles.auth-profile-cooldowns.test.ts @@ -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 }); + } + }); +}); diff --git a/src/agents/auth-profiles/usage.ts b/src/agents/auth-profiles/usage.ts index 43bb660b6..1681cdad2 100644 --- a/src/agents/auth-profiles/usage.ts +++ b/src/agents/auth-profiles/usage.ts @@ -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 { - 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); } diff --git a/src/agents/pi-embedded-runner/run.ts b/src/agents/pi-embedded-runner/run.ts index 746398217..b6b002319 100644 --- a/src/agents/pi-embedded-runner/run.ts +++ b/src/agents/pi-embedded-runner/run.ts @@ -646,6 +646,7 @@ export async function runEmbeddedPiAgent( await markAuthProfileUsed({ store: authStore, profileId: lastProfileId, + model: modelId, agentDir: params.agentDir, }); }