From 9227026af362b5a66dd0f8063f91db8829ad2396 Mon Sep 17 00:00:00 2001
From: YiWang24 <yiwang2457@gmail.com>
Date: Mon, 26 Jan 2026 03:53:09 -0500
Subject: [PATCH] fix(agents): correct imports and API usage for cooldown check

---
 src/agents/model-fallback.ts | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/src/agents/model-fallback.ts b/src/agents/model-fallback.ts
index aca834b89..60827ea00 100644
--- a/src/agents/model-fallback.ts
+++ b/src/agents/model-fallback.ts
@@ -14,9 +14,9 @@ import {
   resolveModelRefFromString,
 } from "./model-selection.js";
 import type { FailoverReason } from "./pi-embedded-helpers.js";
-import { loadAuthProfileStore, isProfileInCooldown } from "./auth-profiles/usage.js";
+import { isProfileInCooldown } from "./auth-profiles/usage.js";
+import { loadAuthProfileStore } from "./auth-profiles/store.js";
 import { resolveAuthProfileOrder } from "./auth-profiles/order.js";
-import { getAgentAuthStoreDir } from "./agent-dirs.js";
 
 type ModelCandidate = {
   provider: string;
@@ -215,9 +215,7 @@ export async function runWithModelFallback<T>(params: {
     fallbacksOverride: params.fallbacksOverride,
   });
 
-  const authStore = params.cfg
-    ? loadAuthProfileStore(getAgentAuthStoreDir(params.cfg, "main"))
-    : null;
+  const authStore = params.cfg ? loadAuthProfileStore() : null;
 
   const attempts: FallbackAttempt[] = [];
   let lastError: unknown;
@@ -227,7 +225,11 @@ export async function runWithModelFallback<T>(params: {
 
     // Skip candidates that are in cooldown
     if (authStore) {
-      const profileIds = resolveAuthProfileOrder(authStore, candidate.provider);
+      const profileIds = resolveAuthProfileOrder({
+        cfg: params.cfg,
+        store: authStore,
+        provider: candidate.provider,
+      });
       const isAnyProfileAvailable = profileIds.some((id) => !isProfileInCooldown(authStore, id));
 
       if (profileIds.length > 0 && !isAnyProfileAvailable) {