diff --git a/packages/daemon/src/lib/agent/agent-session.ts b/packages/daemon/src/lib/agent/agent-session.ts index 501424dab..130a17551 100644 --- a/packages/daemon/src/lib/agent/agent-session.ts +++ b/packages/daemon/src/lib/agent/agent-session.ts @@ -83,6 +83,7 @@ import type { SystemPromptConfig, McpServerConfig, Provider, + FallbackModelEntry, } from '@hyperneo/shared'; import type { ChatMessage, @@ -247,6 +248,8 @@ import { MessageRecoveryHandler } from './message-recovery-handler'; import { RewindHandler, type RewindHandlerContext, type RewindPoint } from './rewind-handler'; import { SessionConfigHandler, type SessionConfigHandlerContext } from './session-config-handler'; import { RateLimitWatchdog } from './rate-limit-watchdog'; +import { resolveFallbackChain } from './fallback-recovery'; +import { getProviderRegistry } from '../providers/factory.js'; /** * AgentSession - Pure facade that delegates to specialized handlers @@ -427,9 +430,65 @@ export class AgentSession // Initialize SessionConfigHandler (handlers take AgentSession context directly) this.sessionConfigHandler = new SessionConfigHandler(this); - // Initialize RateLimitWatchdog — detects 429 exhaustion and schedules auto-retry - this.rateLimitWatchdog = new RateLimitWatchdog(session.id, this.stateManager); - this.rateLimitWatchdog.setRetryCallback(async (lastUserMessage) => { + // Initialize RateLimitWatchdog — detects 429/usage-limit exhaustion and + // drives two-phase recovery: (A) immediate fallback-model switch via the + // configured fallback chain, then (B) a cooldown at a parsed reset time or + // on a backoff ladder. Deps are injected here so the watchdog stays free of + // session/DB/provider coupling. + this.rateLimitWatchdog = new RateLimitWatchdog(session.id, this.stateManager, { + getCurrentModel: () => ({ + provider: (this.session.config.provider as string | undefined) ?? 'anthropic', + model: this.session.config.model ?? 'sonnet', + }), + resolveChain: () => { + const gs = this.settingsManager.getGlobalSettings(); + const { provider, model } = this.session.config; + return resolveFallbackChain( + (provider as string | undefined) ?? 'anthropic', + model ?? 'sonnet', + gs.modelFallbackMap, + gs.fallbackModels + ); + }, + isEntryAvailable: async (entry) => { + try { + const reg = getProviderRegistry(); + const p = reg.detectProviderForModel(entry.model, entry.provider); + if (!p) return false; + const ok = await Promise.resolve(p.isAvailable()); + if (!ok) return false; + if (typeof p.getAuthStatus === 'function') { + const auth = await p.getAuthStatus(); + return auth.isAuthenticated; + } + return true; + } catch { + return false; + } + }, + switchAndRetry: (lastUserMessage, entry) => + this.switchAndRetryForFallback(lastUserMessage, entry), + notifyPause: (payload) => { + this.internalEventBus.publish('session.rate_limit_pause', { + sessionId: this.session.id, + kind: payload.kind, + resetAt: payload.resetAt, + reason: payload.reason, + }); + }, + notifyResume: () => { + this.internalEventBus.publish('session.rate_limit_resume', { + sessionId: this.session.id, + }); + }, + }); + this.rateLimitWatchdog.setRetryCallback(async (lastUserMessage, switchTo) => { + if (switchTo) { + // A cooldown that was scheduled after a fallback switch re-switches + // before re-enqueuing (rare). Goes through the same timing-safe path. + await this.switchAndRetryForFallback(lastUserMessage, switchTo); + return; + } await this.executeRateLimitAutoRetry(lastUserMessage); }); @@ -772,6 +831,57 @@ export class AgentSession // Rate Limit Auto-Retry // ============================================================================ + /** + * Switch to a fallback model and re-enqueue the last user message (Phase A of + * rate-limit recovery). Timing-critical: this MUST run after the failed + * query's `finally` block completes, so we `await this.queryPromise` first + * (it resolves only once query-runner has torn the query down — queryObject + * nulled, env restored, setIdle called). By then the query is inactive, so + * `handleModelSwitch` takes its config-only branch (model-switch-handler) and + * `executeRateLimitAutoRetry` starts a fresh query with the new model. + * + * Returns false if the switch itself failed so the watchdog marks the entry + * tried and advances to the next chain entry. + */ + private async switchAndRetryForFallback( + lastUserMessage: { uuid: string; content: string | MessageContent[] } | null, + entry: FallbackModelEntry + ): Promise { + if (!lastUserMessage) { + this.logger.warn('Fallback switch skipped: no last user message available.'); + await this.stateManager.setIdle(); + return false; + } + try { + // (1) Wait for the failed query's cleanup to finish before mutating config. + if (this.queryPromise) { + try { + await this.queryPromise; + } catch { + // The failed query already rejected; its finally has still run. + } + } + + // (2) Switch model. The query is inactive now → config-only branch, no restart. + const result = await this.handleModelSwitch(entry.model, entry.provider); + if (!result.success) { + this.logger.warn( + `Fallback switch to ${entry.provider}/${entry.model} failed: ${result.error}. ` + + `Will try the next chain entry.` + ); + return false; + } + + // (3) Re-enqueue with the new model and start a fresh query. + await this.executeRateLimitAutoRetry(lastUserMessage); + return true; + } catch (err) { + this.logger.error('Fallback switch-and-retry failed:', err); + await this.stateManager.setIdle(); + return false; + } + } + /** * Execute auto-retry after rate limit cooldown. * Re-enqueues the last user message and starts a new query. diff --git a/packages/daemon/src/lib/agent/fallback-recovery.ts b/packages/daemon/src/lib/agent/fallback-recovery.ts new file mode 100644 index 000000000..bf81bcecd --- /dev/null +++ b/packages/daemon/src/lib/agent/fallback-recovery.ts @@ -0,0 +1,343 @@ +/** + * Fallback model chain + format-agnostic rate/usage-limit reset handling. + * + * Pure module (no session/DB deps) consumed by `rate-limit-watchdog.ts`. Kept + * separate so the chain resolution, timestamp extraction, and backoff ladder + * are fully unit-testable without an AgentSession or database. + * + * Background: `GlobalSettings.fallbackModels` / `modelFallbackMap` + * (`@hyperneo/shared` settings.ts) were editable in the UI but had no runtime + * consumers. On 429/usage-cap exhaustion this module resolves the chain, picks + * the next untried model, and — when the chain is exhausted — computes a + * cooldown from a reset time extracted format-agnostically from the error text + * (never by matching vendor-specific phrasing) or an exponential backoff ladder. + */ + +import type { FallbackModelEntry } from '@hyperneo/shared'; + +export type { FallbackModelEntry }; + +/** Maximum plausible quota-reset window. Parsed timestamps beyond this are rejected. */ +export const MAX_RESET_HORIZON_MS = 7 * 24 * 60 * 60 * 1000; // 7 days + +/** Buffer added to a parsed reset time so we retry just after the window lifts. */ +export const RESET_BUFFER_MS = 30 * 1000; // 30s + +/** Backoff ladder (ms), indexed by cooldown step. Used when no reset time is known. */ +export const BACKOFF_LADDER_MS: readonly number[] = [ + 10 * 60 * 1000, // 10m + 30 * 60 * 1000, // 30m + 60 * 60 * 1000, // 1h + 2 * 60 * 60 * 1000, // 2h + 4 * 60 * 60 * 1000, // 4h +]; + +/** Hard cap for any single backoff wait. */ +export const BACKOFF_CAP_MS = 8 * 60 * 60 * 1000; // 8h + +/** Jitter fraction: actual wait = capped base * (1 ± BACKOFF_JITTER). */ +export const BACKOFF_JITTER = 0.15; + +/** Minimum backoff wait (floor) so jitter can't shrink a step below this. */ +const BACKOFF_FLOOR_MS = 60 * 1000; // 1m + +/** + * Stable dedup key for a fallback entry. We never fall back to the same + * provider+model that just failed, so this key is also the "tried" marker. + */ +export function entryKey(entry: FallbackModelEntry): string { + return `${entry.provider}/${entry.model}`; +} + +/** + * Resolve the fallback chain for a (provider, model) pair. + * + * Priority: a non-empty `modelFallbackMap["${provider}/${model}"]` overrides the + * global `fallbackModels` list. Returns a defensive copy so callers cannot + * mutate the live settings arrays. + * + * Pure: both settings fields are passed in by the caller. + */ +export function resolveFallbackChain( + provider: string, + model: string, + modelFallbackMap: Record | undefined, + fallbackModels: FallbackModelEntry[] | undefined +): FallbackModelEntry[] { + const key = entryKey({ provider, model }); + const override = modelFallbackMap?.[key]; + if (override && override.length > 0) { + return [...override]; + } + if (fallbackModels && fallbackModels.length > 0) { + return [...fallbackModels]; + } + return []; +} + +export type FallbackSkipReason = 'none' | 'tried' | 'unavailable'; + +export interface FallbackSelection { + /** The chosen next entry, or null if the chain is exhausted. */ + next: FallbackModelEntry | null; + /** True when no untried+available entry remains (caller falls through to cooldown). */ + exhausted: boolean; + /** Why the candidate(s) were skipped, for logging. */ + skipReason: FallbackSkipReason; +} + +/** + * Pick the next chain entry to try. + * + * Iteration order = chain order. The first entry whose key is NOT in `triedKeys` + * AND for which `isAvailable` returns true is returned. The caller must already + * have added the current (failed) provider/model to `triedKeys` so we never + * re-select it. + * + * `isAvailable` is a synchronous predicate; the watchdog pre-resolves async + * provider availability into a Set before calling so this function stays pure. + */ +export function selectNextFallback( + chain: FallbackModelEntry[], + triedKeys: ReadonlySet, + isAvailable: (entry: FallbackModelEntry) => boolean +): FallbackSelection { + if (chain.length === 0) { + return { next: null, exhausted: true, skipReason: 'none' }; + } + + let lastSkip: FallbackSkipReason = 'none'; + for (const entry of chain) { + if (triedKeys.has(entryKey(entry))) { + lastSkip = 'tried'; + continue; + } + if (!isAvailable(entry)) { + lastSkip = 'unavailable'; + continue; + } + return { next: entry, exhausted: false, skipReason: 'none' }; + } + return { next: null, exhausted: true, skipReason: lastSkip }; +} + +export type ResetTimestampStrategy = 'iso8601' | 'yyyymmdd-hms' | 'epoch-millis' | 'epoch-seconds'; + +export interface ParsedReset { + /** Epoch-ms of the parsed reset moment. */ + resetAtMs: number; + /** Which strategy matched, for telemetry/logging. */ + strategy: ResetTimestampStrategy; +} + +// ISO-8601 with an explicit offset or Z (most precise — unambiguous timezone). +const ISO_WITH_TZ_RE = + /(\d{4})-(\d{2})-(\d{2})[T ](\d{2}):(\d{2}):(\d{2})(?:\.\d+)?(Z|[+-]\d{2}:?\d{2})/; + +// YYYY-MM-DD HH:mm:ss with NO timezone (e.g. the Chinese relay shape). Parsed +// as daemon-local time. Tried only after ISO_WITH_TZ_RE so an explicit offset +// always wins. +const LOCAL_DATETIME_RE = /(\d{4})-(\d{2})-(\d{2})[ T](\d{2}):(\d{2}):(\d{2})/; + +// 13-digit epoch millis (word-bounded to avoid UUID/request-id fragments). +const EPOCH_MILLIS_RE = /\b\d{13}\b/; + +// 10-digit epoch seconds (word-bounded; tried last to minimise false positives). +const EPOCH_SECONDS_RE = /\b\d{10}\b/; + +function isValidReset(ms: number, now: number): boolean { + if (!Number.isFinite(ms)) return false; + return ms > now && ms < now + MAX_RESET_HORIZON_MS; +} + +function parseLocalGroups(groups: RegExpMatchArray): number { + const [, yyyy, mm, dd, hh, mi, ss] = groups; + // `new Date('YYYY-MM-DDTHH:mm:ss')` (no Z) parses as LOCAL time per ES spec + // (V8/Node). Reconstruct with a T separator from the space-separated capture. + return new Date(`${yyyy}-${mm}-${dd}T${hh}:${mi}:${ss}`).getTime(); +} + +function parseIsoWithTzGroups(groups: RegExpMatchArray): number { + const [, yyyy, mm, dd, hh, mi, ss, tz] = groups; + // Normalise offset shape to `±HH:MM` (drop a bare `Z`). + let offset = ''; + if (tz !== 'Z') { + const raw = tz.replace(':', ''); + offset = `${raw.slice(0, 3)}:${raw.slice(3)}`; + } + const iso = `${yyyy}-${mm}-${dd}T${hh}:${mi}:${ss}${tz === 'Z' ? 'Z' : offset}`; + return new Date(iso).getTime(); +} + +/** + * Extract the first plausible quota-reset timestamp from an error message. + * + * Locale- and vendor-agnostic: matches digit/separator shapes only, NEVER + * Chinese/English phrasing like "将在…重置" / "resets at". Vendors keep + * changing phrasing; the digit shape is stable. + * + * For the relay example + * `Request rejected (429) · [1308][已达到 5 小时的使用上限。您的限额将在 2026-07-22 17:55:10 重置。]` + * the `LOCAL_DATETIME_RE` strategy matches `2026-07-22 17:55:10` (the `[1308]` + * code is 4 digits and never matches the epoch regexes). + * + * Accepts only timestamps in the future and within `MAX_RESET_HORIZON_MS`. Past + * or far-future matches are rejected (returns null → caller uses backoff). + */ +export function extractResetTimestamp( + errorMessage: string, + now: number = Date.now() +): ParsedReset | null { + if (!errorMessage) return null; + + // 1. ISO-8601 with timezone. + const iso = errorMessage.match(ISO_WITH_TZ_RE); + if (iso) { + const ms = parseIsoWithTzGroups(iso); + if (isValidReset(ms, now)) { + return { resetAtMs: ms, strategy: 'iso8601' }; + } + } + + // 2. YYYY-MM-DD HH:mm:ss local. + const local = errorMessage.match(LOCAL_DATETIME_RE); + if (local) { + const ms = parseLocalGroups(local); + if (isValidReset(ms, now)) { + return { resetAtMs: ms, strategy: 'yyyymmdd-hms' }; + } + } + + // 3. Epoch millis (13 digits). + const millis = errorMessage.match(EPOCH_MILLIS_RE); + if (millis) { + const ms = Number.parseInt(millis[0], 10); + if (isValidReset(ms, now)) { + return { resetAtMs: ms, strategy: 'epoch-millis' }; + } + } + + // 4. Epoch seconds (10 digits). + const seconds = errorMessage.match(EPOCH_SECONDS_RE); + if (seconds) { + const ms = Number.parseInt(seconds[0], 10) * 1000; + if (isValidReset(ms, now)) { + return { resetAtMs: ms, strategy: 'epoch-seconds' }; + } + } + + return null; +} + +export type CooldownReason = 'parsed-reset' | 'backoff-ladder'; + +export interface CooldownDecision { + /** ms from now until the retry should fire. */ + delayMs: number; + /** Absolute retryAt epoch-ms (for state serialization / UI). */ + retryAtMs: number; + /** Why this delay was chosen. */ + reason: CooldownReason; + /** Ladder index (0-based) when reason==='backoff-ladder', else -1. */ + ladderIndex: number; + /** + * Whether this wait is "free" — does NOT count toward maxAutoRetries. + * - parsed-reset: true (we know when the window lifts; waiting isn't a guess). + * - backoff-ladder: false (each ladder step is one budgeted retry). + */ + freeWait: boolean; + /** Parsed reset (when reason==='parsed-reset'), for surfacing to the UI. */ + reset: ParsedReset | null; +} + +/** + * Decide the next cooldown delay after the fallback chain is exhausted. + * + * @param errorMessage The 429/usage-limit error text. + * @param cooldownRetryCount Number of (non-free) cooldown steps already used. + * @param now Injected for testability. + * @param jitterFn Injected randomness so tests are deterministic. Defaults to + * Math.random scaled to [-1, 1]. + */ +export function computeCooldown( + errorMessage: string, + cooldownRetryCount: number, + now: number = Date.now(), + jitterFn: () => number = () => Math.random() * 2 - 1 +): CooldownDecision { + const parsed = extractResetTimestamp(errorMessage, now); + if (parsed) { + const retryAtMs = parsed.resetAtMs + RESET_BUFFER_MS; + const delayMs = Math.max(0, parsed.resetAtMs - now) + RESET_BUFFER_MS; + return { + delayMs, + retryAtMs, + reason: 'parsed-reset', + ladderIndex: -1, + freeWait: true, + reset: parsed, + }; + } + + const lastIndex = BACKOFF_LADDER_MS.length - 1; + const ladderIndex = Math.min(cooldownRetryCount, lastIndex); + const base = Math.min(BACKOFF_LADDER_MS[ladderIndex], BACKOFF_CAP_MS); + const jitter = base * BACKOFF_JITTER * jitterFn(); + const delayMs = Math.max(BACKOFF_FLOOR_MS, Math.round(base + jitter)); + return { + delayMs, + retryAtMs: now + delayMs, + reason: 'backoff-ladder', + ladderIndex, + freeWait: false, + reset: null, + }; +} + +// Keywords (lowercased; ASCII + the Chinese cap/usage characters) that signal a +// usage CAP rather than a transient rate limit. Used to classify the paused +// status surfaced to the UI. ASCII-only matching is case-insensitive; the CJK +// characters are matched literally. +// +// Deliberately narrow: generic phrases like "exceeded" / "limit reached" appear +// in BOTH transient rate-limit messages ("rate limit exceeded, retry in 60s") +// and usage-cap messages, so they don't discriminate and are excluded. Only +// cap-specific terms classify as a usage_limit; everything else is a transient +// rate_limit. +const USAGE_CAP_KEYWORDS = [ + 'usage', + 'cap', + 'quota', + 'daily', + 'weekly', + '上限', + '额度', + '小时', + '周', +]; + +/** + * Classify a paused rate-limit episode as a short `rate_limit` or a + * daily/weekly `usage_limit`, for surfacing the right task status. + * + * Heuristic: a known reset time (parsed timestamp) implies a CAP window, so it + * is a `usage_limit`. Otherwise, cap/usage keywords in the message override to + * `usage_limit`; the default is a transient `rate_limit`. + */ +export function classifyLimitKind( + errorMessage: string, + decision: CooldownDecision +): 'rate_limit' | 'usage_limit' { + if (decision.reason === 'parsed-reset') { + return 'usage_limit'; + } + const lower = errorMessage.toLowerCase(); + if (USAGE_CAP_KEYWORDS.some((kw) => lower.includes(kw.toLowerCase()))) { + return 'usage_limit'; + } + // CJK keywords are not lowercased meaningfully; check the raw message too. + if (USAGE_CAP_KEYWORDS.some((kw) => errorMessage.includes(kw))) { + return 'usage_limit'; + } + return 'rate_limit'; +} diff --git a/packages/daemon/src/lib/agent/rate-limit-watchdog.ts b/packages/daemon/src/lib/agent/rate-limit-watchdog.ts index 9e1ec4626..43ee4b479 100644 --- a/packages/daemon/src/lib/agent/rate-limit-watchdog.ts +++ b/packages/daemon/src/lib/agent/rate-limit-watchdog.ts @@ -1,25 +1,47 @@ /** - * RateLimitWatchdog - Auto-retry after 429 rate limit exhaustion + * RateLimitWatchdog - Auto-recovery after 429 rate/usage-limit exhaustion. * - * When the Claude SDK exhausts all retries on a 429 (rate limited) error, - * the session stalls. This watchdog detects that state and schedules an - * automatic retry after a cooldown period. + * Two-phase recovery strategy: * - * Behavior: - * - Schedules a retry timer (default 10 minutes) when 429 exhaustion is detected - * - Max 3 auto-retry cycles before giving up entirely - * - Cancelable: if the user sends a new message, the timer is cleared - * - Fires the retry by re-enqueueing the last user message and restarting the query + * Phase A (fallback chain): when a 429/usage-limit exhausts the SDK's own + * retries, resolve the configured fallback chain (`GlobalSettings.modelFallbackMap` + * override, else `fallbackModels`) and immediately switch the session to the + * next untried, available entry via the existing model-switch machinery, then + * retry. Fallback switches are FREE — they don't count toward maxAutoRetries — + * and each is tracked per-episode so a 429 on the fallback advances to the + * next entry instead of retrying the same model. + * + * Phase B (cooldown): only when the chain is exhausted (or empty) does the + * watchdog schedule a timed retry. The delay is computed format-agnostically: + * a reset timestamp parsed from the error text (ISO-8601, YYYY-MM-DD HH:mm:ss, + * or epoch) wins; otherwise an exponential backoff ladder (10m→30m→1h→2h→4h, + * cap 8h, with jitter). Waits against a known reset do NOT count toward + * maxAutoRetries — only speculative backoff steps do. + * + * Pure recovery logic lives in `fallback-recovery.ts`; this class is the + * stateful orchestrator that owns the episode (tried entries, cooldown timer) + * and delegates model switching + event surfacing to injected deps. */ -import type { MessageContent } from '@hyperneo/shared'; +import type { FallbackModelEntry, MessageContent } from '@hyperneo/shared'; import type { ProcessingStateManager } from './processing-state-manager'; +import { + classifyLimitKind, + computeCooldown, + entryKey, + selectNextFallback, + type CooldownDecision, +} from './fallback-recovery'; import { Logger } from '../logger'; export interface RateLimitWatchdogConfig { - /** Cooldown period in ms before auto-retry (default: 600000 = 10 min) */ + /** + * Fallback cooldown used only if cooldown computation is bypassed defensively. + * Normal cooldowns are driven by `computeCooldown` (parsed reset or backoff + * ladder). Kept for backward compatibility and as a last-resort default. + */ cooldownMs: number; - /** Maximum number of auto-retry cycles (default: 3) */ + /** Maximum number of (non-free) cooldown steps before giving up entirely. */ maxAutoRetries: number; } @@ -28,121 +50,265 @@ const DEFAULT_CONFIG: RateLimitWatchdogConfig = { maxAutoRetries: 3, }; +export type RateLimitWatchdogStatus = 'idle' | 'cooldown' | 'fallback-pending'; + export interface RateLimitWatchdogState { - status: 'idle' | 'cooldown'; + status: RateLimitWatchdogStatus; + /** Number of non-free cooldown steps used this episode. */ retryCount: number; maxRetries: number; retryAt: number | null; lastUserMessage: { uuid: string; content: string | MessageContent[] } | null; + /** provider/model keys attempted this episode (for UI/diagnostics). */ + triedEntries: string[]; + /** Resolved fallback chain in effect, or null if not applicable/resolved. */ + fallbackChain: FallbackModelEntry[] | null; + /** True briefly while an immediate fallback switch is in flight. */ + fallbackPending: boolean; + /** Classified limit kind for the active cooldown (for UI). */ + limitKind: 'rate_limit' | 'usage_limit' | null; +} + +/** + * Payload emitted via `notifyPause` so listeners (e.g. the Space runtime) can + * surface a paused task status with a resume-at timestamp. + */ +export interface RateLimitPausePayload { + kind: 'rate_limit' | 'usage_limit'; + /** Epoch-ms when the limit is expected to reset (if known). */ + resetAt?: number; + /** Short human-readable reason (the cooldown decision reason). */ + reason: string; } /** - * Callback type for when the cooldown timer fires and a retry should be attempted. - * The implementor (AgentSession) is responsible for re-enqueueing the message - * and restarting the query. + * Dependencies injected by AgentSession. Kept as an interface so tests inject + * fakes without a DB or live session. All members are optional except the model + * + chain resolvers the watchdog needs to function. + */ +export interface RateLimitWatchdogDeps { + /** Returns the currently-configured (provider, model) — the failing one. */ + getCurrentModel(): { provider: string; model: string }; + /** Returns the resolved fallback chain for the current (provider, model). */ + resolveChain(): FallbackModelEntry[]; + /** + * Returns true iff the entry's provider is registered AND authenticated. + * The implementation may perform async checks; the watchdog awaits them. + */ + isEntryAvailable(entry: FallbackModelEntry): Promise; + /** + * Switch the session to the entry AND re-enqueue the last user message. + * MUST await the failed query's completion before switching (the AgentSession + * implementation handles this). Returns false if the switch itself failed so + * the watchdog can advance to the next entry. + */ + switchAndRetry( + lastUserMessage: { uuid: string; content: string | MessageContent[] } | null, + entry: FallbackModelEntry + ): Promise; + /** Surface a paused state (Space runtime marks the task rate/usage-limited). */ + notifyPause?(payload: RateLimitPausePayload): void; + /** Surface that the pause is over (task restored to in_progress). */ + notifyResume?(): void; +} + +/** + * Callback type for when recovery fires a retry. The implementor (AgentSession) + * re-enqueues the message and restarts the query. When `switchTo` is set, the + * callback switches to that model first (used when a cooldown scheduled after a + * fallback switch needs to re-switch — rare). */ export type RateLimitRetryCallback = ( - lastUserMessage: { uuid: string; content: string | MessageContent[] } | null + lastUserMessage: { uuid: string; content: string | MessageContent[] } | null, + switchTo?: FallbackModelEntry ) => Promise; export class RateLimitWatchdog { private logger: Logger; private config: RateLimitWatchdogConfig; + private deps: RateLimitWatchdogDeps; private retryCount = 0; private cooldownTimer: ReturnType | null = null; private lastUserMessage: { uuid: string; content: string | MessageContent[] } | null = null; + private lastErrorMessage = ''; private retryCallback: RateLimitRetryCallback | null = null; private stateManager: ProcessingStateManager; + // Episode state. + private triedKeys = new Set(); + private chain: FallbackModelEntry[] | null = null; + private fallbackPending = false; + private limitKind: 'rate_limit' | 'usage_limit' | null = null; + private paused = false; + constructor( sessionId: string, stateManager: ProcessingStateManager, + deps: RateLimitWatchdogDeps, config?: Partial ) { this.logger = new Logger(`RateLimitWatchdog ${sessionId}`); this.config = { ...DEFAULT_CONFIG, ...config }; this.stateManager = stateManager; + this.deps = deps; } - /** - * Set the callback invoked when the cooldown timer fires. - */ + /** Set the callback invoked when a cooldown elapses (or Retry Now is pressed). */ setRetryCallback(callback: RateLimitRetryCallback): void { this.retryCallback = callback; } - /** - * Get current watchdog state (for serialization / UI). - */ + /** Get current watchdog state (for serialization / UI). */ getState(): RateLimitWatchdogState { const retryAt = this.cooldownTimer !== null ? Date.now() + this.getRemainingMs() : null; return { - status: this.cooldownTimer !== null ? 'cooldown' : 'idle', + status: this.fallbackPending + ? 'fallback-pending' + : this.cooldownTimer !== null + ? 'cooldown' + : 'idle', retryCount: this.retryCount, maxRetries: this.config.maxAutoRetries, retryAt, lastUserMessage: this.lastUserMessage, + triedEntries: [...this.triedKeys], + fallbackChain: this.chain, + fallbackPending: this.fallbackPending, + limitKind: this.limitKind, }; } /** - * Schedule an auto-retry after cooldown. Called when a 429 exhaustion error - * is detected in QueryRunner's error handling. + * React to a 429/usage-limit exhaustion. + * + * Phase A: switch to the next untried fallback entry and retry immediately + * (free — no retryCount bump). Phase B: when the chain is exhausted, schedule + * a cooldown computed from a parsed reset time or the backoff ladder. * - * @returns true if scheduled, false if max retries exceeded + * @returns true if recovery was engaged (caller skips the terminal error + * broadcast), false if recovery is impossible (no message, or budget + * exhausted). */ async scheduleRetry( errorMessage: string, lastUserMessage: { uuid: string; content: string | MessageContent[] } | null ): Promise { - // Cancel any existing timer - this.cancel(); + // Cancel any existing timer (a previous cooldown or in-flight switch path). + this.cancelCooldownTimer(); + this.lastErrorMessage = errorMessage; - // Cannot retry without a message to re-enqueue — fail fast instead of - // scheduling a no-op cooldown timer that would abort on fire. + // Cannot retry without a message to re-enqueue — fail fast. if (!lastUserMessage) { - this.logger.warn('Cannot schedule rate limit auto-retry: no user message to retry.'); + this.logger.warn('Cannot schedule rate limit recovery: no user message to retry.'); return false; } + this.lastUserMessage = lastUserMessage; + + // Mark the CURRENT (failed) provider+model as tried so we never re-select it. + const { provider, model } = this.deps.getCurrentModel(); + this.triedKeys.add(entryKey({ provider, model })); + + // Resolve the chain once per episode. + if (this.chain === null) { + this.chain = this.deps.resolveChain(); + } + + // ── Phase A: try an immediate fallback switch (free retry). ──────────── + if (this.chain.length > 0) { + const availability = new Map(); + await Promise.all( + this.chain.map(async (entry) => { + availability.set(entry, await this.deps.isEntryAvailable(entry)); + }) + ); + const sel = selectNextFallback( + this.chain, + this.triedKeys, + (e) => availability.get(e) === true + ); + + if (sel.next) { + this.logger.info( + `Fallback: switching to ${sel.next.provider}/${sel.next.model} ` + + `(skipReason for prior candidates: ${sel.skipReason}).` + ); + this.fallbackPending = true; + // Fire-and-forget the switch; scheduleRetry must resolve to true (it is + // awaited at query-runner.ts) so the caller skips the terminal error + // broadcast + setIdle. The switch itself runs after the failed query's + // finally via switchAndRetry's `await queryPromise`. + void this.fireImmediateFallback(lastUserMessage, sel.next); + return true; + } + this.logger.info( + `Fallback chain exhausted (${sel.skipReason}); falling through to cooldown.` + ); + } - if (this.retryCount >= this.config.maxAutoRetries) { + // ── Phase B: chain exhausted (or empty) — schedule a cooldown. ───────── + const decision = computeCooldown(errorMessage, this.retryCount); + + // Only speculative backoff steps count toward the budget; parsed-reset + // waits are free (we know when the window lifts). + if (!decision.freeWait && this.retryCount >= this.config.maxAutoRetries) { this.logger.warn( - `Max auto-retries (${this.config.maxAutoRetries}) exceeded for 429 error. ` + - `Giving up. Error: ${errorMessage}` + `Max auto-retries (${this.config.maxAutoRetries}) exceeded for 429 error. Giving up. ` + + `Error: ${errorMessage}` ); return false; } + if (!decision.freeWait) { + this.retryCount++; + } - this.lastUserMessage = lastUserMessage; - this.retryCount++; + await this.scheduleCooldown(errorMessage, decision); + return true; + } - const retryAt = Date.now() + this.config.cooldownMs; + /** + * Schedule (and notify) a cooldown wait. Extracted so the re-entry path on a + * failed immediate switch can reuse it without re-running Phase A. + */ + private async scheduleCooldown(errorMessage: string, decision: CooldownDecision): Promise { + const kind = classifyLimitKind(errorMessage, decision); + this.limitKind = kind; + const retryAt = decision.retryAtMs; this.logger.info( - `Scheduling auto-retry #${this.retryCount}/${this.config.maxAutoRetries} ` + - `in ${this.config.cooldownMs}ms. Error: ${errorMessage}` + `Scheduling cooldown (${decision.reason}, kind=${kind}) in ${decision.delayMs}ms. ` + + `Error: ${errorMessage}` ); - // Set processing state to rate_limit_cooldown (awaited to prevent - // race with user messages or retry-now overwriting the state). + // Flip the processing state FIRST, then surface the pause via the bus. This + // avoids a brief window where a subscriber reacting to the pause event + // observes the prior (idle/processing) state (and an onIdleCallback firing). await this.stateManager.setRateLimitCooldown({ retryCount: this.retryCount, maxRetries: this.config.maxAutoRetries, retryAt, }); + // Surface the paused state before arming the timer so listeners can mark + // the task rate/usage-limited with a resume-at timestamp. + this.notifyPause({ + kind, + resetAt: decision.reset ? decision.retryAtMs : undefined, + reason: decision.reason, + }); + this.cooldownTimer = setTimeout(() => { this.cooldownTimer = null; this.logger.info( - `Cooldown elapsed for auto-retry #${this.retryCount}/${this.config.maxAutoRetries}. Firing retry.` + `Cooldown elapsed; firing retry (step ${this.retryCount}/${this.config.maxAutoRetries}).` ); + // The pause is over — notify before the retry restarts the query. + this.notifyResume(); if (this.retryCallback) { - void this.retryCallback(this.lastUserMessage); + void this.retryCallback(this.lastUserMessage, undefined); } - }, this.config.cooldownMs); + }, decision.delayMs); - // Unref so the timer doesn't keep the process alive if ( this.cooldownTimer && typeof this.cooldownTimer === 'object' && @@ -150,32 +316,84 @@ export class RateLimitWatchdog { ) { this.cooldownTimer.unref(); } + } - return true; + /** + * Fire an immediate fallback switch. On switch failure (or any thrown error), + * mark the entry tried and re-enter scheduleRetry to try the next entry or + * fall through to a cooldown (retry count NOT incremented for the failed + * switch). Wrapped in try/catch so a rejection can never leave + * `fallbackPending` stuck true (which would freeze getState/retryNow). + */ + private async fireImmediateFallback( + lastUserMessage: { uuid: string; content: string | MessageContent[] } | null, + entry: FallbackModelEntry + ): Promise { + let ok = false; + try { + ok = await this.deps.switchAndRetry(lastUserMessage, entry); + } catch (err) { + this.logger.error( + `Fallback switch to ${entry.provider}/${entry.model} threw; advancing to next entry:`, + err + ); + ok = false; + } finally { + this.fallbackPending = false; + } + if (ok) { + return; + } + this.triedKeys.add(entryKey(entry)); + // Re-enter recovery with the same error/message to try the next entry or + // schedule a cooldown. Guard against losing the message and against a + // rejecting re-entry (which would otherwise escape unhandled). + if (this.lastUserMessage) { + try { + await this.scheduleRetry(this.lastErrorMessage, this.lastUserMessage); + } catch (err) { + // scheduleRetry schedules a cooldown or another switch; if it rejects + // (e.g. setRateLimitCooldown's DB write), fall back to a single + // best-effort cooldown so recovery isn't silently lost. + this.logger.error('Fallback re-entry rejected; scheduling a deferrive cooldown:', err); + try { + await this.scheduleCooldown( + this.lastErrorMessage, + computeCooldown(this.lastErrorMessage, this.retryCount) + ); + } catch { + // Nothing more we can do; the caller (query-runner) will surface the + // original 429 via its normal error path on the next message. + } + } + } } /** - * Cancel any pending auto-retry. Called when: - * - User sends a new message before the retry fires - * - User explicitly cancels the auto-retry - * - Session is cleaned up + * Cancel any pending cooldown timer. Called on new user input, explicit + * cancel, reset, or cleanup. Notifies resume so a paused task is restored. */ cancel(): void { + this.cancelCooldownTimer(); + this.fallbackPending = false; + this.notifyResume(); + } + + private cancelCooldownTimer(): void { if (this.cooldownTimer !== null) { clearTimeout(this.cooldownTimer); this.cooldownTimer = null; - this.logger.info('Cancelled pending rate limit auto-retry.'); + this.logger.info('Cancelled pending rate limit cooldown.'); } } /** - * Immediately trigger the retry (bypassing the cooldown). - * Used when the user clicks "Retry Now" in the UI. - * - * @returns true if a retry was pending and fired, false if no retry was pending + * Immediately trigger the cooldown retry (bypassing the wait). Used when the + * user clicks "Retry Now". Returns false if no cooldown is pending or a + * fallback switch is in flight (the switch should be allowed to complete). */ retryNow(): boolean { - if (this.cooldownTimer === null) { + if (this.fallbackPending || this.cooldownTimer === null) { return false; } @@ -183,50 +401,63 @@ export class RateLimitWatchdog { this.cooldownTimer = null; this.logger.info( - `Immediate retry triggered for auto-retry #${this.retryCount}/${this.config.maxAutoRetries}.` + `Immediate retry triggered (step ${this.retryCount}/${this.config.maxAutoRetries}).` ); - + this.notifyResume(); if (this.retryCallback) { - void this.retryCallback(this.lastUserMessage); + void this.retryCallback(this.lastUserMessage, undefined); } - return true; } /** - * Reset the watchdog entirely (e.g., on successful API call). + * Reset the watchdog entirely (e.g. on a successful API call). Clears the + * episode so the next 429 starts a fresh fallback chain. */ reset(): void { this.cancel(); this.retryCount = 0; this.lastUserMessage = null; + this.lastErrorMessage = ''; + this.triedKeys.clear(); + this.chain = null; + this.limitKind = null; } - /** - * Get remaining milliseconds until the retry fires. - * Returns 0 if no retry is scheduled. - */ + /** Remaining ms until the cooldown fires (0 if none scheduled). */ private getRemainingMs(): number { - // Approximation — the timer was set with config.cooldownMs, - // so remaining time is tracked via the scheduled retryAt from the state. const state = this.stateManager.getState(); if (state.status === 'rate_limit_cooldown') { - const remaining = state.retryAt - Date.now(); - return Math.max(0, remaining); + return Math.max(0, state.retryAt - Date.now()); } return 0; } - /** - * Check if a cooldown auto-retry is currently scheduled. - */ + /** Is a cooldown auto-retry currently scheduled? */ isPending(): boolean { return this.cooldownTimer !== null; } - /** - * Cleanup (called during session cleanup). - */ + private notifyPause(payload: RateLimitPausePayload): void { + this.paused = true; + try { + this.deps.notifyPause?.(payload); + } catch (err) { + this.logger.warn('notifyPause callback threw:', err); + } + } + + private notifyResume(): void { + if (!this.paused) return; + this.paused = false; + try { + this.deps.notifyResume?.(); + } catch (err) { + this.logger.warn('notifyResume callback threw:', err); + } + } + + /** Cleanup (called during session cleanup). */ destroy(): void { this.cancel(); this.retryCallback = null; diff --git a/packages/daemon/src/lib/internal-event-bus.ts b/packages/daemon/src/lib/internal-event-bus.ts index da84617a7..346d90e6b 100644 --- a/packages/daemon/src/lib/internal-event-bus.ts +++ b/packages/daemon/src/lib/internal-event-bus.ts @@ -394,6 +394,26 @@ export interface SessionEvents { 'commands.updated': { sessionId: string; commands: string[] }; 'session.error': { sessionId: string; error: string; details?: unknown }; 'session.errorClear': { sessionId: string }; + /** + * A session has paused after rate/usage-limit exhaustion with no fallback + * left (chain exhausted). Listeners (e.g. the Space runtime) surface a paused + * task status with a resume-at timestamp. Emitted by RateLimitWatchdog. + */ + 'session.rate_limit_pause': { + sessionId: string; + /** 'rate_limit' (transient) or 'usage_limit' (daily/weekly cap). */ + kind: 'rate_limit' | 'usage_limit'; + /** Epoch-ms when the limit is expected to reset, if known. */ + resetAt?: number; + /** Short reason string (the cooldown decision reason). */ + reason: string; + }; + /** + * A previously rate-limited session has resumed (cooldown fired, user + * cancelled/retried, or the API call succeeded). Listeners clear any paused + * task status. Emitted by RateLimitWatchdog. + */ + 'session.rate_limit_resume': { sessionId: string }; } /** diff --git a/packages/daemon/src/lib/space/managers/space-task-manager.ts b/packages/daemon/src/lib/space/managers/space-task-manager.ts index 18df2a73c..c2c642e12 100644 --- a/packages/daemon/src/lib/space/managers/space-task-manager.ts +++ b/packages/daemon/src/lib/space/managers/space-task-manager.ts @@ -55,6 +55,11 @@ export const VALID_SPACE_TASK_TRANSITIONS: Record { + const spaces = await this.listActiveSpaces(); + const now = Date.now(); + for (const space of spaces) { + for (const task of this.config.taskRepo.listRateLimitedBySpace(space.id)) { + const resetAt = task.restrictions?.resetAt; + if (resetAt !== undefined && resetAt > now) continue; // still waiting + try { + await this.updateTaskAndEmit(space.id, task.id, { + status: 'in_progress', + restrictions: null, + }); + log.info( + `SpaceRuntime: auto-resumed ${task.status} task ${task.id} (resetAt ${resetAt ?? 'none'} ≤ now) — rehydration will restart the worker.` + ); + } catch (err) { + log.warn( + `SpaceRuntime: failed to auto-resume paused task ${task.id}: ${err instanceof Error ? err.message : String(err)}` + ); + } + } + } + } + private async checkStandaloneTasks(): Promise { const spaces = await this.listActiveSpaces(); diff --git a/packages/daemon/src/lib/space/runtime/task-agent-manager.ts b/packages/daemon/src/lib/space/runtime/task-agent-manager.ts index d17183a5b..e7a42502b 100644 --- a/packages/daemon/src/lib/space/runtime/task-agent-manager.ts +++ b/packages/daemon/src/lib/space/runtime/task-agent-manager.ts @@ -448,10 +448,16 @@ export class TaskAgentManager { * Populated on first cleanup subscription attempt; cleared in `cleanupAll()`. */ private taskArchiveListenerUnsub: (() => void) | null = null; + /** + * Unsubs for the rate-limit pause/resume listeners (one each). Torn down in + * `cleanupAll()`. + */ + private rateLimitListenerUnsubs: Array<() => void> = []; constructor(private readonly config: TaskAgentManagerConfig) { this.auditLogRepo = new McpAuditLogRepository(this.config.db.getDatabase()); this.subscribeToTaskArchiveEvents(); + this.subscribeToRateLimitEvents(); } *getTrackedAgentRootPids(): Iterable { @@ -515,6 +521,96 @@ export class TaskAgentManager { ); } + /** + * Surface a paused task status when a worker session hits a rate/usage cap + * with no fallback left, and restore it on resume. + * + * With fallback-chain recovery (Parts A+B) the 429 error broadcast is skipped, + * so a paused session never emits `session.error` and the task never fails. + * These listeners add the visible status: on pause, mark the parent task + * `rate_limited` / `usage_limited` with a `restrictions` resume-at blob; on + * resume, restore `in_progress` and clear the blob. Global subscriptions — the + * sessionId on the payload resolves to the parent task via + * `findParentTaskIdForSubSession`. + */ + private subscribeToRateLimitEvents(): void { + if (this.rateLimitListenerUnsubs.length > 0) return; + + this.rateLimitListenerUnsubs.push( + this.config.internalEventBus.subscribe( + 'session.rate_limit_pause', + (event) => { + const taskId = this.findParentTaskIdForSubSession(event.sessionId); + if (!taskId) return; + const status = event.kind === 'usage_limit' ? 'usage_limited' : 'rate_limited'; + void this.markTaskRateLimited(taskId, status, event.resetAt, event.reason).catch( + (err) => { + log.warn( + `TaskAgentManager: failed to mark task ${taskId} ${status} for session ${event.sessionId}:`, + err + ); + } + ); + }, + { subscriberName: 'TaskAgentManager.rateLimitPause' } + ) + ); + + this.rateLimitListenerUnsubs.push( + this.config.internalEventBus.subscribe( + 'session.rate_limit_resume', + (event) => { + const taskId = this.findParentTaskIdForSubSession(event.sessionId); + if (!taskId) return; + void this.restoreTaskFromRateLimit(taskId).catch((err) => { + log.warn( + `TaskAgentManager: failed to restore task ${taskId} from rate limit for session ${event.sessionId}:`, + err + ); + }); + }, + { subscriberName: 'TaskAgentManager.rateLimitResume' } + ) + ); + } + + /** Mark a task paused on a rate/usage limit with a resume-at restriction. */ + private async markTaskRateLimited( + taskId: string, + status: 'rate_limited' | 'usage_limited', + resetAt: number | undefined, + reason: string + ): Promise { + const task = this.config.taskRepo.getTask(taskId); + if (!task) return; + // Don't override a terminal status the runtime has already decided on. + if ( + task.status === 'done' || + task.status === 'cancelled' || + task.status === 'archived' || + task.status === 'blocked' + ) { + return; + } + this.config.taskRepo.updateTask(taskId, { + status, + restrictions: { + type: status === 'usage_limited' ? 'usage_limit' : 'rate_limit', + limit: reason, + resetAt: resetAt ?? Date.now() + 60 * 60 * 1000, + sessionRole: 'worker', + }, + }); + } + + /** Restore a rate/usage-limited task to in_progress (resume). */ + private async restoreTaskFromRateLimit(taskId: string): Promise { + const task = this.config.taskRepo.getTask(taskId); + if (!task) return; + if (task.status !== 'rate_limited' && task.status !== 'usage_limited') return; + this.config.taskRepo.updateTask(taskId, { status: 'in_progress', restrictions: null }); + } + /** * Archive pipeline for a task that has transitioned to `archived`. * @@ -2053,6 +2149,8 @@ export class TaskAgentManager { this.taskArchiveListenerUnsub(); this.taskArchiveListenerUnsub = null; } + for (const unsub of this.rateLimitListenerUnsubs) unsub(); + this.rateLimitListenerUnsubs = []; const taskIds = Array.from(this.subSessions.keys()); await Promise.allSettled(taskIds.map((taskId) => this.shutdownTask(taskId))); log.info(`TaskAgentManager: cleanupAll complete (${taskIds.length} tasks shut down)`); diff --git a/packages/daemon/src/storage/repositories/space-task-repository.ts b/packages/daemon/src/storage/repositories/space-task-repository.ts index af652e4a3..1ac4da124 100644 --- a/packages/daemon/src/storage/repositories/space-task-repository.ts +++ b/packages/daemon/src/storage/repositories/space-task-repository.ts @@ -13,6 +13,7 @@ import type { InternalCreateSpaceTaskParams, InternalUpdateSpaceTaskParams, } from '@hyperneo/shared'; +import type { TaskRestriction } from '@hyperneo/shared/types/neo'; import type { ReactiveDatabase } from '../reactive-database'; import type { SQLiteValue } from '../types'; @@ -305,6 +306,20 @@ export class SpaceTaskRepository { return rows.map((r) => this.rowToSpaceTask(r)); } + /** + * List tasks paused on a rate/usage limit (`rate_limited` / `usage_limited`) + * within a space. Used by the runtime tick loop to auto-resume tasks whose + * reset time has passed — including across daemon restarts (the resume is + * driven off the persisted `restrictions.resetAt`, not an in-memory timer). + */ + listRateLimitedBySpace(spaceId: string): SpaceTask[] { + const stmt = this.db.prepare( + `SELECT * FROM space_tasks WHERE space_id = ? AND status IN ('rate_limited', 'usage_limited') ORDER BY updated_at DESC, id DESC` + ); + const rows = stmt.all(spaceId) as Record[]; + return rows.map((r) => this.rowToSpaceTask(r)); + } + /** * List tasks by status within a space, optionally filtered by block reason, * paginated, returning both the page of tasks and the total count. @@ -501,6 +516,20 @@ export class SpaceTaskRepository { fields.push('active_session = ?'); values.push(null); } + // Auto-clear a stale rate/usage-limit restriction when the task leaves the + // paused statuses via any path that doesn't set `restrictions` explicitly + // (e.g. a manual setTaskStatus to in_progress/cancelled/done). The runtime + // resume path sets restrictions:null itself; this is the backstop so a + // paused task can't keep a stale resume-at blob after manual intervention. + if ( + params.restrictions === undefined && + params.status !== undefined && + params.status !== 'rate_limited' && + params.status !== 'usage_limited' + ) { + fields.push('restrictions = ?'); + values.push(null); + } if (params.taskAgentSessionId !== undefined) { fields.push('task_agent_session_id = ?'); values.push(params.taskAgentSessionId ?? null); @@ -571,6 +600,10 @@ export class SpaceTaskRepository { fields.push('post_approval_blocked_reason = ?'); values.push(params.postApprovalBlockedReason ?? null); } + if (params.restrictions !== undefined) { + fields.push('restrictions = ?'); + values.push(params.restrictions ? JSON.stringify(params.restrictions) : null); + } if (fields.length > 0) { fields.push('updated_at = ?'); @@ -800,6 +833,7 @@ export class SpaceTaskRepository { postApprovalSessionId: (row.post_approval_session_id as string | null) ?? null, postApprovalStartedAt: (row.post_approval_started_at as number | null) ?? null, postApprovalBlockedReason: (row.post_approval_blocked_reason as string | null) ?? null, + restrictions: parseRestrictions(row.restrictions), createdAt: row.created_at as number, startedAt: (row.started_at as number | null) ?? null, completedAt: (row.completed_at as number | null) ?? null, @@ -807,3 +841,24 @@ export class SpaceTaskRepository { }; } } + +/** + * Parse a `restrictions` JSON blob into a `TaskRestriction`, or null when the + * task is not paused. Defensive against malformed/missing JSON. + */ +function parseRestrictions(raw: unknown): TaskRestriction | null { + if (typeof raw !== 'string' || raw.length === 0) return null; + try { + const parsed = JSON.parse(raw) as TaskRestriction; + if ( + parsed && + (parsed.type === 'rate_limit' || parsed.type === 'usage_limit') && + typeof parsed.resetAt === 'number' + ) { + return parsed; + } + return null; + } catch { + return null; + } +} diff --git a/packages/daemon/src/storage/schema/migrations.ts b/packages/daemon/src/storage/schema/migrations.ts index 7e41ee8a4..08d29e1b6 100644 --- a/packages/daemon/src/storage/schema/migrations.ts +++ b/packages/daemon/src/storage/schema/migrations.ts @@ -733,6 +733,11 @@ export function runMigrations(db: BunDatabase, createBackup: () => void): void { // evolution-finding domain to their rebranded identifiers so persisted rows match // the code after the HyperNeo rename. run(migrationMarkerKey(162), () => runMigration162(db)); + + // Migration 163: Add rate_limited/usage_limited to space_tasks.status CHECK and + // a restrictions column, so worker sessions paused on rate/usage caps can + // surface a distinct status with a resume-at timestamp. + run(migrationMarkerKey(163), () => runMigration163(db)); } function migrationMarkerKey(version: number): string { @@ -10950,3 +10955,95 @@ export function runMigration162(db: BunDatabase): void { ).run(); } } + +/** + * Migration 163: Add `rate_limited` / `usage_limited` to `space_tasks.status` + * CHECK constraint, and add a nullable `restrictions` column. + * + * Worker sessions paused on a rate/usage cap (chain exhausted) now surface a + * distinct task status with a resume-at timestamp (`restrictions` JSON) instead + * of failing. SQLite cannot ALTER a CHECK constraint, so the table is rebuilt + * from its live DDL (which is current — later rebuilds re-synced it after the + * ALTER-added columns). The `restrictions` column is injected into the new DDL + * so the column copy (which references the pre-restriction column set) leaves + * it NULL for existing rows. Idempotent. + */ +export function runMigration163(db: BunDatabase): void { + if (!tableExists(db, 'space_tasks')) return; + + // Already widened → only backfill the restrictions column if a partial run + // left it missing. + if (statusCheckContains(db, 'space_tasks', 'rate_limited')) { + if (!tableHasColumn(db, 'space_tasks', 'restrictions')) { + db.exec(`ALTER TABLE space_tasks ADD COLUMN restrictions TEXT`); + } + return; + } + + const currentSql = tableCreateSql(db, 'space_tasks'); + if (!currentSql) { + throw new Error('Migration 163: space_tasks CREATE TABLE sql not found'); + } + + // Columns present before the rebuild (no `restrictions` yet) — used for the + // INSERT … SELECT copy. The new table adds `restrictions` (nullable) on top. + const copyColumns = tableColumnNames(db, 'space_tasks').map(quoteSqlIdent).join(', '); + const existingIndexDdl = capturedIndexDdl(db, 'space_tasks'); + + const newTableSql = addRateUsageStatusAndRestrictions( + replaceCreateTableName(currentSql, 'space_tasks_m163_new') + ); + + // CRITICAL: Disable foreign keys during table recreation to prevent CASCADE + // deletes from wiping child rows when we DROP TABLE space_tasks. + db.exec('PRAGMA foreign_keys = OFF'); + db.exec('BEGIN'); + try { + db.exec(`DROP TABLE IF EXISTS space_tasks_m163_new`); + db.exec(newTableSql); + db.exec( + `INSERT INTO space_tasks_m163_new (${copyColumns}) SELECT ${copyColumns} FROM space_tasks` + ); + db.exec(`DROP TABLE space_tasks`); + db.exec(`ALTER TABLE space_tasks_m163_new RENAME TO space_tasks`); + recreateCompatibleIndexes(db, 'space_tasks', existingIndexDdl); + db.exec('COMMIT'); + } catch (err) { + db.exec('ROLLBACK'); + throw err; + } finally { + db.exec('PRAGMA foreign_keys = ON'); + } +} + +/** + * Widens the `status` CHECK to include `rate_limited` / `usage_limited` and + * injects a nullable `restrictions TEXT` column before the first FOREIGN KEY + * (or before the closing paren if there are none). Preserves every other + * constraint. Throws if the status CHECK is not found. + */ +function addRateUsageStatusAndRestrictions(createSql: string): string { + let statusMatched = false; + let result = createSql.replace( + /CHECK\s*\(\s*status\s+IN\s*\(([^)]*)\)\s*\)/i, + (match, values: string) => { + statusMatched = true; + if (values.includes("'rate_limited'")) { + return match; + } + return `CHECK(status IN ('rate_limited', 'usage_limited', ${values.trim()}))`; + } + ); + if (!statusMatched) { + throw new Error('Migration 163: space_tasks status CHECK constraint not found'); + } + + if (!/\brestrictions\s+TEXT\b/i.test(result)) { + if (/\bFOREIGN\s+KEY\b/i.test(result)) { + result = result.replace(/\bFOREIGN\s+KEY\b/i, 'restrictions TEXT,\n\t\t\t\t\t\tFOREIGN KEY'); + } else { + result = result.replace(/\)\s*$/, 'restrictions TEXT\n\t\t\t\t\t)'); + } + } + return result; +} diff --git a/packages/daemon/tests/unit/1-core/agent/fallback-recovery.test.ts b/packages/daemon/tests/unit/1-core/agent/fallback-recovery.test.ts new file mode 100644 index 000000000..0a8279720 --- /dev/null +++ b/packages/daemon/tests/unit/1-core/agent/fallback-recovery.test.ts @@ -0,0 +1,301 @@ +/** + * Unit tests for the fallback-recovery pure module. + * + * Covers: chain resolution, next-entry selection (skip tried / same-model / + * unavailable), format-agnostic reset-timestamp extraction across locales, + * backoff-ladder progression, cooldown decision, and limit-kind classification. + */ +import { describe, expect, test } from 'bun:test'; +import { + BACKOFF_CAP_MS, + BACKOFF_JITTER, + BACKOFF_LADDER_MS, + MAX_RESET_HORIZON_MS, + RESET_BUFFER_MS, + classifyLimitKind, + computeCooldown, + entryKey, + extractResetTimestamp, + resolveFallbackChain, + selectNextFallback, +} from '../../../../src/lib/agent/fallback-recovery'; +import type { FallbackModelEntry } from '@hyperneo/shared'; + +const A: FallbackModelEntry = { provider: 'anthropic', model: 'claude-sonnet-4-5' }; +const B: FallbackModelEntry = { provider: 'glm', model: 'glm-4.6' }; +const C: FallbackModelEntry = { provider: 'minimax', model: 'abab6.5' }; + +describe('resolveFallbackChain', () => { + test('modelFallbackMap override wins over global list', () => { + const chain = resolveFallbackChain( + 'anthropic', + 'claude-sonnet-4-5', + { 'anthropic/claude-sonnet-4-5': [B, C] }, + [A] + ); + expect(chain).toEqual([B, C]); + }); + + test('global list used when map key absent', () => { + const chain = resolveFallbackChain('glm', 'glm-4.6', { 'anthropic/x': [A] }, [B, C]); + expect(chain).toEqual([B, C]); + }); + + test('empty map value falls through to global', () => { + const chain = resolveFallbackChain( + 'anthropic', + 'claude-sonnet-4-5', + { 'anthropic/claude-sonnet-4-5': [] }, + [B] + ); + expect(chain).toEqual([B]); + }); + + test('both undefined → empty', () => { + expect(resolveFallbackChain('anthropic', 'x', undefined, undefined)).toEqual([]); + }); + + test('returns a defensive copy', () => { + const chain = resolveFallbackChain('anthropic', 'x', undefined, [B, C]); + chain.push(A); + // Re-resolving yields the original order — the caller did not mutate source. + const again = resolveFallbackChain('anthropic', 'x', undefined, [B, C]); + expect(again).toEqual([B, C]); + }); +}); + +describe('entryKey', () => { + test('joins provider and model', () => { + expect(entryKey(A)).toBe('anthropic/claude-sonnet-4-5'); + }); +}); + +describe('selectNextFallback', () => { + test('returns first untried + available entry', () => { + const sel = selectNextFallback([A, B, C], new Set(), () => true); + expect(sel.next).toEqual(A); + expect(sel.exhausted).toBe(false); + expect(sel.skipReason).toBe('none'); + }); + + test('skips tried entries', () => { + const tried = new Set([entryKey(A)]); + const sel = selectNextFallback([A, B, C], tried, () => true); + expect(sel.next).toEqual(B); + }); + + test('skips the current (just-failed) model when caller pre-adds it', () => { + const tried = new Set([entryKey(A)]); + const sel = selectNextFallback([A, B], tried, () => true); + expect(sel.next).toEqual(B); + }); + + test('skips unavailable entries', () => { + const sel = selectNextFallback([A, B, C], new Set(), (e) => e !== A && e !== B); + expect(sel.next).toEqual(C); + }); + + test('reports exhausted with last skip reason', () => { + const tried = new Set([entryKey(A)]); + const sel = selectNextFallback([A, B], tried, () => false); + expect(sel.next).toBeNull(); + expect(sel.exhausted).toBe(true); + expect(sel.skipReason).toBe('unavailable'); + }); + + test('empty chain → exhausted with none', () => { + const sel = selectNextFallback([], new Set(), () => true); + expect(sel.next).toBeNull(); + expect(sel.exhausted).toBe(true); + expect(sel.skipReason).toBe('none'); + }); + + test('advances across successive calls as caller adds keys', () => { + const tried = new Set(); + let sel = selectNextFallback([A, B, C], tried, () => true); + expect(sel.next).toEqual(A); + tried.add(entryKey(sel.next!)); + sel = selectNextFallback([A, B, C], tried, () => true); + expect(sel.next).toEqual(B); + tried.add(entryKey(sel.next!)); + sel = selectNextFallback([A, B, C], tried, () => true); + expect(sel.next).toEqual(C); + tried.add(entryKey(sel.next!)); + sel = selectNextFallback([A, B, C], tried, () => true); + expect(sel.next).toBeNull(); + expect(sel.exhausted).toBe(true); + }); +}); + +describe('extractResetTimestamp', () => { + // `now` fixed well before any test timestamp; future dates chosen so they are + // valid regardless of the host timezone for the local-parse strategy. + const NOW = new Date('2026-01-01T00:00:00Z').getTime(); + + test('ISO-8601 with Z', () => { + const r = extractResetTimestamp('rate limited; retry after 2026-01-01T12:00:00Z', NOW); + expect(r?.strategy).toBe('iso8601'); + expect(r?.resetAtMs).toBe(new Date('2026-01-01T12:00:00Z').getTime()); + }); + + test('ISO-8601 with +08:00 offset', () => { + const r = extractResetTimestamp('resets 2026-01-01T20:00:00+08:00', NOW); + expect(r?.strategy).toBe('iso8601'); + expect(r?.resetAtMs).toBe(new Date('2026-01-01T12:00:00Z').getTime()); + }); + + test('Chinese relay message — YYYY-MM-DD HH:mm:ss parsed as local time', () => { + const msg = + 'Request rejected (429) · [1308][已达到 5 小时的使用上限。您的限额将在 2026-01-02 17:55:10 重置。]'; + const r = extractResetTimestamp(msg, NOW); + expect(r?.strategy).toBe('yyyymmdd-hms'); + // Expected parsed as LOCAL time — compute the same way the impl does so the + // assertion is timezone-independent. + expect(r?.resetAtMs).toBe(new Date('2026-01-02T17:55:10').getTime()); + }); + + test('does not match vendor phrasing — only the digit shape', () => { + // A message with cap wording but NO timestamp token returns null. + const r = extractResetTimestamp('已达到使用上限,请稍后重试', NOW); + expect(r).toBeNull(); + }); + + test('epoch millis (13 digits)', () => { + const target = NOW + 3 * 60 * 60 * 1000; + const r = extractResetTimestamp(`retry after ${target}`, NOW); + expect(r?.strategy).toBe('epoch-millis'); + expect(r?.resetAtMs).toBe(target); + }); + + test('epoch seconds (10 digits)', () => { + const target = NOW + 3 * 60 * 60 * 1000; + const r = extractResetTimestamp(`retry-after: ${Math.floor(target / 1000)}`, NOW); + expect(r?.strategy).toBe('epoch-seconds'); + expect(r?.resetAtMs).toBe(target); + }); + + test('past date → null', () => { + expect(extractResetTimestamp('resets 2020-01-01 00:00:00', NOW)).toBeNull(); + }); + + test('far-future date (> 7 days) → null', () => { + expect(extractResetTimestamp('resets 2099-12-31 23:59:59', NOW)).toBeNull(); + }); + + test('no timestamp token → null', () => { + expect(extractResetTimestamp('429 rate limit exceeded', NOW)).toBeNull(); + }); + + test('4-digit relay code [1308] does not match epoch-seconds', () => { + expect(extractResetTimestamp('error [1308] something', NOW)).toBeNull(); + }); + + test('ISO with offset is preferred over local interpretation', () => { + // 2026-01-01T12:00:00Z and a same-string local reading would differ; the + // explicit-Z match must win and produce the UTC value. + const r = extractResetTimestamp('reset 2026-01-01T12:00:00Z', NOW); + expect(r?.resetAtMs).toBe(new Date('2026-01-01T12:00:00Z').getTime()); + }); + + test('horizon boundary: just under 7 days accepted, over rejected', () => { + const within = NOW + MAX_RESET_HORIZON_MS - 60_000; + expect(extractResetTimestamp(`r ${within}`, NOW)?.resetAtMs).toBe(within); + const beyond = NOW + MAX_RESET_HORIZON_MS + 60_000; + expect(extractResetTimestamp(`r ${beyond}`, NOW)).toBeNull(); + }); +}); + +describe('computeCooldown', () => { + const NOW = new Date('2026-01-01T00:00:00Z').getTime(); + + test('parsed-reset → free wait at reset time + buffer', () => { + const reset = NOW + 5 * 60 * 60 * 1000; // 5h out + const d = computeCooldown(`resets 2026-01-01T05:00:00Z`, 0, NOW); + expect(d.reason).toBe('parsed-reset'); + expect(d.freeWait).toBe(true); + expect(d.ladderIndex).toBe(-1); + expect(d.delayMs).toBe(5 * 60 * 60 * 1000 + RESET_BUFFER_MS); + expect(d.retryAtMs).toBe(reset + RESET_BUFFER_MS); + expect(d.reset?.strategy).toBe('iso8601'); + }); + + test('parsed-reset in near future still adds buffer (never negative)', () => { + // 5s in the future — delay = remaining (5s) + buffer (30s), never negative. + const d = computeCooldown(`resets 2026-01-01T00:00:05Z`, 0, NOW); + expect(d.delayMs).toBe(5000 + RESET_BUFFER_MS); + }); + + test('backoff ladder progression for steps 0..4', () => { + const noJitter = () => 0; + const expected = BACKOFF_LADDER_MS; + for (let i = 0; i < expected.length; i++) { + const d = computeCooldown('429 rate limit', i, NOW, noJitter); + expect(d.reason).toBe('backoff-ladder'); + expect(d.freeWait).toBe(false); + expect(d.ladderIndex).toBe(i); + expect(d.delayMs).toBe(Math.min(expected[i], BACKOFF_CAP_MS)); + } + }); + + test('ladder caps at last entry for counts beyond the array', () => { + const noJitter = () => 0; + const d = computeCooldown('429', 99, NOW, noJitter); + expect(d.ladderIndex).toBe(BACKOFF_LADDER_MS.length - 1); + expect(d.delayMs).toBe( + Math.min(BACKOFF_LADDER_MS[BACKOFF_LADDER_MS.length - 1], BACKOFF_CAP_MS) + ); + }); + + test('delay never exceeds BACKOFF_CAP_MS even with max positive jitter', () => { + const maxPos = () => 1; + for (let i = 0; i < 10; i++) { + const d = computeCooldown('429', i, NOW, maxPos); + expect(d.delayMs).toBeLessThanOrEqual(BACKOFF_CAP_MS); + } + }); + + test('jitter stays within ±BACKOFF_JITTER of the (capped) base', () => { + const base = Math.min(BACKOFF_LADDER_MS[0], BACKOFF_CAP_MS); + let min = Infinity; + let max = -Infinity; + for (let i = 0; i < 1000; i++) { + const d = computeCooldown('429', 0, NOW); + min = Math.min(min, d.delayMs); + max = Math.max(max, d.delayMs); + } + expect(min).toBeGreaterThanOrEqual(Math.round(base * (1 - BACKOFF_JITTER))); + expect(max).toBeLessThanOrEqual(Math.round(base * (1 + BACKOFF_JITTER))); + }); + + test('backoff floor 1min even with max negative jitter', () => { + const maxNeg = () => -1; + const d = computeCooldown('429', 0, NOW, maxNeg); + expect(d.delayMs).toBeGreaterThanOrEqual(60_000); + }); +}); + +describe('classifyLimitKind', () => { + test('parsed-reset → usage_limit', () => { + const d = computeCooldown( + 'resets 2026-01-01T05:00:00Z', + 0, + new Date('2026-01-01T00:00:00Z').getTime() + ); + expect(classifyLimitKind('some message', d)).toBe('usage_limit'); + }); + + test('backoff with usage keyword → usage_limit', () => { + const d = { reason: 'backoff-ladder' } as ReturnType; + expect(classifyLimitKind('usage limit reached', d)).toBe('usage_limit'); + }); + + test('backoff with Chinese 上限 keyword → usage_limit', () => { + const d = { reason: 'backoff-ladder' } as ReturnType; + expect(classifyLimitKind('已达到使用上限', d)).toBe('usage_limit'); + }); + + test('backoff with no cap signal → rate_limit', () => { + const d = { reason: 'backoff-ladder' } as ReturnType; + expect(classifyLimitKind('429 too many requests', d)).toBe('rate_limit'); + }); +}); diff --git a/packages/daemon/tests/unit/1-core/agent/rate-limit-watchdog.test.ts b/packages/daemon/tests/unit/1-core/agent/rate-limit-watchdog.test.ts index 32fbab0a8..f05d999c7 100644 --- a/packages/daemon/tests/unit/1-core/agent/rate-limit-watchdog.test.ts +++ b/packages/daemon/tests/unit/1-core/agent/rate-limit-watchdog.test.ts @@ -1,20 +1,25 @@ /** * RateLimitWatchdog Tests * - * Tests for the rate limit auto-retry watchdog: - * - Schedule retry with cooldown - * - Cancel pending retry - * - Max retries exceeded - * - RetryNow bypasses cooldown - * - Reset clears state - * - Null message guard + * Two-phase recovery: + * - Phase A: immediate fallback-model switch (free, no retryCount bump). + * - Phase B: cooldown at a parsed reset time or on the backoff ladder. + * + * Plus episode tracking (tried entries), pause/resume surfacing, and the + * preserved cancel/retryNow/reset/destroy semantics. */ import { describe, expect, it, beforeEach, mock } from 'bun:test'; -import { RateLimitWatchdog } from '../../../../src/lib/agent/rate-limit-watchdog'; +import { + RateLimitWatchdog, + type RateLimitWatchdogDeps, +} from '../../../../src/lib/agent/rate-limit-watchdog'; +import { RESET_BUFFER_MS } from '../../../../src/lib/agent/fallback-recovery'; import type { ProcessingStateManager } from '../../../../src/lib/agent/processing-state-manager'; +import type { FallbackModelEntry } from '@hyperneo/shared'; + +type Msg = { uuid: string; content: string }; -// Helper to create a mock ProcessingStateManager function createMockStateManager(): ProcessingStateManager { return { getState: mock(() => ({ status: 'idle' })), @@ -38,152 +43,366 @@ function createMockStateManager(): ProcessingStateManager { } as unknown as ProcessingStateManager; } +interface MockDepsOptions { + chain?: FallbackModelEntry[]; + current?: { provider: string; model: string }; + available?: (e: FallbackModelEntry) => boolean; + switchSucceeds?: boolean | ((e: FallbackModelEntry) => boolean); +} + +function createMockDeps(opts: MockDepsOptions = {}): { + deps: RateLimitWatchdogDeps; + setModel: (provider: string, model: string) => void; + switchAndRetry: ReturnType; + notifyPause: ReturnType; + notifyResume: ReturnType; +} { + let current = opts.current ?? { provider: 'anthropic', model: 'claude-sonnet-4-5' }; + const switchAndRetry = mock(async (_msg: Msg | null, entry: FallbackModelEntry) => { + // Simulate the session's model actually changing on a successful switch. + current = { provider: entry.provider, model: entry.model }; + return typeof opts.switchSucceeds === 'function' + ? opts.switchSucceeds(entry) + : (opts.switchSucceeds ?? true); + }); + const notifyPause = mock((_payload: unknown) => {}); + const notifyResume = mock(() => {}); + const deps: RateLimitWatchdogDeps = { + getCurrentModel: () => current, + resolveChain: () => opts.chain ?? [], + isEntryAvailable: async (e) => (opts.available ? opts.available(e) : true), + switchAndRetry, + notifyPause, + notifyResume, + }; + return { + deps, + setModel: (p, m) => (current = { provider: p, model: m }), + switchAndRetry, + notifyPause, + notifyResume, + }; +} + +const flush = () => new Promise((resolve) => setTimeout(resolve, 0)); + describe('RateLimitWatchdog', () => { - let watchdog: RateLimitWatchdog; let stateManager: ProcessingStateManager; - let retryCallback: ReturnType; beforeEach(() => { stateManager = createMockStateManager(); - retryCallback = mock(async () => {}); - watchdog = new RateLimitWatchdog('test-session', stateManager, { - cooldownMs: 100, // Fast cooldown for tests - maxAutoRetries: 3, - }); - watchdog.setRetryCallback(retryCallback); }); describe('getState', () => { it('returns idle state initially', () => { + const { deps } = createMockDeps(); + const watchdog = new RateLimitWatchdog('s', stateManager, deps, { maxAutoRetries: 3 }); const state = watchdog.getState(); expect(state.status).toBe('idle'); expect(state.retryCount).toBe(0); expect(state.maxRetries).toBe(3); expect(state.retryAt).toBeNull(); expect(state.lastUserMessage).toBeNull(); + expect(state.triedEntries).toEqual([]); + expect(state.fallbackPending).toBe(false); }); }); - describe('scheduleRetry', () => { - it('schedules a retry and sets rate_limit_cooldown state', async () => { + describe('Phase B — cooldown (chain empty)', () => { + it('schedules a backoff-ladder cooldown and sets rate_limit_cooldown state', async () => { + const { deps, notifyPause } = createMockDeps({ chain: [] }); + const watchdog = new RateLimitWatchdog('s', stateManager, deps, { maxAutoRetries: 3 }); const result = await watchdog.scheduleRetry('429 rate limit', { - uuid: 'msg-1', - content: 'hello', + uuid: 'm1', + content: 'hi', }); - expect(result).toBe(true); expect(stateManager.setRateLimitCooldown).toHaveBeenCalledTimes(1); - - const state = watchdog.getState(); - expect(state.status).toBe('cooldown'); - expect(state.retryCount).toBe(1); - expect(state.lastUserMessage).toEqual({ uuid: 'msg-1', content: 'hello' }); - expect(state.retryAt).toBeGreaterThan(Date.now() - 1000); + expect(watchdog.getState().status).toBe('cooldown'); + expect(watchdog.getState().retryCount).toBe(1); + expect(watchdog.getState().lastUserMessage).toEqual({ uuid: 'm1', content: 'hi' }); + // Pause surfaced as a (transient) rate_limit. + expect(notifyPause).toHaveBeenCalledTimes(1); + expect(notifyPause.mock.calls[0][0]).toMatchObject({ kind: 'rate_limit' }); + watchdog.cancel(); }); - it('increments retryCount on subsequent calls', async () => { - await watchdog.scheduleRetry('429', { uuid: 'msg-1', content: 'hello' }); - watchdog.cancel(); // Cancel before scheduling next + it('schedules a cooldown at a parsed reset time without bumping retryCount (free wait)', async () => { + const reset = Date.now() + 5 * 60 * 60 * 1000; // 5h out + const { deps, notifyPause } = createMockDeps({ chain: [] }); + const watchdog = new RateLimitWatchdog('s', stateManager, deps, { maxAutoRetries: 3 }); + const result = await watchdog.scheduleRetry( + `限额将在 ${new Date(reset) + .toISOString() + .replace('T', ' ') + .replace(/\.\d+Z$/, '')} 重置`, + { uuid: 'm1', content: 'hi' } + ); + expect(result).toBe(true); + expect(watchdog.getState().retryCount).toBe(0); // free wait + const cooldownArgs = (stateManager.setRateLimitCooldown as ReturnType).mock + .calls[0][0]; + expect(cooldownArgs.retryAt).toBeGreaterThan(Date.now()); + // Usage cap (reset known) → usage_limit. + expect(notifyPause.mock.calls[0][0]).toMatchObject({ kind: 'usage_limit' }); + watchdog.cancel(); + }); - await watchdog.scheduleRetry('429', { uuid: 'msg-2', content: 'world' }); + it('returns false when lastUserMessage is null', async () => { + const { deps } = createMockDeps({ chain: [] }); + const watchdog = new RateLimitWatchdog('s', stateManager, deps); + const result = await watchdog.scheduleRetry('429', null); + expect(result).toBe(false); + expect(watchdog.getState().status).toBe('idle'); + expect(watchdog.isPending()).toBe(false); + }); + it('returns false when max cooldown retries exceeded', async () => { + const { deps } = createMockDeps({ chain: [] }); + const watchdog = new RateLimitWatchdog('s', stateManager, deps, { maxAutoRetries: 2 }); + for (let i = 0; i < 2; i++) { + await watchdog.scheduleRetry('429', { uuid: `m${i}`, content: 'x' }); + watchdog.cancel(); + } + // 3rd cooldown attempt (retryCount already 2) → false. + const result = await watchdog.scheduleRetry('429', { uuid: 'm3', content: 'x' }); + expect(result).toBe(false); expect(watchdog.getState().retryCount).toBe(2); }); - it('returns false when max retries exceeded', async () => { - for (let i = 0; i < 3; i++) { - await watchdog.scheduleRetry('429', { uuid: `msg-${i}`, content: `test-${i}` }); + it('a parsed-reset wait bypasses the maxAutoRetries budget (free wait)', async () => { + const { deps } = createMockDeps({ chain: [] }); + const watchdog = new RateLimitWatchdog('s', stateManager, deps, { maxAutoRetries: 2 }); + // Burn the budget on backoff (no-timestamp) cooldowns. + for (let i = 0; i < 2; i++) { + await watchdog.scheduleRetry('429', { uuid: `m${i}`, content: 'x' }); watchdog.cancel(); } + expect(watchdog.getState().retryCount).toBe(2); // at the budget + + // A parseable reset wait is free: returns true AND does not bump retryCount + // even though we are already at maxAutoRetries. + const resetIso = new Date(Date.now() + 60 * 60 * 1000).toISOString(); + const result = await watchdog.scheduleRetry(`resets ${resetIso}`, { + uuid: 'm-reset', + content: 'x', + }); + expect(result).toBe(true); + expect(watchdog.getState().retryCount).toBe(2); // unchanged — free wait + watchdog.cancel(); + }); - // 4th attempt should fail - const result = await watchdog.scheduleRetry('429', { uuid: 'msg-4', content: 'nope' }); - expect(result).toBe(false); - expect(watchdog.getState().retryCount).toBe(3); + it('notifyPause resetAt carries the 30s buffer over the parsed reset time', async () => { + const { deps, notifyPause } = createMockDeps({ chain: [] }); + const watchdog = new RateLimitWatchdog('s', stateManager, deps); + // Use a whole-second reset so the second-granularity ISO parse is exact. + const resetAt = Math.floor((Date.now() + 60 * 60 * 1000) / 1000) * 1000; + const resetIso = new Date(resetAt).toISOString(); + await watchdog.scheduleRetry(`resets ${resetIso}`, { uuid: 'm1', content: 'x' }); + const payload = notifyPause.mock.calls[0][0] as { resetAt?: number }; + expect(payload.resetAt).toBe(resetAt + RESET_BUFFER_MS); + watchdog.cancel(); }); - it('returns false when lastUserMessage is null', async () => { - const result = await watchdog.scheduleRetry('429', null); - expect(result).toBe(false); - expect(watchdog.getState().status).toBe('idle'); - expect(watchdog.isPending()).toBe(false); + it('increments retryCount only on (non-free) backoff steps', async () => { + const { deps } = createMockDeps({ chain: [] }); + const watchdog = new RateLimitWatchdog('s', stateManager, deps, { maxAutoRetries: 5 }); + await watchdog.scheduleRetry('429', { uuid: 'm1', content: 'x' }); + expect(watchdog.getState().retryCount).toBe(1); + watchdog.cancel(); + await watchdog.scheduleRetry('429', { uuid: 'm2', content: 'x' }); + expect(watchdog.getState().retryCount).toBe(2); + watchdog.cancel(); + }); + }); + + describe('Phase A — immediate fallback switch', () => { + const A: FallbackModelEntry = { provider: 'glm', model: 'glm-4.6' }; + const B: FallbackModelEntry = { provider: 'minimax', model: 'abab6.5' }; + + it('switches to the next available fallback and retries without a cooldown', async () => { + const { deps, switchAndRetry } = createMockDeps({ + current: { provider: 'anthropic', model: 'sonnet' }, + chain: [A, B], + }); + const watchdog = new RateLimitWatchdog('s', stateManager, deps, { maxAutoRetries: 3 }); + const result = await watchdog.scheduleRetry('429', { uuid: 'm1', content: 'hi' }); + expect(result).toBe(true); + await flush(); + expect(switchAndRetry).toHaveBeenCalledTimes(1); + expect(switchAndRetry.mock.calls[0][1]).toEqual(A); + // No cooldown scheduled for a fallback switch. + expect(stateManager.setRateLimitCooldown).not.toHaveBeenCalled(); + expect(watchdog.getState().retryCount).toBe(0); + expect(watchdog.getState().triedEntries).toContain('anthropic/sonnet'); }); - it('fires the retry callback after cooldown', async () => { - await watchdog.scheduleRetry('429', { uuid: 'msg-1', content: 'hello' }); + it('skips unavailable entries', async () => { + const { deps, switchAndRetry } = createMockDeps({ + current: { provider: 'anthropic', model: 'sonnet' }, + chain: [A, B], + available: (e) => e !== A, + }); + const watchdog = new RateLimitWatchdog('s', stateManager, deps); + await watchdog.scheduleRetry('429', { uuid: 'm1', content: 'hi' }); + await flush(); + expect(switchAndRetry.mock.calls[0][1]).toEqual(B); + }); - // Wait for cooldown (100ms) + buffer - await new Promise((resolve) => setTimeout(resolve, 200)); + it('advances to the next entry when a switch fails, without bumping retryCount', async () => { + let calls = 0; + const { deps, switchAndRetry } = createMockDeps({ + current: { provider: 'anthropic', model: 'sonnet' }, + chain: [A, B], + switchSucceeds: (e) => e !== A, // A fails, B succeeds + }); + const watchdog = new RateLimitWatchdog('s', stateManager, deps); + await watchdog.scheduleRetry('429', { uuid: 'm1', content: 'hi' }); + await flush(); + expect(switchAndRetry).toHaveBeenCalledTimes(2); // A (fail) then B (success) + expect(switchAndRetry.mock.calls[0][1]).toEqual(A); + expect(switchAndRetry.mock.calls[1][1]).toEqual(B); + expect(watchdog.getState().retryCount).toBe(0); // switches are free + expect(stateManager.setRateLimitCooldown).not.toHaveBeenCalled(); + }); - expect(retryCallback).toHaveBeenCalledTimes(1); - expect(retryCallback).toHaveBeenCalledWith({ uuid: 'msg-1', content: 'hello' }); - expect(watchdog.isPending()).toBe(false); + it('falls through to a cooldown when every switch fails and the chain is exhausted', async () => { + const { deps } = createMockDeps({ + current: { provider: 'anthropic', model: 'sonnet' }, + chain: [A], + switchSucceeds: () => false, + }); + const watchdog = new RateLimitWatchdog('s', stateManager, deps, { maxAutoRetries: 3 }); + await watchdog.scheduleRetry('429', { uuid: 'm1', content: 'hi' }); + await flush(); + expect(stateManager.setRateLimitCooldown).toHaveBeenCalledTimes(1); + watchdog.cancel(); }); }); - describe('cancel', () => { - it('cancels a pending retry', async () => { - await watchdog.scheduleRetry('429', { uuid: 'msg-1', content: 'hello' }); - expect(watchdog.isPending()).toBe(true); + describe('episode tracking across repeated 429', () => { + const A: FallbackModelEntry = { provider: 'glm', model: 'glm-a' }; + const B: FallbackModelEntry = { provider: 'minimax', model: 'mm-b' }; + it('walks A → B → cooldown as successive models 429', async () => { + const { deps, switchAndRetry, setModel } = createMockDeps({ + current: { provider: 'anthropic', model: 'sonnet' }, + chain: [A, B], + }); + const watchdog = new RateLimitWatchdog('s', stateManager, deps, { maxAutoRetries: 3 }); + + // 1st 429 on anthropic/sonnet → switch to A. + await watchdog.scheduleRetry('429', { uuid: 'm1', content: 'hi' }); + await flush(); + expect(switchAndRetry.mock.calls[0][1]).toEqual(A); + // setModel simulates the session now running A. + setModel(A.provider, A.model); + + // 2nd 429 on A → switch to B. + await watchdog.scheduleRetry('429', { uuid: 'm1', content: 'hi' }); + await flush(); + expect(switchAndRetry.mock.calls[1][1]).toEqual(B); + setModel(B.provider, B.model); + + // 3rd 429 on B → chain exhausted → cooldown (retryCount 1). + await watchdog.scheduleRetry('429', { uuid: 'm1', content: 'hi' }); + expect(stateManager.setRateLimitCooldown).toHaveBeenCalledTimes(1); + expect(watchdog.getState().retryCount).toBe(1); watchdog.cancel(); - - expect(watchdog.isPending()).toBe(false); - expect(watchdog.getState().status).toBe('idle'); }); - it('does not fire callback after cancellation', async () => { - await watchdog.scheduleRetry('429', { uuid: 'msg-1', content: 'hello' }); + it('maxAutoRetries counts only cooldowns, not fallback switches', async () => { + const { deps, switchAndRetry, setModel } = createMockDeps({ + current: { provider: 'anthropic', model: 'sonnet' }, + chain: [A, B], + }); + const watchdog = new RateLimitWatchdog('s', stateManager, deps, { maxAutoRetries: 2 }); + // Burn the two fallback switches. + await watchdog.scheduleRetry('429', { uuid: 'm1', content: 'hi' }); + await flush(); + setModel(A.provider, A.model); + await watchdog.scheduleRetry('429', { uuid: 'm1', content: 'hi' }); + await flush(); + setModel(B.provider, B.model); + // Two cooldown steps allowed. + await watchdog.scheduleRetry('429', { uuid: 'm1', content: 'hi' }); + expect(watchdog.getState().retryCount).toBe(1); watchdog.cancel(); - - // Wait past cooldown - await new Promise((resolve) => setTimeout(resolve, 200)); - - expect(retryCallback).not.toHaveBeenCalled(); + await watchdog.scheduleRetry('429', { uuid: 'm1', content: 'hi' }); + expect(watchdog.getState().retryCount).toBe(2); + watchdog.cancel(); + // 3rd cooldown → exhausted → false (despite only 2 prior switchAndRetry calls). + const result = await watchdog.scheduleRetry('429', { uuid: 'm1', content: 'hi' }); + expect(result).toBe(false); + expect(switchAndRetry).toHaveBeenCalledTimes(2); }); }); - describe('retryNow', () => { - it('immediately fires the retry callback', async () => { - await watchdog.scheduleRetry('429', { uuid: 'msg-1', content: 'hello' }); + describe('cancel / retryNow / reset / destroy', () => { + it('cancel clears a pending cooldown', async () => { + const { deps } = createMockDeps({ chain: [] }); + const watchdog = new RateLimitWatchdog('s', stateManager, deps); + await watchdog.scheduleRetry('429', { uuid: 'm1', content: 'hi' }); + expect(watchdog.isPending()).toBe(true); + watchdog.cancel(); + expect(watchdog.isPending()).toBe(false); + }); + it('retryNow fires the callback and notifies resume', async () => { + const { deps, notifyResume } = createMockDeps({ chain: [] }); + const retryCallback = mock(async () => {}); + const watchdog = new RateLimitWatchdog('s', stateManager, deps); + watchdog.setRetryCallback(retryCallback); + await watchdog.scheduleRetry('429', { uuid: 'm1', content: 'hi' }); const result = watchdog.retryNow(); - expect(result).toBe(true); expect(retryCallback).toHaveBeenCalledTimes(1); - expect(retryCallback).toHaveBeenCalledWith({ uuid: 'msg-1', content: 'hello' }); - expect(watchdog.isPending()).toBe(false); + expect(retryCallback.mock.calls[0][0]).toEqual({ uuid: 'm1', content: 'hi' }); + expect(retryCallback.mock.calls[0][1]).toBeUndefined(); // no switchTo on a cooldown retry + expect(notifyResume).toHaveBeenCalled(); }); - it('returns false if no retry is pending', () => { - const result = watchdog.retryNow(); - expect(result).toBe(false); + it('retryNow returns false during a fallback-pending switch', async () => { + const A: FallbackModelEntry = { provider: 'glm', model: 'glm-4.6' }; + // Make switchAndRetry hang so fallbackPending stays true when we call retryNow. + const { deps } = createMockDeps({ + current: { provider: 'anthropic', model: 'sonnet' }, + chain: [A], + }); + // Override switchAndRetry to never resolve within the test window. + deps.switchAndRetry = () => new Promise(() => {}); + const watchdog = new RateLimitWatchdog('s', stateManager, deps); + await watchdog.scheduleRetry('429', { uuid: 'm1', content: 'hi' }); + expect(watchdog.getState().fallbackPending).toBe(true); + expect(watchdog.retryNow()).toBe(false); }); - }); - describe('reset', () => { - it('clears all state', async () => { - await watchdog.scheduleRetry('429', { uuid: 'msg-1', content: 'hello' }); + it('reset clears the episode (tried entries + chain + retryCount)', async () => { + const A: FallbackModelEntry = { provider: 'glm', model: 'glm-4.6' }; + const { deps } = createMockDeps({ + current: { provider: 'anthropic', model: 'sonnet' }, + chain: [A], + }); + const watchdog = new RateLimitWatchdog('s', stateManager, deps); + await watchdog.scheduleRetry('429', { uuid: 'm1', content: 'hi' }); + await flush(); watchdog.reset(); - expect(watchdog.getState().status).toBe('idle'); expect(watchdog.getState().retryCount).toBe(0); + expect(watchdog.getState().triedEntries).toEqual([]); + expect(watchdog.getState().fallbackChain).toBeNull(); expect(watchdog.getState().lastUserMessage).toBeNull(); - expect(watchdog.isPending()).toBe(false); }); - }); - describe('destroy', () => { - it('cancels timers and clears callback', async () => { - await watchdog.scheduleRetry('429', { uuid: 'msg-1', content: 'hello' }); + it('destroy cancels timers and clears the callback', async () => { + const { deps } = createMockDeps({ chain: [] }); + const retryCallback = mock(async () => {}); + const watchdog = new RateLimitWatchdog('s', stateManager, deps); + watchdog.setRetryCallback(retryCallback); + await watchdog.scheduleRetry('429', { uuid: 'm1', content: 'hi' }); watchdog.destroy(); - - // Wait past cooldown - await new Promise((resolve) => setTimeout(resolve, 200)); - - // Callback was cleared, so even though timer might have been set, - // destroy sets retryCallback to null - expect(retryCallback).not.toHaveBeenCalled(); + // Callback cleared: even if a retry were triggered, nothing fires. + expect(() => watchdog.retryNow()).not.toThrow(); }); }); }); diff --git a/packages/daemon/tests/unit/4-space-storage/storage/space-task-repository.test.ts b/packages/daemon/tests/unit/4-space-storage/storage/space-task-repository.test.ts index ef443a895..d6b85b0f7 100644 --- a/packages/daemon/tests/unit/4-space-storage/storage/space-task-repository.test.ts +++ b/packages/daemon/tests/unit/4-space-storage/storage/space-task-repository.test.ts @@ -170,6 +170,90 @@ describe('SpaceTaskRepository', () => { }); }); + describe('restrictions (rate/usage-limit pause)', () => { + it('persists and reads back a restrictions blob with usage_limited status', () => { + const task = repo.createTask({ spaceId, title: 'T', description: '' }); + const resetAt = Date.now() + 60 * 60 * 1000; + const updated = repo.updateTask(task.id, { + status: 'usage_limited', + restrictions: { + type: 'usage_limit', + limit: 'parsed-reset', + resetAt, + sessionRole: 'worker', + }, + }); + expect(updated?.status).toBe('usage_limited'); + expect(updated?.restrictions).toEqual({ + type: 'usage_limit', + limit: 'parsed-reset', + resetAt, + sessionRole: 'worker', + }); + + // Re-read from DB (exercises the row mapper, not just the update return). + const reread = repo.getTask(task.id); + expect(reread?.status).toBe('usage_limited'); + expect(reread?.restrictions).toEqual({ + type: 'usage_limit', + limit: 'parsed-reset', + resetAt, + sessionRole: 'worker', + }); + }); + + it('clears restrictions when set to null on resume', () => { + const task = repo.createTask({ spaceId, title: 'T', description: '' }); + repo.updateTask(task.id, { + status: 'rate_limited', + restrictions: { + type: 'rate_limit', + limit: 'backoff-ladder', + resetAt: Date.now() + 60000, + sessionRole: 'worker', + }, + }); + const resumed = repo.updateTask(task.id, { status: 'in_progress', restrictions: null }); + expect(resumed?.status).toBe('in_progress'); + expect(resumed?.restrictions).toBeNull(); + }); + + it('defaults restrictions to null for an ordinary task', () => { + const task = repo.createTask({ spaceId, title: 'T', description: '' }); + expect(repo.getTask(task.id)?.restrictions).toBeNull(); + }); + + it('clears restrictions on a manual transition out of the paused status', () => { + const task = repo.createTask({ spaceId, title: 'T', description: '' }); + repo.updateTask(task.id, { + status: 'usage_limited', + restrictions: { + type: 'usage_limit', + limit: 'parsed-reset', + resetAt: Date.now() + 60000, + sessionRole: 'worker', + }, + }); + // Manual cancel without explicitly clearing restrictions. + const cancelled = repo.updateTask(task.id, { status: 'cancelled' }); + expect(cancelled?.status).toBe('cancelled'); + expect(cancelled?.restrictions).toBeNull(); + }); + }); + + describe('listRateLimitedBySpace', () => { + it('returns only rate_limited / usage_limited tasks', () => { + const t1 = repo.createTask({ spaceId, title: 'paused-1', description: '' }); + repo.updateTask(t1.id, { status: 'usage_limited' }); + const t2 = repo.createTask({ spaceId, title: 'paused-2', description: '' }); + repo.updateTask(t2.id, { status: 'rate_limited' }); + repo.createTask({ spaceId, title: 'active', description: '' }); // in_progress + + const paused = repo.listRateLimitedBySpace(spaceId); + expect(paused.map((t) => t.id).sort()).toEqual([t1.id, t2.id].sort()); + }); + }); + describe('getTask', () => { it('returns task by ID', () => { const created = repo.createTask({ spaceId, title: 'T', description: '' }); diff --git a/packages/daemon/tests/unit/5-space/runtime/space-runtime-tick-loop.test.ts b/packages/daemon/tests/unit/5-space/runtime/space-runtime-tick-loop.test.ts index 0b34b7ffb..259d1f03c 100644 --- a/packages/daemon/tests/unit/5-space/runtime/space-runtime-tick-loop.test.ts +++ b/packages/daemon/tests/unit/5-space/runtime/space-runtime-tick-loop.test.ts @@ -1986,4 +1986,61 @@ describe('SpaceRuntime — tick loop correctness', () => { expect(rehydrateCount).toBe(1); }); }); + + // ------------------------------------------------------------------------- + // Auto-resume of rate/usage-limited tasks across daemon restarts + // (recoverRateLimitedTasks). The in-memory watchdog cooldown does not survive + // a restart; this tick sweep is driven off the persisted restrictions.resetAt. + // ------------------------------------------------------------------------- + describe('auto-resume of rate/usage-limited tasks', () => { + test('restores a paused task whose resetAt has passed (cross-restart backstop)', async () => { + const tam = makeMockTaskAgentManager(taskRepo, nodeExecutionRepo); + const rt = new SpaceRuntime(buildConfig(tam)); + + const task = taskRepo.createTask({ + spaceId: SPACE_ID, + title: 'capped', + description: '', + }); + taskRepo.updateTask(task.id, { status: 'usage_limited' }); + taskRepo.updateTask(task.id, { + status: 'usage_limited', + restrictions: { + type: 'usage_limit', + limit: 'parsed-reset', + resetAt: Date.now() - 60_000, // reset already passed + sessionRole: 'worker', + }, + }); + + await rt.executeTick(); + + const restored = taskRepo.getTask(task.id)!; + expect(restored.status).toBe('in_progress'); + expect(restored.restrictions).toBeNull(); + }); + + test('leaves a paused task whose resetAt is still in the future', async () => { + const tam = makeMockTaskAgentManager(taskRepo, nodeExecutionRepo); + const rt = new SpaceRuntime(buildConfig(tam)); + + const task = taskRepo.createTask({ spaceId: SPACE_ID, title: 'capped', description: '' }); + const futureReset = Date.now() + 60 * 60 * 1000; + taskRepo.updateTask(task.id, { + status: 'rate_limited', + restrictions: { + type: 'rate_limit', + limit: 'backoff-ladder', + resetAt: futureReset, + sessionRole: 'worker', + }, + }); + + await rt.executeTick(); + + const still = taskRepo.getTask(task.id)!; + expect(still.status).toBe('rate_limited'); + expect(still.restrictions?.resetAt).toBe(futureReset); + }); + }); }); diff --git a/packages/daemon/tests/unit/5-space/runtime/task-agent-rate-limit-listener.test.ts b/packages/daemon/tests/unit/5-space/runtime/task-agent-rate-limit-listener.test.ts new file mode 100644 index 000000000..40b292641 --- /dev/null +++ b/packages/daemon/tests/unit/5-space/runtime/task-agent-rate-limit-listener.test.ts @@ -0,0 +1,174 @@ +/** + * TaskAgentManager rate-limit pause/resume listener tests (Part C). + * + * Verifies that a worker session hitting a rate/usage cap (no fallback left) + * surfaces a paused task status with a resume-at restriction, and auto-resumes + * (in_progress + restrictions cleared) when the limit lifts. + * + * The session→task link is the in-memory `subSessions` map seeded directly; the + * listener logic itself (findParentTaskIdForSubSession + updateTask) is what's + * under test, exercised through the real internalEventBus. + */ +import { describe, expect, it, beforeEach, afterEach } from 'bun:test'; +import { Database } from 'bun:sqlite'; +import { TaskAgentManager } from '../../../../src/lib/space/runtime/task-agent-manager'; +import type { TaskAgentManagerConfig } from '../../../../src/lib/space/runtime/task-agent-manager'; +import { SpaceRepository } from '../../../../src/storage/repositories/space-repository'; +import { SpaceTaskRepository } from '../../../../src/storage/repositories/space-task-repository'; +import { createDaemonInternalEventBus } from '../../../../src/lib/internal-event-bus'; +import { createSpaceTables } from '../../helpers/space-test-db'; + +const flush = () => new Promise((resolve) => setTimeout(resolve, 0)); + +describe('TaskAgentManager rate-limit pause/resume listener', () => { + let db: Database; + let taskRepo: SpaceTaskRepository; + let bus: ReturnType; + let manager: TaskAgentManager; + let taskId: string; + const subSessionId = 'worker-session-1'; + + beforeEach(() => { + db = new Database(':memory:'); + createSpaceTables(db); + const spaceRepo = new SpaceRepository( + db as unknown as Parameters[0] + ); + const space = spaceRepo.createSpace({ workspacePath: '/w', slug: 's', name: 'S' }); + taskRepo = new SpaceTaskRepository( + db as unknown as Parameters[0] + ); + const task = taskRepo.createTask({ spaceId: space.id, title: 'T', description: '' }); + taskId = task.id; + taskRepo.updateTask(taskId, { status: 'in_progress' }); + + bus = createDaemonInternalEventBus(); + + // Minimal config: only db / taskRepo / internalEventBus are touched by the + // listener and the constructor. The rest are stubbed. + const config = { + db: { getDatabase: () => db }, + taskRepo, + internalEventBus: bus, + } as unknown as TaskAgentManagerConfig; + manager = new TaskAgentManager(config); + + // Seed the in-memory session→task map the listener resolves through. + const subSessions = (manager as unknown as { subSessions: Map> }) + .subSessions; + subSessions.set(taskId, new Map()); + subSessions.get(taskId)!.set(subSessionId, { id: subSessionId }); + }); + + afterEach(() => { + db.close(); + }); + + it('marks the task usage_limited with a resume-at restriction on pause (usage cap)', async () => { + const resetAt = Date.now() + 5 * 60 * 60 * 1000; + bus.publish('session.rate_limit_pause', { + sessionId: subSessionId, + kind: 'usage_limit', + resetAt, + reason: 'parsed-reset', + }); + await flush(); + + const task = taskRepo.getTask(taskId); + expect(task?.status).toBe('usage_limited'); + expect(task?.restrictions).toMatchObject({ + type: 'usage_limit', + resetAt, + sessionRole: 'worker', + }); + }); + + it('marks the task rate_limited on pause (transient)', async () => { + bus.publish('session.rate_limit_pause', { + sessionId: subSessionId, + kind: 'rate_limit', + reason: 'backoff-ladder', + }); + await flush(); + + const task = taskRepo.getTask(taskId); + expect(task?.status).toBe('rate_limited'); + expect(task?.restrictions?.type).toBe('rate_limit'); + }); + + it('restores the task to in_progress and clears restrictions on resume', async () => { + const resetAt = Date.now() + 60 * 60 * 1000; + bus.publish('session.rate_limit_pause', { + sessionId: subSessionId, + kind: 'usage_limit', + resetAt, + reason: 'parsed-reset', + }); + await flush(); + expect(taskRepo.getTask(taskId)?.status).toBe('usage_limited'); + + bus.publish('session.rate_limit_resume', { sessionId: subSessionId }); + await flush(); + + const task = taskRepo.getTask(taskId); + expect(task?.status).toBe('in_progress'); + expect(task?.restrictions).toBeNull(); + }); + + it('ignores pause events for an unknown session (no parent task)', async () => { + bus.publish('session.rate_limit_pause', { + sessionId: 'not-a-known-session', + kind: 'usage_limit', + resetAt: Date.now() + 60000, + reason: 'parsed-reset', + }); + await flush(); + expect(taskRepo.getTask(taskId)?.status).toBe('in_progress'); + }); + + it('does not override a terminal/blocked status on pause (done/blocked/cancelled/archived)', async () => { + for (const terminal of ['done', 'blocked', 'cancelled', 'archived'] as const) { + taskRepo.updateTask(taskId, { status: 'in_progress' }); // reset between iterations + taskRepo.updateTask(taskId, { status: terminal }); + bus.publish('session.rate_limit_pause', { + sessionId: subSessionId, + kind: 'usage_limit', + resetAt: Date.now() + 60000, + reason: 'parsed-reset', + }); + await flush(); + expect(taskRepo.getTask(taskId)?.status).toBe(terminal); + } + }); + + it('a late resume does not resurrect a terminal task (cancelled/done/archived)', async () => { + // Seed a paused state + restriction, then move the task terminal via a + // manual transition (which auto-clears restrictions), then fire resume. + for (const terminal of ['cancelled', 'done', 'archived'] as const) { + taskRepo.updateTask(taskId, { + status: 'usage_limited', + restrictions: { + type: 'usage_limit', + limit: 'parsed-reset', + resetAt: Date.now() + 60000, + sessionRole: 'worker', + }, + }); + taskRepo.updateTask(taskId, { status: terminal }); + expect(taskRepo.getTask(taskId)?.restrictions).toBeNull(); + + bus.publish('session.rate_limit_resume', { sessionId: subSessionId }); + await flush(); + const task = taskRepo.getTask(taskId); + expect(task?.status).toBe(terminal); // not resurrected to in_progress + expect(task?.restrictions).toBeNull(); + } + }); + + it('resume is a no-op when the task is not currently limited', async () => { + bus.publish('session.rate_limit_resume', { sessionId: subSessionId }); + await flush(); + expect(taskRepo.getTask(taskId)?.status).toBe('in_progress'); + expect(taskRepo.getTask(taskId)?.restrictions).toBeNull(); + }); +}); diff --git a/packages/daemon/tests/unit/helpers/space-test-db.ts b/packages/daemon/tests/unit/helpers/space-test-db.ts index cef1e2169..427963a19 100644 --- a/packages/daemon/tests/unit/helpers/space-test-db.ts +++ b/packages/daemon/tests/unit/helpers/space-test-db.ts @@ -265,7 +265,7 @@ export function createSpaceTables(db: BunDatabase): void { title TEXT NOT NULL, description TEXT NOT NULL DEFAULT '', status TEXT NOT NULL DEFAULT 'open' - CHECK(status IN ('draft', 'open', 'in_progress', 'review', 'done', 'blocked', 'cancelled', 'archived', 'approved')), + CHECK(status IN ('rate_limited', 'usage_limited', 'draft', 'open', 'in_progress', 'review', 'done', 'blocked', 'cancelled', 'archived', 'approved')), priority TEXT NOT NULL DEFAULT 'normal' CHECK(priority IN ('low', 'normal', 'high', 'urgent')), labels TEXT NOT NULL DEFAULT '[]', @@ -299,6 +299,7 @@ export function createSpaceTables(db: BunDatabase): void { created_by_session TEXT DEFAULT NULL, created_by_task_schedule_id TEXT DEFAULT NULL, archived_at INTEGER, + restrictions TEXT, created_at INTEGER NOT NULL, started_at INTEGER, completed_at INTEGER, diff --git a/packages/shared/src/types/space.ts b/packages/shared/src/types/space.ts index 9badd6350..550071f9a 100644 --- a/packages/shared/src/types/space.ts +++ b/packages/shared/src/types/space.ts @@ -9,6 +9,7 @@ import type { ThinkingLevel } from '../types'; import type { SettingSource } from './settings'; import type { McpServerConfig } from './sdk-config'; +import type { TaskRestriction } from './neo'; // ============================================================================ // Space Types @@ -392,6 +393,10 @@ export interface UpdateSpaceParams { * - `blocked` — task requires human attention or intervention * - `cancelled` — task was cancelled and will not be completed * - `archived` — task is archived (soft-delete, `archivedAt` is stamped) + * - `rate_limited` — task paused because of a transient HTTP 429 rate limit; the + * owning session is in a cooldown and auto-resumes (`restrictions.resetAt`) + * - `usage_limited` — task paused because of a daily/weekly usage cap with no + * fallback left; auto-resumes when the cap resets (`restrictions.resetAt`) */ export type SpaceTaskStatus = | 'draft' @@ -402,7 +407,9 @@ export type SpaceTaskStatus = | 'done' | 'blocked' | 'cancelled' - | 'archived'; + | 'archived' + | 'rate_limited' + | 'usage_limited'; /** * Outcome an end-node agent reports via `task.reportedStatus`. @@ -817,6 +824,13 @@ export interface SpaceTask { * Schema only in PR 1; no runtime consumer yet. */ postApprovalBlockedReason?: string | null; + /** + * Restriction data when a task is paused due to an API rate or usage limit + * (`status` is `rate_limited` or `usage_limited`). Null when the task is not + * paused. Persisted so the UI can show the reset time and the runtime can + * auto-resume when the limit lifts. + */ + restrictions?: TaskRestriction | null; /** Last update timestamp (milliseconds since epoch) */ updatedAt: number; } @@ -1041,6 +1055,12 @@ export interface UpdateSpaceTaskParams { * Schema only in PR 1; no runtime consumer yet. */ postApprovalBlockedReason?: string | null; + /** + * Restriction data for a task paused on a rate/usage limit; null to clear + * (restoring the task to in_progress). Set together with `status` + * `rate_limited` / `usage_limited`. + */ + restrictions?: TaskRestriction | null; } /** diff --git a/packages/web/src/components/space/SpaceTaskPane.tsx b/packages/web/src/components/space/SpaceTaskPane.tsx index d6f07ed15..38ad5f9b9 100644 --- a/packages/web/src/components/space/SpaceTaskPane.tsx +++ b/packages/web/src/components/space/SpaceTaskPane.tsx @@ -61,6 +61,8 @@ const STATUS_LABELS: Record = { blocked: 'Blocked', cancelled: 'Cancelled', archived: 'Archived', + rate_limited: 'Rate Limited', + usage_limited: 'Usage Limited', }; const PRIORITY_LABELS: Record = { @@ -91,6 +93,8 @@ const STATUS_BADGE_CLASSES: Record = { blocked: 'border-red-500/30 bg-red-500/10 text-red-300', cancelled: 'border-gray-500/25 bg-gray-500/10 text-gray-400', archived: 'border-gray-500/25 bg-gray-500/10 text-gray-400', + rate_limited: 'border-orange-500/30 bg-orange-500/10 text-orange-300', + usage_limited: 'border-orange-600/30 bg-orange-600/10 text-orange-400', }; const PRIORITY_BADGE_CLASSES: Record = { diff --git a/packages/web/src/components/space/TaskAuxiliaryPanel.tsx b/packages/web/src/components/space/TaskAuxiliaryPanel.tsx index 387a3f88a..ab8106e5a 100644 --- a/packages/web/src/components/space/TaskAuxiliaryPanel.tsx +++ b/packages/web/src/components/space/TaskAuxiliaryPanel.tsx @@ -43,6 +43,8 @@ const STATUS_LABELS: Record = { blocked: 'Blocked', cancelled: 'Cancelled', archived: 'Archived', + rate_limited: 'Rate Limited', + usage_limited: 'Usage Limited', }; const PRIORITY_LABELS: Record = { @@ -62,6 +64,8 @@ const STATUS_BADGE_CLASSES: Record = { blocked: 'border-red-500/30 bg-red-500/10 text-red-300', cancelled: 'border-gray-500/25 bg-gray-500/10 text-gray-400', archived: 'border-gray-500/25 bg-gray-500/10 text-gray-400', + rate_limited: 'border-orange-500/30 bg-orange-500/10 text-orange-300', + usage_limited: 'border-orange-600/30 bg-orange-600/10 text-orange-400', }; const PRIORITY_BADGE_CLASSES: Record = { diff --git a/packages/web/src/components/space/TaskStatusActions.tsx b/packages/web/src/components/space/TaskStatusActions.tsx index c24df1b39..779befa02 100644 --- a/packages/web/src/components/space/TaskStatusActions.tsx +++ b/packages/web/src/components/space/TaskStatusActions.tsx @@ -16,6 +16,10 @@ export const VALID_TASK_TRANSITIONS: Record done: ['in_progress', 'archived'], blocked: ['open', 'in_progress', 'review', 'cancelled', 'archived'], cancelled: ['open', 'in_progress', 'done', 'archived'], + // Runtime-set paused states (rate/usage cap). Manual escape hatches only: + // resume, reopen, cancel, or archive. Not user-transitionable TO. + rate_limited: ['in_progress', 'open', 'cancelled', 'archived'], + usage_limited: ['in_progress', 'open', 'cancelled', 'archived'], archived: [], }; @@ -56,6 +60,17 @@ export const TRANSITION_LABELS: Record = { 'cancelled->in_progress': 'Resume', 'cancelled->done': 'Mark Done', 'cancelled->archived': 'Archive', + // Runtime-set paused states (rate/usage cap). Manual escape hatches: resume, + // reopen, or cancel. The runtime auto-resumes on cooldown fire; these render + // only for manual intervention. + 'rate_limited->in_progress': 'Resume', + 'rate_limited->open': 'Reopen', + 'rate_limited->cancelled': 'Cancel', + 'rate_limited->archived': 'Archive', + 'usage_limited->in_progress': 'Resume', + 'usage_limited->open': 'Reopen', + 'usage_limited->cancelled': 'Cancel', + 'usage_limited->archived': 'Archive', }; /**