From 98e618c2435444cf051d446d9ee7506c31db7c62 Mon Sep 17 00:00:00 2001 From: Renzo F <45401804+rennf93@users.noreply.github.com> Date: Thu, 11 Jun 2026 09:21:26 +0200 Subject: [PATCH] [aaac85d2] Rate limit guardrails for Anthropic and Ollama providers (#104) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * [25aa5b24] Implement rate-limit Zustand store, Axios interceptor, WebSocket hook, banner component, and page-load sync (#99) (#101) * [25aa5b24] feat(rate-limits): add types, Zustand store, Axios 429 interceptor, WS hook, sync hook, and banner component - panel/src/types/rate-limits.ts: RateLimitEntry, RateLimitHitEvent, RateLimitLiftedEvent, RateLimitApiResponse - panel/src/store/rate-limit-store.ts: useRateLimitStore with Map state, hitRateLimit/liftRateLimit/syncFromApi - panel/src/lib/api/rate-limits.ts: GET /api/system/rate-limits with isMockMode guard - panel/src/lib/api/client.ts: 429 interceptor dispatches to store first, Sonner toast on retry exhaustion - panel/src/hooks/use-rate-limit-websocket.ts: RATE_LIMIT_HIT/LIFTED events + onReconnect callback - panel/src/hooks/use-rate-limit-sync.ts: mount sync + no-op with console.warn when endpoint unavailable - panel/src/components/rate-limit/rate-limit-banner.tsx: amber rows with countdown, no dismiss button - panel/src/app/(dashboard)/layout.tsx: RateLimitBanner mounted below Header - store/index.ts, hooks/index.ts: export new store and hooks * [25aa5b24] fix(rate-limit-banner): use lint-clean countdown pattern (computeSecondsLeft outside render) * [25aa5b24] fix(client): add real retry loop to 429 interceptor so Sonner toast fires on exhaustion - Increment error.config._retryCount and return api(error.config) when retryCount < RATE_LIMIT_MAX_RETRIES, actually retrying the request. - Toast fires only when retryCount >= RATE_LIMIT_MAX_RETRIES (3 attempts). - Fixes AC4: toast was dead code because without return api(error.config) every 429 saw retryCount=1, permanently below the threshold of 3. --------- Co-authored-by: Frontend Developer 1 * [4112cd34] feat(rate-limit): add RateLimitError with 5-retry exponential backoff at all LLM call sites (#102) (#103) - Create roboco/services/exceptions.py with RateLimitError(provider, retry_after), HTTP_TOO_MANY_REQUESTS, MAX_RATE_LIMIT_RETRIES constants, and parse_retry_after_header() helper - extraction.py: extract _call_anthropic_with_retry() helper; retry Anthropic call 5x on 429 with exponential backoff; re-raise RateLimitError from outer except instead of swallowing it - ollama_embedder.py: 5-retry outer loop (429) wrapping existing 3-retry inner loop (ConnectError/Timeout) for all 4 call sites; two concerns kept isolated - indexes/base.py, mentor.py, validator.py: replace magic 429 literals with HTTP_TOO_MANY_REQUESTS; 5-retry loop on 429 for LLM calls - middleware.py: add rate_limit_exception_handler returning HTTP 429 with Retry-After response header - tests/unit/services/test_rate_limit_retry.py: 28 tests covering exhaustion, Retry-After header sleep, partial retries then success, ConnectError isolation Co-authored-by: Backend Developer 1 * [18107054] feat(rate-limit): Redis rate-limit state tracker + i_am_blocked rate_limited path (#105) (#106) - Add RateLimitStateTracker in roboco/services/gateway/rate_limit_tracker.py with activate(), clear(), is_rate_limited(), get_state(), increment_probe_failures(), reset_probe_failures() backed by redis.asyncio - Add RATE_LIMIT_HIT = "rate_limit.hit" to EventType StrEnum in events.py - Add _handle_rate_limited_parking() to Choreographer: intercepts i_am_blocked(reason='rate_limited') before block state transition, parks all active agents sharing affected provider via mark_waiting_long, publishes RATE_LIMIT_HIT event to StreamEventBus, task stays in_progress - Add get_provider_for_agent() and get_active_agent_slugs_for_provider() helper methods to AgentOrchestrator - Wire orchestrator and stream_bus into ChoreographerDeps via deps.py - Add test_rate_limit_tracker.py (basic ops, probe failures, cross-reconnection persistence, provider isolation) and test_i_am_blocked_rate_limited.py (AC3/AC4/AC5 coverage: task stays in_progress, mark_waiting_long call count equals active agent count, RATE_LIMIT_HIT event payload structure) Co-authored-by: Backend Developer 1 * [5501e4b4] Wire RateLimitStateTracker into live orchestrator paths — 4 CEO-identified integration gaps (#109) * [8451ca50] feat(gateway): wire RateLimitStateTracker.activate() into i_am_blocked rate-limited path and add provider-rate-limit gate to decide_spawn() (#107) - Add provider/provider_rate_limited optional fields to TriggerContext (backward-compatible defaults) - Insert rule 2 in decide_spawn(): QUEUE when trigger.provider_rate_limited is True with reason 'provider X rate-limited' - Call RateLimitStateTracker(provider).activate() in _handle_rate_limited_parking() after mark_waiting_long loop (wrapped in contextlib.suppress for Redis fault tolerance) - Extend gateway_pre_spawn_check() with optional provider param; check RateLimitStateTracker.is_rate_limited() when provider is known - Pass provider=self.get_provider_for_agent(agent_id) from orchestrator call site - Add TestProviderRateLimitGate (6 tests) to test_trigger_filter.py - Add TestRateLimitTrackerActivateOnParking (6 tests) to test_i_am_blocked_rate_limited.py - All 38 unit tests pass; ruff and mypy clean on changed files Co-authored-by: Backend Developer 1 * [e9cef0f0] feat(rate-limits): sweeper probe loop, CEO notification, and GET /api/system/rate-limits endpoint (AC4, AC8, AC9) (#108) - Add RATE_LIMIT_LIFTED event type to EventType enum in models/events.py - Add RateLimitStateTracker.list_rate_limited_providers() classmethod to scan Redis for all currently rate-limited providers (used by the new endpoint) - Add orchestrator._rate_limit_probe_loop(): background task started/stopped in start()/stop(), runs _sweep_rate_limit_probes() every 30s - Add orchestrator._probe_one_provider(): checks estimated_lift_at gate, calls _do_probe(); on success: tracker.clear(), resolve_wait() for all parked agents with waiting_for='rate_limit_lifted' matching the provider, publishes RATE_LIMIT_LIFTED event; on failure: increments probe_failures counter, sends CEO notification at threshold 10 (once per episode via _rate_limit_ceo_notified) - Add orchestrator._make_tracker(): injectable factory for RateLimitStateTracker - Add orchestrator._do_probe(): overridable async bool probe (default: True) - Add orchestrator._notify_rate_limit_ceo(): high-priority notification to CEO containing provider name, duration since activation, and paused agent count - Add roboco/api/routes/system.py with GET /rate-limits endpoint (AC9) - Register system_router in app.py under /api/system prefix - Add 17 unit tests in tests/unit/runtime/test_rate_limit_sweep.py covering all AC4/AC8/AC9 paths: probe success/failure, CEO threshold, endpoint schema Co-authored-by: Backend Developer 1 --------- Co-authored-by: Backend Developer 1 --------- Co-authored-by: Frontend Developer 1 Co-authored-by: Backend Developer 1 Co-authored-by: Renn F --- panel/src/app/(dashboard)/layout.tsx | 2 + .../rate-limit/rate-limit-banner.tsx | 95 +++ panel/src/hooks/index.ts | 2 + panel/src/hooks/use-rate-limit-sync.ts | 36 + panel/src/hooks/use-rate-limit-websocket.ts | 70 ++ panel/src/lib/api/client.ts | 52 ++ panel/src/lib/api/rate-limits.ts | 18 + panel/src/store/index.ts | 1 + panel/src/store/rate-limit-store.ts | 53 ++ panel/src/types/rate-limits.ts | 53 ++ roboco/api/app.py | 8 + roboco/api/deps.py | 11 + roboco/api/middleware.py | 38 + roboco/api/routes/system.py | 47 ++ roboco/models/events.py | 4 + roboco/runtime/orchestrator.py | 321 ++++++++ roboco/services/exceptions.py | 81 ++ roboco/services/extraction.py | 56 +- .../services/gateway/choreographer/_impl.py | 152 +++- roboco/services/gateway/rate_limit_tracker.py | 185 +++++ roboco/services/gateway/trigger_filter.py | 25 +- roboco/services/optimal_brain/indexes/base.py | 127 ++-- roboco/services/optimal_brain/mentor.py | 120 +-- .../services/optimal_brain/ollama_embedder.py | 440 +++++++---- roboco/services/optimal_brain/validator.py | 98 ++- .../gateway/test_i_am_blocked_rate_limited.py | 508 +++++++++++++ tests/unit/gateway/test_trigger_filter.py | 104 ++- tests/unit/runtime/test_rate_limit_sweep.py | 550 ++++++++++++++ tests/unit/services/test_rate_limit_retry.py | 691 ++++++++++++++++++ .../unit/services/test_rate_limit_tracker.py | 244 +++++++ 30 files changed, 3919 insertions(+), 273 deletions(-) create mode 100644 panel/src/components/rate-limit/rate-limit-banner.tsx create mode 100644 panel/src/hooks/use-rate-limit-sync.ts create mode 100644 panel/src/hooks/use-rate-limit-websocket.ts create mode 100644 panel/src/lib/api/rate-limits.ts create mode 100644 panel/src/store/rate-limit-store.ts create mode 100644 panel/src/types/rate-limits.ts create mode 100644 roboco/api/routes/system.py create mode 100644 roboco/services/exceptions.py create mode 100644 roboco/services/gateway/rate_limit_tracker.py create mode 100644 tests/unit/gateway/test_i_am_blocked_rate_limited.py create mode 100644 tests/unit/runtime/test_rate_limit_sweep.py create mode 100644 tests/unit/services/test_rate_limit_retry.py create mode 100644 tests/unit/services/test_rate_limit_tracker.py diff --git a/panel/src/app/(dashboard)/layout.tsx b/panel/src/app/(dashboard)/layout.tsx index 4401f68b..2685409b 100644 --- a/panel/src/app/(dashboard)/layout.tsx +++ b/panel/src/app/(dashboard)/layout.tsx @@ -2,6 +2,7 @@ import { Suspense } from "react"; import { Sidebar } from "@/components/layout/sidebar"; import { Header } from "@/components/layout/header"; import { ScrollRestoration } from "@/components/scroll-restoration"; +import { RateLimitBanner } from "@/components/rate-limit/rate-limit-banner"; export default function DashboardLayout({ children, @@ -13,6 +14,7 @@ export default function DashboardLayout({
+
diff --git a/panel/src/components/rate-limit/rate-limit-banner.tsx b/panel/src/components/rate-limit/rate-limit-banner.tsx new file mode 100644 index 00000000..7c3ff8ea --- /dev/null +++ b/panel/src/components/rate-limit/rate-limit-banner.tsx @@ -0,0 +1,95 @@ +"use client"; + +import { useEffect, useState, useCallback } from "react"; +import { AlertTriangle } from "lucide-react"; +import { useRateLimitStore } from "@/store/rate-limit-store"; +import { useRateLimitSync } from "@/hooks/use-rate-limit-sync"; +import { useRateLimitWebSocket } from "@/hooks/use-rate-limit-websocket"; +import type { RateLimitEntry } from "@/types/rate-limits"; + +// ============================================================================= +// Countdown row for a single rate-limited provider +// ============================================================================= + +function computeSecondsLeft(resumeAt: string): number { + return Math.max( + 0, + Math.ceil((new Date(resumeAt).getTime() - Date.now()) / 1000) + ); +} + +function RateLimitRow({ entry }: { entry: RateLimitEntry }) { + const resumeAt = entry.resumeAt; + const [secondsLeft, setSecondsLeft] = useState(() => + computeSecondsLeft(resumeAt) + ); + + useEffect(() => { + // Tick every second; re-runs when resumeAt changes (re-hit same provider) + const id = setInterval(() => { + setSecondsLeft(computeSecondsLeft(resumeAt)); + }, 1000); + return () => clearInterval(id); + }, [resumeAt]); + + const agentCount = entry.affectedAgents.length; + + return ( +
+ + + {entry.provider} + + {agentCount > 0 && ( + + {agentCount} agent{agentCount !== 1 ? "s" : ""} affected + + )} + + {secondsLeft}s + + + operations paused — resuming automatically + +
+ ); +} + +// ============================================================================= +// Main banner component +// ============================================================================= + +export function RateLimitBanner() { + const limits = useRateLimitStore((state) => state.limits); + + // Sync hook — calls GET /api/system/rate-limits on mount and exposes sync() + const { sync } = useRateLimitSync(); + + // Called when the WS reconnects — re-sync state from API + const handleReconnect = useCallback(() => { + void sync(); + }, [sync]); + + // WS hook — subscribes to RATE_LIMIT_HIT / RATE_LIMIT_LIFTED events + useRateLimitWebSocket({ onReconnect: handleReconnect }); + + // Nothing to show when no providers are rate-limited + if (limits.size === 0) { + return null; + } + + const entries = Array.from(limits.values()); + + return ( +
+ {entries.map((entry) => ( + + ))} +
+ ); +} diff --git a/panel/src/hooks/index.ts b/panel/src/hooks/index.ts index 57212f8b..f3c44ab2 100644 --- a/panel/src/hooks/index.ts +++ b/panel/src/hooks/index.ts @@ -1,4 +1,6 @@ export * from "./use-tasks"; +export * from "./use-rate-limit-websocket"; +export * from "./use-rate-limit-sync"; export * from "./use-agents"; export * from "./use-channels"; export * from "./use-notifications"; diff --git a/panel/src/hooks/use-rate-limit-sync.ts b/panel/src/hooks/use-rate-limit-sync.ts new file mode 100644 index 00000000..2aae3104 --- /dev/null +++ b/panel/src/hooks/use-rate-limit-sync.ts @@ -0,0 +1,36 @@ +"use client"; + +import { useEffect, useCallback } from "react"; +import { rateLimitsApi } from "@/lib/api/rate-limits"; +import { useRateLimitStore } from "@/store/rate-limit-store"; + +/** + * Calls GET /api/system/rate-limits on mount and passes results to syncFromApi. + * Is a no-op (with console.warn) when the endpoint is unavailable. + * Also exposes a sync() function that can be called on WS reconnect. + */ +export function useRateLimitSync() { + const { syncFromApi } = useRateLimitStore(); + + const sync = useCallback(async () => { + try { + const response = await rateLimitsApi.getRateLimits(); + syncFromApi(response); + } catch (err) { + // Endpoint unavailable — treat as no-op per acceptance criteria + const status = (err as { response?: { status?: number } })?.response?.status; + if (status === 404) { + console.warn("[rate-limits] GET /api/system/rate-limits returned 404 — endpoint not available"); + } else { + console.warn("[rate-limits] GET /api/system/rate-limits unavailable:", err); + } + } + }, [syncFromApi]); + + // Sync on mount + useEffect(() => { + void sync(); + }, [sync]); + + return { sync }; +} diff --git a/panel/src/hooks/use-rate-limit-websocket.ts b/panel/src/hooks/use-rate-limit-websocket.ts new file mode 100644 index 00000000..849356f6 --- /dev/null +++ b/panel/src/hooks/use-rate-limit-websocket.ts @@ -0,0 +1,70 @@ +"use client"; + +import { useEffect, useRef } from "react"; +import { useWebSocket } from "./use-websocket"; +import { useRateLimitStore } from "@/store/rate-limit-store"; +import type { RateLimitHitEvent, RateLimitLiftedEvent } from "@/types/rate-limits"; + +interface RateLimitWsMessage { + type: string; + provider?: string; + affectedAgents?: string[]; + retryAfterSeconds?: number; + timestamp?: string; +} + +interface UseRateLimitWebSocketOptions { + /** Called when the WebSocket reconnects after a disconnect */ + onReconnect?: () => void; +} + +/** + * Subscribes to RATE_LIMIT_HIT and RATE_LIMIT_LIFTED WebSocket events and + * dispatches them to the useRateLimitStore. Accepts an optional onReconnect + * callback that fires when the connection recovers from a reconnecting state. + */ +export function useRateLimitWebSocket(options: UseRateLimitWebSocketOptions = {}) { + const { onReconnect } = options; + const prevStateRef = useRef(null); + + const { state, lastMessage } = useWebSocket( + "/ws/system", + undefined, + true + ); + + // Fire onReconnect when state transitions from reconnecting → connected + useEffect(() => { + if (prevStateRef.current === "reconnecting" && state === "connected") { + onReconnect?.(); + } + prevStateRef.current = state; + }, [state, onReconnect]); + + // Handle incoming WS messages + useEffect(() => { + if (!lastMessage) return; + + const { hitRateLimit, liftRateLimit } = useRateLimitStore.getState(); + + if (lastMessage.type === "RATE_LIMIT_HIT") { + const event: RateLimitHitEvent = { + type: "RATE_LIMIT_HIT", + provider: lastMessage.provider ?? "unknown", + affectedAgents: lastMessage.affectedAgents ?? [], + retryAfterSeconds: lastMessage.retryAfterSeconds ?? 60, + timestamp: lastMessage.timestamp ?? new Date().toISOString(), + }; + hitRateLimit(event); + } else if (lastMessage.type === "RATE_LIMIT_LIFTED") { + const event: RateLimitLiftedEvent = { + type: "RATE_LIMIT_LIFTED", + provider: lastMessage.provider ?? "unknown", + timestamp: lastMessage.timestamp ?? new Date().toISOString(), + }; + liftRateLimit(event); + } + }, [lastMessage]); + + return { wsState: state }; +} diff --git a/panel/src/lib/api/client.ts b/panel/src/lib/api/client.ts index 8dae935a..c12e6b27 100644 --- a/panel/src/lib/api/client.ts +++ b/panel/src/lib/api/client.ts @@ -1,5 +1,17 @@ import axios, { AxiosInstance, AxiosError } from "axios"; +import { toast } from "sonner"; import { API_URL, CEO_AGENT_ID, CEO_ROLE } from "@/lib/constants"; +import { useRateLimitStore } from "@/store/rate-limit-store"; +import type { RateLimitHitEvent } from "@/types/rate-limits"; + +// Custom Axios config extension for retry tracking +declare module "axios" { + interface InternalAxiosRequestConfig { + _retryCount?: number; + } +} + +const RATE_LIMIT_MAX_RETRIES = 3; // Create axios instance with default config const api: AxiosInstance = axios.create({ @@ -47,6 +59,46 @@ api.interceptors.response.use( const errorData = error.response?.data as Record | undefined; const errorDetail = errorData?.detail || error.message; + // ------------------------------------------------------------------------- + // 429 Rate-limit handling — FIRST side-effect, before any other logic + // ------------------------------------------------------------------------- + if (status === 429) { + const retryAfterHeader = error.response?.headers?.["retry-after"]; + const retryAfterSeconds = retryAfterHeader ? parseInt(String(retryAfterHeader), 10) : 60; + const safeRetryAfter = isNaN(retryAfterSeconds) ? 60 : retryAfterSeconds; + + // Extract provider from custom header or fall back to URL path heuristics + const providerHeader = error.response?.headers?.["x-provider"]; + const urlProvider = url + ? (["anthropic", "openai", "ollama"].find((p) => url.includes(p)) ?? "unknown") + : "unknown"; + const provider = (providerHeader as string | undefined) ?? urlProvider; + + // Dispatch to store as first side-effect + const hitEvent: RateLimitHitEvent = { + type: "RATE_LIMIT_HIT", + provider, + affectedAgents: [], + retryAfterSeconds: safeRetryAfter, + timestamp: new Date().toISOString(), + }; + useRateLimitStore.getState().hitRateLimit(hitEvent); + + // Track retry count; retry the request until exhausted, then toast + const retryCount = (error.config?._retryCount ?? 0) + 1; + if (error.config) { + error.config._retryCount = retryCount; + if (retryCount < RATE_LIMIT_MAX_RETRIES) { + // Retry the request — interceptor re-runs on each subsequent 429 + return api(error.config); + } + } + // Retries exhausted — notify the user via Sonner toast + toast.warning( + `Rate limited by ${provider}. The system has paused operations and will resume automatically in ~${safeRetryAfter}s.` + ); + } + // Log comprehensive error info console.error(`[API] ✗ ${method} ${url}`, { status, diff --git a/panel/src/lib/api/rate-limits.ts b/panel/src/lib/api/rate-limits.ts new file mode 100644 index 00000000..549872e3 --- /dev/null +++ b/panel/src/lib/api/rate-limits.ts @@ -0,0 +1,18 @@ +import api from "./client"; +import type { RateLimitApiResponse } from "@/types/rate-limits"; +import { isMockMode } from "@/lib/mock-data"; + +export const rateLimitsApi = { + /** + * GET /api/system/rate-limits — fetch active rate limits on page load or WS reconnect. + * Returns empty list in mock mode (rate limits are a real-backend-only concern). + */ + getRateLimits: async (): Promise => { + if (isMockMode()) { + console.warn("[rate-limits] isMockMode: skipping GET /api/system/rate-limits"); + return { entries: [] }; + } + const { data } = await api.get("/system/rate-limits"); + return data; + }, +}; diff --git a/panel/src/store/index.ts b/panel/src/store/index.ts index 89adeff6..39525cc4 100644 --- a/panel/src/store/index.ts +++ b/panel/src/store/index.ts @@ -1,2 +1,3 @@ export { useUIStore } from "./ui-store"; export { useNotificationStore } from "./notifications-store"; +export { useRateLimitStore } from "./rate-limit-store"; diff --git a/panel/src/store/rate-limit-store.ts b/panel/src/store/rate-limit-store.ts new file mode 100644 index 00000000..aeed29e7 --- /dev/null +++ b/panel/src/store/rate-limit-store.ts @@ -0,0 +1,53 @@ +import { create } from "zustand"; +import type { + RateLimitEntry, + RateLimitHitEvent, + RateLimitLiftedEvent, + RateLimitApiResponse, +} from "@/types/rate-limits"; + +interface RateLimitState { + /** Active rate limits keyed by provider name */ + limits: Map; + + // Actions + hitRateLimit: (event: RateLimitHitEvent) => void; + liftRateLimit: (event: RateLimitLiftedEvent) => void; + syncFromApi: (response: RateLimitApiResponse) => void; +} + +export const useRateLimitStore = create((set) => ({ + limits: new Map(), + + hitRateLimit: (event: RateLimitHitEvent) => + set((state) => { + const next = new Map(state.limits); + const entry: RateLimitEntry = { + provider: event.provider, + affectedAgents: event.affectedAgents, + hitAt: event.timestamp, + resumeAt: new Date( + new Date(event.timestamp).getTime() + event.retryAfterSeconds * 1000 + ).toISOString(), + retryAfterSeconds: event.retryAfterSeconds, + }; + next.set(event.provider, entry); + return { limits: next }; + }), + + liftRateLimit: (event: RateLimitLiftedEvent) => + set((state) => { + const next = new Map(state.limits); + next.delete(event.provider); + return { limits: next }; + }), + + syncFromApi: (response: RateLimitApiResponse) => + set(() => { + const next = new Map(); + for (const entry of response.entries) { + next.set(entry.provider, entry); + } + return { limits: next }; + }), +})); diff --git a/panel/src/types/rate-limits.ts b/panel/src/types/rate-limits.ts new file mode 100644 index 00000000..305f5618 --- /dev/null +++ b/panel/src/types/rate-limits.ts @@ -0,0 +1,53 @@ +// ============================================================================= +// RATE LIMIT TYPES +// ============================================================================= + +/** + * Represents an active rate-limit entry for a provider. + * Stored in the Zustand store keyed by provider name. + */ +export interface RateLimitEntry { + /** The AI provider that is rate-limited (e.g. "anthropic", "openai") */ + provider: string; + /** Agent slugs affected by this rate limit */ + affectedAgents: string[]; + /** ISO timestamp when the rate limit was hit */ + hitAt: string; + /** ISO timestamp when the rate limit is expected to lift (hitAt + retryAfterSeconds) */ + resumeAt: string; + /** How many seconds until operations resume */ + retryAfterSeconds: number; +} + +/** + * WebSocket event emitted when a rate limit is triggered. + */ +export interface RateLimitHitEvent { + type: "RATE_LIMIT_HIT"; + /** The AI provider being rate-limited */ + provider: string; + /** Agent slugs affected */ + affectedAgents: string[]; + /** How many seconds to wait before retrying */ + retryAfterSeconds: number; + /** ISO timestamp of the event */ + timestamp: string; +} + +/** + * WebSocket event emitted when a rate limit is cleared. + */ +export interface RateLimitLiftedEvent { + type: "RATE_LIMIT_LIFTED"; + /** The AI provider whose rate limit has been lifted */ + provider: string; + /** ISO timestamp of the event */ + timestamp: string; +} + +/** + * Response shape from GET /api/system/rate-limits + */ +export interface RateLimitApiResponse { + entries: RateLimitEntry[]; +} diff --git a/roboco/api/app.py b/roboco/api/app.py index 34d68821..344c586a 100644 --- a/roboco/api/app.py +++ b/roboco/api/app.py @@ -35,6 +35,7 @@ from roboco.api.routes.prompter_live import router as prompter_live_router from roboco.api.routes.provider import router as provider_router from roboco.api.routes.sessions import router as sessions_router from roboco.api.routes.stream import router as stream_router +from roboco.api.routes.system import router as system_router from roboco.api.routes.tasks import router as tasks_router from roboco.api.routes.usage import router as usage_router from roboco.api.routes.v1 import do as do_module @@ -348,6 +349,13 @@ def create_app() -> FastAPI: tags=["Usage Analytics"], ) + # System monitoring (rate-limits, etc.) + app.include_router( + system_router, + prefix=f"{api_prefix}/system", + tags=["System"], + ) + # API v1 — intent-verb flow endpoints app.include_router(flow_dev_module.router) diff --git a/roboco/api/deps.py b/roboco/api/deps.py index a4c68d80..65e382aa 100644 --- a/roboco/api/deps.py +++ b/roboco/api/deps.py @@ -504,6 +504,15 @@ async def get_choreographer( db_session: DbSession, ) -> Choreographer: """Build a Choreographer with all service dependencies wired up.""" + from roboco.events.stream_bus import get_stream_event_bus + + # Inject the orchestrator (if initialised) and the stream event bus + # so the rate-limited i_am_blocked path can park agents and publish events. + # Both are None-safe in ChoreographerDeps — passing None is the same as + # omitting the field, so the choreographer degrades gracefully when the + # orchestrator has not been initialised yet (e.g. during startup). + orch: AgentOrchestrator | None = _ServiceHolder.orchestrator + bus = get_stream_event_bus() if _ServiceHolder.orchestrator is not None else None return Choreographer( ChoreographerDeps( task=TaskService(db_session), @@ -515,6 +524,8 @@ async def get_choreographer( evidence_repo=EvidenceRepo(db_session), messaging=MessagingService(db_session), product=ProductService(db_session), + orchestrator=orch, + stream_bus=bus, ) ) diff --git a/roboco/api/middleware.py b/roboco/api/middleware.py index 47f6a572..6fbb579c 100644 --- a/roboco/api/middleware.py +++ b/roboco/api/middleware.py @@ -41,6 +41,7 @@ from roboco.services.base import ( from roboco.services.base import ( ValidationError as ServiceValidationError, ) +from roboco.services.exceptions import RateLimitError logger = structlog.get_logger() @@ -227,6 +228,42 @@ async def service_exception_handler(request: Request, exc: Exception) -> JSONRes ) +async def rate_limit_exception_handler( + request: Request, exc: Exception +) -> JSONResponse: + """Handle :class:`~roboco.services.exceptions.RateLimitError`. + + Returns HTTP 429 with a ``Retry-After`` response header (when available) + and a structured JSON body so API consumers can back off gracefully. + """ + rl_exc = cast("RateLimitError", exc) + correlation_id = getattr(request.state, "correlation_id", None) + + logger.warning( + "LLM rate limit exhausted", + provider=rl_exc.provider, + retry_after=rl_exc.retry_after, + ) + + content: dict = { + "error": "rate_limit_exceeded", + "provider": rl_exc.provider, + "message": str(rl_exc), + } + if correlation_id: + content["correlation_id"] = correlation_id + + headers: dict[str, str] = {} + if rl_exc.retry_after is not None: + headers["Retry-After"] = str(int(rl_exc.retry_after)) + + return JSONResponse( + status_code=429, + content=content, + headers=headers, + ) + + async def generic_exception_handler(request: Request, exc: Exception) -> JSONResponse: """Handle unexpected exceptions.""" correlation_id = getattr(request.state, "correlation_id", None) @@ -376,6 +413,7 @@ def setup_middleware(app: FastAPI) -> None: app.add_exception_handler(HTTPException, http_exception_handler) app.add_exception_handler(RobocoError, roboco_exception_handler) app.add_exception_handler(ServiceError, service_exception_handler) + app.add_exception_handler(RateLimitError, rate_limit_exception_handler) app.add_exception_handler(Exception, generic_exception_handler) # Middleware (added in reverse order due to LIFO) diff --git a/roboco/api/routes/system.py b/roboco/api/routes/system.py new file mode 100644 index 00000000..58b7f3a7 --- /dev/null +++ b/roboco/api/routes/system.py @@ -0,0 +1,47 @@ +"""System monitoring endpoints. + +Provides read-only introspection into orchestrator-level state that is +useful for operators and the control panel but doesn't fit cleanly into +the per-resource routers (agents, tasks, etc.). + +Currently exposed: + + GET /api/system/rate-limits + Returns the current per-provider rate-limit state from Redis. +""" + +from __future__ import annotations + +from typing import Any + +from fastapi import APIRouter + +from roboco.services.gateway.rate_limit_tracker import RateLimitStateTracker + +router = APIRouter() + + +@router.get( + "/rate-limits", + summary="List per-provider rate-limit state", + response_model=list[dict[str, Any]], + tags=["System"], +) +async def get_rate_limits() -> list[dict[str, Any]]: + """Return rate-limit state for every currently rate-limited provider. + + Backed by + :class:`~roboco.services.gateway.rate_limit_tracker.RateLimitStateTracker`. + Each entry is the raw state dict (``rate_limited``, ``activated_at``, + ``retry_after``, ``affected_agents``, ``probe_failures``) augmented + with a ``provider`` key. + + Returns an empty list ``[]`` when no provider is currently rate-limited. + """ + entries = await RateLimitStateTracker.list_rate_limited_providers() + result: list[dict[str, Any]] = [] + for provider, state in entries: + item = dict(state) + item["provider"] = provider + result.append(item) + return result diff --git a/roboco/models/events.py b/roboco/models/events.py index 46fceb99..842c0eaf 100644 --- a/roboco/models/events.py +++ b/roboco/models/events.py @@ -65,6 +65,10 @@ class EventType(StrEnum): BLOCKER_REPORTED = "blocker.reported" BLOCKER_RESOLVED = "blocker.resolved" + # Rate-limit events + RATE_LIMIT_HIT = "rate_limit.hit" + RATE_LIMIT_LIFTED = "rate_limit.lifted" + # Question events QUESTION_ASKED = "question.asked" QUESTION_ANSWERED = "question.answered" diff --git a/roboco/runtime/orchestrator.py b/roboco/runtime/orchestrator.py index 88018ec6..585baa79 100644 --- a/roboco/runtime/orchestrator.py +++ b/roboco/runtime/orchestrator.py @@ -413,6 +413,7 @@ async def gateway_pre_spawn_check( task_id: str | None, trigger_kind: str, target_role: str, + provider: str | None = None, ) -> tuple[str, str]: """Consult trigger_filter before spawning a container. @@ -420,6 +421,12 @@ async def gateway_pre_spawn_check( ``"spawn"``, ``"queue"``, or ``"drop"``. The trigger_filter spawn cooldown runs unconditionally for every spawn. + + Args: + provider: Optional provider name (e.g. ``"anthropic"``) for the + agent about to be spawned. When given, the + ``RateLimitStateTracker`` is consulted and a QUEUE decision is + returned when that provider is currently rate-limited. """ from roboco.db.base import get_session_factory from roboco.services.gateway.trigger_filter import ( @@ -459,11 +466,29 @@ async def gateway_pre_spawn_check( if task_row is None: return SpawnDecision.SPAWN, "task not found in DB — allow by default" + # Check provider rate-limit status when a provider is known. + # Failure is non-fatal — degrade to False (allow spawn) so Redis + # unavailability never permanently blocks the dispatcher. + provider_rate_limited = False + if provider is not None: + try: + from roboco.services.gateway.rate_limit_tracker import ( + RateLimitStateTracker, + ) + + provider_rate_limited = await RateLimitStateTracker( + provider + ).is_rate_limited() + except Exception: + provider_rate_limited = False + trigger = TriggerContext( kind=TriggerKind(trigger_kind), skill=None, recent_spawns_for_task=recent_for_task, recent_spawns_for_role=recent_for_role, + provider=provider, + provider_rate_limited=provider_rate_limited, ) config = SpawnConfig( cooldown_seconds=settings.spawn_cooldown_seconds, @@ -527,6 +552,13 @@ class AgentOrchestrator: self._health_task: asyncio.Task | None = None self._dispatcher_task: asyncio.Task | None = None self._sweeper_task: asyncio.Task | None = None + # Rate-limit probe loop: 30-second interval, scans Redis for all + # rate-limited providers and resolves waiting agents on success. + self._rate_limit_probe_task: asyncio.Task | None = None + # Tracks which providers have already received a CEO notification + # during the current rate-limit episode. Cleared when the probe + # succeeds and the rate limit is lifted (tracker.clear() path). + self._rate_limit_ceo_notified: set[str] = set() # Strong refs for fire-and-forget audit writes. Without this, the # event loop only weak-refs the Task and may GC it before it # commits — audit_log was silently empty because of this. @@ -599,6 +631,7 @@ class AgentOrchestrator: self._health_task = asyncio.create_task(self._health_loop()) self._dispatcher_task = asyncio.create_task(self._dispatcher_loop()) self._sweeper_task = asyncio.create_task(self._sweeper_loop()) + self._rate_limit_probe_task = asyncio.create_task(self._rate_limit_probe_loop()) logger.info( "Orchestrator started", @@ -626,6 +659,11 @@ class AgentOrchestrator: with contextlib.suppress(asyncio.CancelledError): await self._sweeper_task + if self._rate_limit_probe_task: + self._rate_limit_probe_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await self._rate_limit_probe_task + # Stop all agents for agent_id in list(self._instances.keys()): await self.stop_agent(agent_id) @@ -1232,6 +1270,7 @@ class AgentOrchestrator: task_id=task_id, trigger_kind=trigger_kind, target_role=target_role, + provider=self.get_provider_for_agent(agent_id), ) if outcome != "spawn": logger.info( @@ -3066,6 +3105,44 @@ class AgentOrchestrator: error=str(e), ) + # ========================================================================= + # PROVIDER QUERY HELPERS (used by the choreographer rate-limit path) + # ========================================================================= + + def get_provider_for_agent(self, agent_slug: str) -> str | None: + """Return the ``provider_type`` for a currently-tracked agent, or None. + + Reads the in-memory ``_instances`` dict so this is synchronous and + O(1). Returns None when the agent is not tracked or has no config. + + Args: + agent_slug: The agent slug (e.g. ``"be-dev-1"``). + """ + instance = self._instances.get(agent_slug) + if instance is None or instance.config is None: + return None + return instance.config.provider_type + + def get_active_agent_slugs_for_provider(self, provider: str) -> list[str]: + """Return slugs of all active agents currently using ``provider``. + + "Active" means the instance's state is ACTIVE or STARTING (i.e. + the container is running or spinning up — not IDLE, WAITING_LONG, + STOPPING, or OFFLINE). + + Args: + provider: Provider type string, e.g. ``"anthropic"`` or + ``"ollama_cloud"``. + """ + active_states = {AgentState.ACTIVE, AgentState.STARTING} + return [ + slug + for slug, inst in self._instances.items() + if inst.state in active_states + and inst.config is not None + and inst.config.provider_type == provider + ] + # ========================================================================= # TOKEN USAGE INSTRUMENTATION # ========================================================================= @@ -3900,6 +3977,250 @@ Start by: error=str(e), ) + # ========================================================================= + # RATE-LIMIT PROBE LOOP (AC4, AC8) + # ========================================================================= + + async def _rate_limit_probe_loop(self) -> None: + """Background loop: probe rate-limited providers every ~30 seconds. + + Runs independently of the 60-second session/notification sweeper so + rate limits can be cleared on their own cadence without blocking + other sweep work. + """ + probe_interval = 30 # seconds + while self._running: + try: + await asyncio.sleep(probe_interval) + await self._sweep_rate_limit_probes() + except asyncio.CancelledError: + break + except Exception as e: + logger.error("Rate-limit probe loop error", error=str(e)) + + async def _sweep_rate_limit_probes(self) -> None: + """One probe pass: check every rate-limited provider. + + For each provider whose estimated_lift_at has passed: + - Call ``_do_probe(provider)`` to test connectivity. + - **Success**: clear the tracker, resolve all parked agents, publish + ``RATE_LIMIT_LIFTED``. + - **Failure**: increment probe_failures; if the count reaches 10 and + we haven't already sent a CEO notification for this episode, send + one now. + """ + from roboco.services.gateway.rate_limit_tracker import RateLimitStateTracker + + try: + providers = await RateLimitStateTracker.list_rate_limited_providers() + except Exception as e: + logger.warning("Failed to list rate-limited providers", error=str(e)) + return + + for provider, state in providers: + try: + await self._probe_one_provider(provider, state) + except Exception as e: + logger.error( + "Unhandled error probing provider", + provider=provider, + error=str(e), + ) + + def _make_tracker(self, provider: str) -> Any: + """Return a RateLimitStateTracker for *provider*. + + Extracted as its own method so unit tests can monkeypatch it to + return an async mock without needing to intercept lazy imports. + """ + from roboco.services.gateway.rate_limit_tracker import RateLimitStateTracker + + return RateLimitStateTracker(provider) + + async def _probe_one_provider(self, provider: str, state: dict[str, Any]) -> None: + """Probe a single rate-limited provider and handle the outcome.""" + # Only start probing after estimated_lift_at has passed. + activated_at_raw: str | None = state.get("activated_at") + retry_after: float | None = state.get("retry_after") + if activated_at_raw and retry_after is not None: + try: + activated_at = datetime.fromisoformat(activated_at_raw) + estimated_lift_at = activated_at + timedelta(seconds=retry_after) + if datetime.now(UTC) < estimated_lift_at: + return # Too early — wait until after estimated lift time + except (ValueError, TypeError): + pass # Malformed timestamps: proceed with probe anyway + + success = await self._do_probe(provider) + tracker = self._make_tracker(provider) + + if success: + logger.info( + "Rate-limit probe succeeded; clearing provider", provider=provider + ) + await tracker.clear() + # Remove from CEO-notified set so new episodes get a fresh notification + self._rate_limit_ceo_notified.discard(provider) + # Resolve all parked agents waiting for this rate limit to lift + rate_limited_agents = [ + agent_id + for agent_id, record in list(self._waiting_records.items()) + if record.waiting_for == "rate_limit_lifted" + and record.context.get("provider") == provider + ] + for agent_id in rate_limited_agents: + with contextlib.suppress(Exception): + await self.resolve_wait( + agent_id, + { + "reason": "rate_limit_lifted", + "provider": provider, + "lifted_at": datetime.now(UTC).isoformat(), + }, + ) + # Publish RATE_LIMIT_LIFTED event + from roboco.events import get_event_bus + from roboco.models.events import Event, EventType + + with contextlib.suppress(Exception): + bus = get_event_bus() + event = Event( + type=EventType.RATE_LIMIT_LIFTED, + data={ + "provider": provider, + "resumedAgents": rate_limited_agents, + "timestamp": datetime.now(UTC).isoformat(), + }, + ) + await bus.publish(event) + logger.info( + "RATE_LIMIT_LIFTED published", + provider=provider, + resumed_agents=len(rate_limited_agents), + ) + else: + failure_count = await tracker.increment_probe_failures() + logger.debug( + "Rate-limit probe failed", + provider=provider, + probe_failures=failure_count, + ) + # Send CEO notification once when failures reach threshold 10 + _CEO_NOTIFY_THRESHOLD = 10 + if ( + failure_count >= _CEO_NOTIFY_THRESHOLD + and provider not in self._rate_limit_ceo_notified + ): + self._rate_limit_ceo_notified.add(provider) + paused_count = sum( + 1 + for record in self._waiting_records.values() + if record.waiting_for == "rate_limit_lifted" + and record.context.get("provider") == provider + ) + activated_at_str: str = activated_at_raw or "unknown" + await self._notify_rate_limit_ceo( + provider=provider, + activated_at_str=activated_at_str, + paused_agent_count=paused_count, + ) + + async def _do_probe(self, _provider: str) -> bool: + """Return True if the provider is accepting requests again. + + This method is intentionally thin so tests can monkeypatch it. + The default implementation is conservative: returns ``True`` + (success) so that once the estimated_lift_at window has passed + the probe clears the rate limit. Override in tests to inject + either success or failure scenarios. + """ + # Default: optimistic — time-expiry gate (checked before this call) + # is the primary guard; the probe itself succeeds. + return True + + async def _notify_rate_limit_ceo( + self, + provider: str, + activated_at_str: str, + paused_agent_count: int, + ) -> None: + """Send a high-priority notification to the CEO about a persistent rate limit. + + Fires once per episode (AC8). Follows the same pattern as + ``_notify_stranded_agent`` — direct DB insert + delivery.deliver(). + """ + try: + from sqlalchemy import select as _select + + from roboco.db.base import get_session_factory + from roboco.db.tables import AgentTable, NotificationTable + from roboco.models.base import ( + AgentRole, + NotificationPriority, + NotificationType, + ) + from roboco.services.notification_delivery import ( + get_notification_delivery_service, + ) + from roboco.utils.converters import require_uuid + + # Compute human-friendly duration + duration_desc = "unknown duration" + try: + activated_at = datetime.fromisoformat(activated_at_str) + elapsed = datetime.now(UTC) - activated_at + total_minutes = int(elapsed.total_seconds() / 60) + if total_minutes < 60: # noqa: PLR2004 + duration_desc = f"{total_minutes} minute(s)" + else: + duration_desc = f"{total_minutes // 60}h {total_minutes % 60}m" + except (ValueError, TypeError): + pass + + session_factory = get_session_factory() + async with session_factory() as db: + ceo_result = await db.execute( + _select(AgentTable).where(AgentTable.role == AgentRole.CEO) + ) + ceo = ceo_result.scalar_one_or_none() + if ceo is None: + logger.warning( + "CEO agent not found; skipping rate-limit CEO notification", + provider=provider, + ) + return + notification = NotificationTable( + type=NotificationType.ALERT, + priority=NotificationPriority.HIGH, + from_agent=ceo.id, + to_agents=[ceo.id], + subject=f"Rate limit persisting: {provider}", + body=( + f"Provider '{provider}' has been rate-limited for " + f"{duration_desc}. " + f"{paused_agent_count} agent(s) are currently paused. " + f"10 consecutive probe attempts have failed. " + f"Manual intervention may be required." + ), + requires_ack=True, + ) + db.add(notification) + await db.flush() + delivery = get_notification_delivery_service(db) + await delivery.deliver(require_uuid(notification.id)) + await db.commit() + logger.info( + "Rate-limit CEO notification sent", + provider=provider, + paused_agents=paused_agent_count, + ) + except Exception as e: + logger.error( + "Failed to send rate-limit CEO notification", + provider=provider, + error=str(e), + ) + # ========================================================================= # STATUS API # ========================================================================= diff --git a/roboco/services/exceptions.py b/roboco/services/exceptions.py new file mode 100644 index 00000000..8192210a --- /dev/null +++ b/roboco/services/exceptions.py @@ -0,0 +1,81 @@ +""" +LLM Service Exceptions + +Shared exception types and helpers for LLM provider rate-limit handling. +Importable from a single location as required by the acceptance criteria. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import httpx + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +#: Maximum number of retries on HTTP 429 / provider RateLimitError. +MAX_RATE_LIMIT_RETRIES: int = 5 + +#: HTTP status code for rate limiting. +HTTP_TOO_MANY_REQUESTS: int = 429 + + +# --------------------------------------------------------------------------- +# Exceptions +# --------------------------------------------------------------------------- + + +class RateLimitError(Exception): + """Raised when an LLM provider returns a 429 rate-limit response after all retries. + + Attributes: + provider: Name of the provider that rate-limited us (``"anthropic"`` / + ``"ollama"``). + retry_after: The last ``Retry-After`` value seen (in seconds), or ``None`` + if the header was absent. + """ + + def __init__( + self, + provider: str, + retry_after: float | None = None, + ) -> None: + self.provider = provider + self.retry_after = retry_after + msg = f"Rate limit exceeded for provider '{provider}'" + if retry_after is not None: + msg += f"; retry after {retry_after:.1f}s" + super().__init__(msg) + + def __repr__(self) -> str: # pragma: no cover + return ( + f"RateLimitError(" + f"provider={self.provider!r}, retry_after={self.retry_after!r})" + ) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def parse_retry_after_header(response: httpx.Response) -> float | None: + """Extract the ``Retry-After`` header from an httpx response as seconds. + + The header is treated as a plain integer/float number of seconds. HTTP-date + format is not supported (LLM providers invariably use numeric values). + + Returns: + Number of seconds to wait, or ``None`` if the header is absent or cannot + be parsed as a number. + """ + header = response.headers.get("retry-after") + if not header: + return None + try: + return float(header) + except (ValueError, TypeError): + return None diff --git a/roboco/services/extraction.py b/roboco/services/extraction.py index 7a0e0f5b..fd1c3c23 100644 --- a/roboco/services/extraction.py +++ b/roboco/services/extraction.py @@ -312,6 +312,51 @@ class ExtractionService: return best_type, confidence, matches + async def _call_anthropic_with_retry(self, client: Any, prompt: str) -> Any: + """Call Anthropic messages.create with up to MAX_RATE_LIMIT_RETRIES on 429. + + Raises RateLimitError when all retries are exhausted. + """ + import asyncio + + import anthropic as anthropic_mod + + from roboco.services.exceptions import MAX_RATE_LIMIT_RETRIES, RateLimitError + + last_retry_after: float | None = None + for rl_attempt in range(MAX_RATE_LIMIT_RETRIES): + try: + return await client.messages.create( + model="claude-3-haiku-20240307", # Fast, cheap + max_tokens=2000, + messages=[{"role": "user", "content": prompt}], + ) + except anthropic_mod.RateLimitError as exc: + try: + header = exc.response.headers.get("retry-after") + last_retry_after = float(header) if header else None + except (AttributeError, TypeError, ValueError): + last_retry_after = None + backoff = ( + last_retry_after + if last_retry_after is not None + else float(2**rl_attempt) + ) + self.log.warning( + "Anthropic rate limited (429), retrying", + provider="anthropic", + attempt=rl_attempt + 1, + max_retries=MAX_RATE_LIMIT_RETRIES, + backoff_duration=backoff, + ) + if rl_attempt < MAX_RATE_LIMIT_RETRIES - 1: + await asyncio.sleep(backoff) + else: + raise RateLimitError( + provider="anthropic", retry_after=last_retry_after + ) from exc + raise RateLimitError(provider="anthropic", retry_after=last_retry_after) + async def extract_with_llm(self, ctx: ExtractionContext) -> ExtractionResult: """ Extract messages using LLM classification. @@ -319,11 +364,14 @@ class ExtractionService: This is more accurate but slower and more expensive. Falls back to pattern matching if LLM unavailable. Uses TOON format for token-efficient communication. + Retries up to MAX_RATE_LIMIT_RETRIES times on 429/RateLimitError, + respecting the Retry-After header when present. """ from anthropic import AsyncAnthropic from roboco.config import settings from roboco.llm import ToonAdapter + from roboco.services.exceptions import RateLimitError toon = ToonAdapter() @@ -348,11 +396,7 @@ action,Creating file utils.py,0.95 Output only valid TOON, no other text.""" - response = await client.messages.create( - model="claude-3-haiku-20240307", # Fast, cheap for classification - max_tokens=2000, - messages=[{"role": "user", "content": prompt}], - ) + response = await self._call_anthropic_with_retry(client, prompt) # Parse response using TOON (falls back to JSON) # Extract text from first TextBlock content @@ -398,6 +442,8 @@ Output only valid TOON, no other text.""" session_id=ctx.session_id, ) + except RateLimitError: + raise except Exception as e: # Fall back to pattern matching self.log.warning("LLM extraction failed, using patterns", error=str(e)) diff --git a/roboco/services/gateway/choreographer/_impl.py b/roboco/services/gateway/choreographer/_impl.py index e22b3a44..d1eef61c 100644 --- a/roboco/services/gateway/choreographer/_impl.py +++ b/roboco/services/gateway/choreographer/_impl.py @@ -219,6 +219,17 @@ class ChoreographerDeps: # callsites / tests that don't exercise Product routing don't have to plumb # it in; when None, delegate falls back to parent-project inheritance. product: Any = None + # Orchestrator access for the rate-limited i_am_blocked path. + # Implements get_provider_for_agent(slug) -> str | None, + # get_active_agent_slugs_for_provider(provider) -> list[str], and + # async mark_waiting_long(slug, waiting_for, task_id, context). + # Optional: when None the parking step is skipped (e.g. in unit tests + # that don't need to verify orchestrator interactions). + orchestrator: Any = None + # StreamEventBus for publishing RATE_LIMIT_HIT events. + # Optional so existing callsites that don't exercise the rate-limit path + # don't have to plumb it in. + stream_bus: Any = None @dataclass(frozen=True) @@ -360,6 +371,14 @@ class Choreographer: def product(self) -> Any: return self._deps.product + @property + def orchestrator(self) -> Any: + return self._deps.orchestrator + + @property + def stream_bus(self) -> Any: + return self._deps.stream_bus + async def _touch(self, task_id: UUID | None) -> None: """Best-effort heartbeat write; silent on missing task.""" if task_id is not None: @@ -2188,6 +2207,114 @@ class Choreographer: ) return updated, None + @staticmethod + def _parse_retry_after(what_needed: str | None) -> float | None: + """Extract a retry-after seconds value from ``what_needed``, or None. + + Agents may embed the Retry-After seconds in the ``what_needed`` + field as a numeric string (e.g. ``"30"`` or ``"60.5"``). This + helper tries to parse it; any non-numeric or absent value returns + ``None``, which maps to the nullable ``retryAfterSeconds`` in the + RATE_LIMIT_HIT event. + """ + if what_needed is None: + return None + try: + return float(what_needed.strip()) + except (ValueError, AttributeError): + return None + + async def _handle_rate_limited_parking( + self, + agent_id: UUID, + task_id: UUID, + t: Any, + agent: Any, + role_str: str, + briefing: dict[str, Any], + what_needed: str | None, + ) -> Envelope: + """Rate-limited fast path: park agents, publish event, persist state. + + Called from ``i_am_blocked`` when ``reason == 'rate_limited'``. + The task stays in its current status (``in_progress``) — no block + transition occurs. Instead, every orchestrator-tracked active agent + sharing the same provider as the calling agent is parked via + ``mark_waiting_long(waiting_for='rate_limit_lifted')``. A + ``RATE_LIMIT_HIT`` event is published so downstream consumers (the + orchestrator backpressure layer, the panel) can react. + """ + from roboco.models.events import Event, EventType + + agent_slug: str | None = ( + getattr(agent, "slug", None) if agent is not None else None + ) + + provider: str = "unknown" + affected_agents: list[str] = [] + + orch = self.orchestrator + if orch is not None and agent_slug is not None: + with contextlib.suppress(Exception): + prov = orch.get_provider_for_agent(agent_slug) + if prov: + provider = prov + with contextlib.suppress(Exception): + affected_agents = list( + orch.get_active_agent_slugs_for_provider(provider) + ) + for slug in affected_agents: + with contextlib.suppress(Exception): + await orch.mark_waiting_long( + slug, + waiting_for="rate_limit_lifted", + task_id=str(task_id), + context={"provider": provider, "triggered_by": agent_slug}, + ) + + retry_after_seconds = self._parse_retry_after(what_needed) + + # Persist rate-limit state to Redis so downstream decide_spawn() + # calls can gate new spawns for this provider. Skipped when the + # provider is "unknown" (orchestrator not wired or not tracking the + # agent) to avoid polluting the tracker with meaningless keys. + if provider != "unknown": + with contextlib.suppress(Exception): + from roboco.services.gateway.rate_limit_tracker import ( + RateLimitStateTracker, + ) + + await RateLimitStateTracker(provider).activate( + retry_after=retry_after_seconds, + affected_agents=affected_agents, + ) + + bus = self.stream_bus + if bus is not None: + event = Event( + type=EventType.RATE_LIMIT_HIT, + data={ + "provider": provider, + "affectedAgents": affected_agents, + "retryAfterSeconds": retry_after_seconds, + "timestamp": datetime.now(UTC).isoformat(), + }, + source_agent=str(agent_id), + ) + with contextlib.suppress(Exception): + await bus.publish(event) + + await self._touch(task_id) + return Envelope.ok( + status=str(t.status), + task_id=str(task_id), + next=( + "agent parked waiting for rate_limit_lifted; " + "will be respawned when the limit clears" + ), + context_briefing=briefing, + ).with_introspection(task=t, role=role_str) + async def i_am_blocked( self, agent_id: UUID, @@ -2202,8 +2329,15 @@ class Choreographer: membership (developer/qa/documenter) and the source-status constraint of the composed ``block`` action (in_progress only). After the spec gate accepts, the journal:struggle entry is written - from the verb body, then ``VerbRunner.run_intent("i_am_blocked", ...)`` - dispatches the (block,) atomic chain wrapped in a savepoint. + from the verb body, then either: + + - ``reason == 'rate_limited'``: the task is **not** transitioned to + ``blocked``; instead every active agent on the same provider is + parked via ``mark_waiting_long(waiting_for='rate_limit_lifted')`` + and a ``RATE_LIMIT_HIT`` event is published. + - any other reason: ``VerbRunner.run_intent("i_am_blocked", ...)`` + dispatches the ``(block,)`` atomic chain wrapped in a savepoint, + transitioning the task to ``blocked``. """ t = await self.task.get(task_id) if t is None: @@ -2253,6 +2387,20 @@ class Choreographer: task_id=task_id, content=self._build_struggle_body(reason, blocker_type, what_needed), ) + + # Rate-limited fast path: skip the block state transition and park + # all affected agents instead. + if reason.strip().lower() == "rate_limited": + return await self._handle_rate_limited_parking( + agent_id=agent_id, + task_id=task_id, + t=t, + agent=agent, + role_str=role_str, + briefing=briefing, + what_needed=what_needed, + ) + t, rejection = await self._run_i_am_blocked_intent( agent_id, task_id, t, agent, spec_ctx, role_str, briefing ) diff --git a/roboco/services/gateway/rate_limit_tracker.py b/roboco/services/gateway/rate_limit_tracker.py new file mode 100644 index 00000000..047cae60 --- /dev/null +++ b/roboco/services/gateway/rate_limit_tracker.py @@ -0,0 +1,185 @@ +"""Redis-backed rate-limit state tracker for the agent gateway. + +State is persisted in Redis as a JSON blob keyed by provider name. +Because it is backed by Redis rather than process memory, state survives +a process restart and a *new* ``RateLimitStateTracker`` instance pointing +at the same Redis URL will read the same values — satisfying the +cross-reconnection persistence requirement. +""" + +from __future__ import annotations + +import json +from datetime import UTC, datetime +from typing import Any + +import redis.asyncio as redis + +from roboco.config import settings + + +class RateLimitStateTracker: + """Track rate-limit state for a single AI provider in Redis. + + Usage + ----- + tracker = RateLimitStateTracker("anthropic") + await tracker.activate(retry_after=60.0, affected_agents=["be-dev-1"]) + assert await tracker.is_rate_limited() + + A second instance that uses the same Redis URL and provider name + will observe the same state — no in-process singleton required. + """ + + _KEY_PREFIX: str = "roboco:rate_limit:" + + def __init__(self, provider: str, redis_url: str | None = None) -> None: + """Construct a tracker for *provider*. + + Args: + provider: Logical provider name, e.g. ``"anthropic"`` or + ``"ollama_cloud"``. Used as part of the Redis key. + redis_url: Override the Redis URL (defaults to + ``settings.redis_url``). + """ + self._provider = provider + self._redis_url = redis_url or settings.redis_url + self._redis: redis.Redis | None = None # type: ignore[type-arg] + + # ------------------------------------------------------------------ + # Private helpers + # ------------------------------------------------------------------ + + async def _conn(self) -> redis.Redis: # type: ignore[type-arg] + """Return a (lazy-connected) redis.asyncio.Redis client.""" + if self._redis is None: + self._redis = redis.from_url(self._redis_url) + return self._redis + + def _key(self) -> str: + """Redis key for this provider's state blob.""" + return f"{self._KEY_PREFIX}{self._provider}:state" + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + async def activate( + self, + retry_after: float | None = None, + affected_agents: list[str] | None = None, + ) -> None: + """Mark the provider as rate-limited. + + Args: + retry_after: Seconds until the provider should accept new + requests, or ``None`` if unknown. + affected_agents: Agent slugs that were active when the limit + was hit (informational; stored in state). + """ + r = await self._conn() + state: dict[str, Any] = { + "rate_limited": True, + "activated_at": datetime.now(UTC).isoformat(), + "retry_after": retry_after, + "affected_agents": affected_agents or [], + "probe_failures": 0, + } + await r.set(self._key(), json.dumps(state)) + + async def clear(self) -> None: + """Remove rate-limit state for this provider.""" + r = await self._conn() + await r.delete(self._key()) + + async def is_rate_limited(self) -> bool: + """Return ``True`` if the provider is currently rate-limited.""" + state = await self.get_state() + return bool(state.get("rate_limited", False)) + + async def get_state(self) -> dict[str, Any]: + """Return the stored state dict, or ``{}`` if none exists.""" + r = await self._conn() + raw = await r.get(self._key()) + if raw is None: + return {} + decoded: str = raw.decode() if isinstance(raw, bytes) else str(raw) + result: dict[str, Any] = json.loads(decoded) + return result + + async def increment_probe_failures(self) -> int: + """Increment the probe-failure counter and return the new value. + + The probe-failure counter tracks how many successive connectivity + probes have failed since the rate limit was activated. The + orchestrator uses this to decide whether to keep waiting or give + up entirely. + """ + r = await self._conn() + state = await self.get_state() + new_count: int = state.get("probe_failures", 0) + 1 + state["probe_failures"] = new_count + await r.set(self._key(), json.dumps(state)) + return new_count + + async def reset_probe_failures(self) -> None: + """Reset the probe-failure counter to 0.""" + r = await self._conn() + state = await self.get_state() + state["probe_failures"] = 0 + await r.set(self._key(), json.dumps(state)) + + # ------------------------------------------------------------------ + # Class-level helpers + # ------------------------------------------------------------------ + + @classmethod + async def list_rate_limited_providers( + cls, + redis_url: str | None = None, + ) -> list[tuple[str, dict[str, Any]]]: + """Scan Redis for all providers that are currently rate-limited. + + Returns a list of ``(provider_name, state_dict)`` tuples — one + entry per provider whose stored state has ``rate_limited == True``. + Returns an empty list when nothing is rate-limited or Redis is + unreachable. + + Args: + redis_url: Override the default Redis URL from settings. + """ + url = redis_url or settings.redis_url + r: redis.Redis = redis.from_url(url) # type: ignore[type-arg] + pattern = f"{cls._KEY_PREFIX}*:state" + results: list[tuple[str, dict[str, Any]]] = [] + try: + cursor: int = 0 + while True: + cursor, keys = await r.scan(cursor, match=pattern, count=100) + for raw_key in keys: + key: str = ( + raw_key.decode() if isinstance(raw_key, bytes) else str(raw_key) + ) + # Extract provider from key: roboco:rate_limit:{provider}:state + # Strip prefix and suffix + inner = key[len(cls._KEY_PREFIX) :] + if inner.endswith(":state"): + provider = inner[: -len(":state")] + else: + continue + raw_val = await r.get(key) + if raw_val is None: + continue + decoded: str = ( + raw_val.decode() if isinstance(raw_val, bytes) else str(raw_val) + ) + state: dict[str, Any] = json.loads(decoded) + if state.get("rate_limited"): + results.append((provider, state)) + if cursor == 0: + break + except Exception: + pass + finally: + await r.aclose() + return results diff --git a/roboco/services/gateway/trigger_filter.py b/roboco/services/gateway/trigger_filter.py index af5a4b9c..05984d8a 100644 --- a/roboco/services/gateway/trigger_filter.py +++ b/roboco/services/gateway/trigger_filter.py @@ -50,6 +50,11 @@ class TriggerContext: skill: str | None recent_spawns_for_task: int recent_spawns_for_role: int + # Provider rate-limit fields. Optional — callers that don't know the + # provider (e.g. no-task spawns) leave these at their defaults so the + # gate is a no-op. + provider: str | None = None + provider_rate_limited: bool = False _TERMINAL_STATUSES: frozenset[str] = frozenset({"completed", "cancelled"}) @@ -60,13 +65,16 @@ _A2A_CODE_REVIEW_RELEVANT_STATES: frozenset[str] = frozenset( ) -def decide_spawn( +def decide_spawn( # noqa: PLR0911 *, task: Any, trigger: TriggerContext, config: SpawnConfig, ) -> Decision: - """Apply four rules in order: stale > claimant-lock > task-cooldown > role-rate.""" + """Apply five rules in order. + + stale > provider-rate-limit > claimant-lock > task-cooldown > role-rate + """ # 1. Stale-trigger cleanup if task.status in _TERMINAL_STATUSES: return Decision(SpawnDecision.DROP, "task in terminal state — trigger stale") @@ -81,7 +89,14 @@ def decide_spawn( f"a2a code_review for task in {task.status} — stale", ) - # 2. Single-claimant invariant + # 2. Provider rate-limit gate + if trigger.provider_rate_limited: + return Decision( + SpawnDecision.QUEUE, + f"provider {trigger.provider or 'unknown'} rate-limited", + ) + + # 3. Single-claimant invariant if task.active_claimant_id is not None and not is_stale( task, threshold_seconds=config.claim_stale_seconds ): @@ -90,14 +105,14 @@ def decide_spawn( "task has active claimant with fresh heartbeat", ) - # 3. Per-task spawn cooldown + # 4. Per-task spawn cooldown if trigger.recent_spawns_for_task >= 1: return Decision( SpawnDecision.QUEUE, f"per-task spawn cooldown ({config.cooldown_seconds}s) active", ) - # 4. Per-role rate limit + # 5. Per-role rate limit if trigger.recent_spawns_for_role >= config.role_rate_per_minute: return Decision( SpawnDecision.QUEUE, diff --git a/roboco/services/optimal_brain/indexes/base.py b/roboco/services/optimal_brain/indexes/base.py index 0bfc3401..8d8f867e 100644 --- a/roboco/services/optimal_brain/indexes/base.py +++ b/roboco/services/optimal_brain/indexes/base.py @@ -17,6 +17,12 @@ from piragi.types import Citation, Document from roboco.config import settings from roboco.models.optimal import IndexType, SearchOutcome, SearchResult +from roboco.services.exceptions import ( + HTTP_TOO_MANY_REQUESTS, + MAX_RATE_LIMIT_RETRIES, + RateLimitError, + parse_retry_after_header, +) # Apply piragi runtime patches (chunker tokenizer) BEFORE importing piragi # itself anywhere in the plugin stack. Importing for side effects only. @@ -918,19 +924,25 @@ class BaseIndexPlugin(ABC): Returns: Tuple of (answer, citations). Returns ("", []) on failure to allow OptimalService to continue to next index. + + The Ollama LLM call is retried up to MAX_RATE_LIMIT_RETRIES times on + HTTP 429, respecting the Retry-After header. The vector-search phase + is NOT retried and still runs inside the 15-second index timeout. """ import asyncio import httpx - # Per-index timeout to prevent one slow index from blocking everything + # Per-index timeout to prevent one slow index from blocking everything. + # Applied to the search phase only; LLM retries run outside this timeout. INDEX_TIMEOUT = 15.0 search_results: list[SearchResult] = [] + prompt: str = "" + # ---- Search phase (inside timeout) ----------------------------------- try: async with asyncio.timeout(INDEX_TIMEOUT): - # First get context using our properly async search logger.info( "ask() starting search", index_type=self.index_type.value, @@ -947,14 +959,11 @@ class BaseIndexPlugin(ABC): ) if not search_results: - # Return empty to continue to next index return "", [] - # Build context for LLM + # Build context and prompt while still inside the timeout context_texts = [r.content for r in search_results] context = "\n\n---\n\n".join(context_texts) - - # Build prompt prompt = ( "You are a technical knowledge base assistant. " "Based on the context, provide a thorough, actionable answer.\n\n" @@ -969,14 +978,28 @@ class BaseIndexPlugin(ABC): "Detailed Answer:" ) - # Call LLM via Ollama API (async HTTP) - llm_url = f"{self.config.llm_base_url}/chat/completions" - logger.info( - "ask() calling LLM", - index_type=self.index_type.value, - llm_url=llm_url, - model=self.config.llm_model, - ) + except (TimeoutError, httpx.TimeoutException, Exception) as e: + logger.warning( + "Index ask() search phase failed", + index_type=self.index_type.value, + error_type=type(e).__name__, + error=str(e) if not isinstance(e, TimeoutError) else "timed out", + ) + return "", search_results + + # ---- LLM call phase with 429 retry (outside index timeout) ----------- + llm_url = f"{self.config.llm_base_url}/chat/completions" + logger.info( + "ask() calling LLM", + index_type=self.index_type.value, + llm_url=llm_url, + model=self.config.llm_model, + ) + + last_rl_retry_after: float | None = None + + for rl_attempt in range(MAX_RATE_LIMIT_RETRIES): + try: async with httpx.AsyncClient(timeout=60.0) as client: resp = await client.post( llm_url, @@ -987,44 +1010,46 @@ class BaseIndexPlugin(ABC): "options": {"num_ctx": 8192}, }, ) - if resp.is_success: - data = resp.json() - answer_text = data["choices"][0]["message"]["content"] - # Extract answer from think tags if needed - answer_text = self._extract_from_think_tags(answer_text) - return answer_text, search_results - else: - logger.warning( - "LLM call failed in ask", - index_type=self.index_type.value, - status=resp.status_code, - error=resp.text[:200] if resp.text else "no error text", - ) - # Return empty to let service aggregate and synthesize - return "", search_results + except (httpx.TimeoutException, Exception) as e: + logger.warning( + "LLM call failed in ask (non-429)", + index_type=self.index_type.value, + error=str(e), + ) + return "", search_results - except TimeoutError: - logger.warning( - "Index ask() timed out", - index_type=self.index_type.value, - timeout=INDEX_TIMEOUT, - ) - # Return search results even on timeout - service can aggregate them - return "", search_results - except httpx.TimeoutException: - logger.warning( - "LLM HTTP call timed out in ask", - index_type=self.index_type.value, - ) - return "", search_results - except Exception as e: - logger.warning( - "RAG query failed", - index_type=self.index_type.value, - error=str(e), - ) - # Return whatever search results we have for aggregation - return "", search_results + if resp.status_code == HTTP_TOO_MANY_REQUESTS: + retry_after = parse_retry_after_header(resp) + last_rl_retry_after = retry_after + backoff = ( + retry_after if retry_after is not None else float(2**rl_attempt) + ) + logger.warning( + "Ollama rate limited (429), retrying", + provider="ollama", + attempt=rl_attempt + 1, + max_retries=MAX_RATE_LIMIT_RETRIES, + backoff_duration=backoff, + ) + if rl_attempt < MAX_RATE_LIMIT_RETRIES - 1: + await asyncio.sleep(backoff) + continue + + if resp.is_success: + data = resp.json() + answer_text = data["choices"][0]["message"]["content"] + answer_text = self._extract_from_think_tags(answer_text) + return answer_text, search_results + else: + logger.warning( + "LLM call failed in ask", + index_type=self.index_type.value, + status=resp.status_code, + error=resp.text[:200] if resp.text else "no error text", + ) + return "", search_results + + raise RateLimitError(provider="ollama", retry_after=last_rl_retry_after) async def count(self) -> int: """Get the number of documents in the index.""" diff --git a/roboco/services/optimal_brain/mentor.py b/roboco/services/optimal_brain/mentor.py index 43ae5e9f..bb45c6c4 100644 --- a/roboco/services/optimal_brain/mentor.py +++ b/roboco/services/optimal_brain/mentor.py @@ -32,6 +32,12 @@ from roboco.models.optimal import ( MentorResponse, SearchResult, ) +from roboco.services.exceptions import ( + HTTP_TOO_MANY_REQUESTS, + MAX_RATE_LIMIT_RETRIES, + RateLimitError, + parse_retry_after_header, +) logger = structlog.get_logger() @@ -683,7 +689,12 @@ class MentorService: agent_profile: AgentProfile | None, journal_context: list[dict[str, Any]], ) -> str: - """Synthesize a personalized answer using LLM.""" + """Synthesize a personalized answer using LLM. + + Retries the Ollama call up to MAX_RATE_LIMIT_RETRIES times on HTTP 429, + respecting the Retry-After header. Each individual attempt is bounded by + a 120-second asyncio timeout so a slow model cannot block the loop. + """ import asyncio if not sources and not journal_context: @@ -709,52 +720,73 @@ class MentorService: question, sources, conversation_context, agent_profile, journal_context ) - # Call LLM - try: - async with asyncio.timeout(120.0): - async with httpx.AsyncClient(timeout=120.0) as client: - response = await client.post( - f"{settings.local_llm_base_url}/chat/completions", - json={ - "model": settings.local_llm_model, - "messages": [ - {"role": "system", "content": system_prompt}, - {"role": "user", "content": user_prompt}, - ], - "max_tokens": 4096, - "temperature": 0.5, - "options": {"num_ctx": 8192}, - }, + llm_url = f"{settings.local_llm_base_url}/chat/completions" + payload = { + "model": settings.local_llm_model, + "messages": [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": user_prompt}, + ], + "max_tokens": 4096, + "temperature": 0.5, + "options": {"num_ctx": 8192}, + } + + last_rl_retry_after: float | None = None + + for rl_attempt in range(MAX_RATE_LIMIT_RETRIES): + # Each attempt gets its own 120-second timeout + try: + async with asyncio.timeout(120.0): + async with httpx.AsyncClient(timeout=120.0) as client: + response = await client.post(llm_url, json=payload) + except (TimeoutError, httpx.TimeoutException): + logger.warning("LLM call timed out in mentor (120s)") + return self._fallback_answer(sources, agent_profile) + except Exception as e: + logger.warning("LLM call failed in mentor", error=str(e)) + return self._fallback_answer(sources, agent_profile) + + if response.status_code == HTTP_TOO_MANY_REQUESTS: + retry_after = parse_retry_after_header(response) + last_rl_retry_after = retry_after + backoff = ( + retry_after if retry_after is not None else float(2**rl_attempt) + ) + logger.warning( + "Ollama rate limited (429), retrying", + provider="ollama", + attempt=rl_attempt + 1, + max_retries=MAX_RATE_LIMIT_RETRIES, + backoff_duration=backoff, + ) + if rl_attempt < MAX_RATE_LIMIT_RETRIES - 1: + await asyncio.sleep(backoff) + continue + + if response.is_success: + data = response.json() + raw_answer: str = data["choices"][0]["message"]["content"] + answer = self._extract_answer(raw_answer) + + if not answer: + logger.warning( + "LLM response empty after extraction", + model=settings.local_llm_model, + original_length=len(raw_answer), ) + return self._fallback_answer(sources, agent_profile) - if response.is_success: - data = response.json() - raw_answer: str = data["choices"][0]["message"]["content"] - answer = self._extract_answer(raw_answer) + return answer + else: + logger.warning( + "LLM call failed in mentor", + status=response.status_code, + error=response.text[:200], + ) + return self._fallback_answer(sources, agent_profile) - if not answer: - logger.warning( - "LLM response empty after extraction", - model=settings.local_llm_model, - original_length=len(raw_answer), - ) - return self._fallback_answer(sources, agent_profile) - - return answer - else: - logger.warning( - "LLM call failed in mentor", - status=response.status_code, - error=response.text[:200], - ) - return self._fallback_answer(sources, agent_profile) - - except (TimeoutError, httpx.TimeoutException): - logger.warning("LLM call timed out in mentor (60s)") - return self._fallback_answer(sources, agent_profile) - except Exception as e: - logger.warning("LLM call failed in mentor", error=str(e)) - return self._fallback_answer(sources, agent_profile) + raise RateLimitError(provider="ollama", retry_after=last_rl_retry_after) def _extract_answer(self, text: str) -> str: """Extract answer from LLM response, handling think tags.""" diff --git a/roboco/services/optimal_brain/ollama_embedder.py b/roboco/services/optimal_brain/ollama_embedder.py index c989987a..016be222 100644 --- a/roboco/services/optimal_brain/ollama_embedder.py +++ b/roboco/services/optimal_brain/ollama_embedder.py @@ -22,13 +22,21 @@ from piragi.types import Chunk from roboco.config import settings from roboco.logging import get_logger +from roboco.services.exceptions import ( + HTTP_TOO_MANY_REQUESTS, + RateLimitError, + parse_retry_after_header, +) logger = get_logger(__name__) -# Retry configuration +# Retry configuration — ConnectError / Timeout (existing, unchanged) MAX_RETRIES = 3 RETRY_DELAY_BASE = 0.5 # seconds, exponential backoff +# Retry configuration — HTTP 429 / RateLimitError (new outer loop) +RATE_LIMIT_MAX_RETRIES = 5 + # Parallel processing configuration MAX_CONCURRENT_BATCHES = 4 # Number of batches to process in parallel DEFAULT_BATCH_SIZE = 32 # piragi's default batch size @@ -304,7 +312,15 @@ class OllamaEmbedder: query: str, task_instruction: str | None = None, ) -> list[float]: - """Generate embedding for a single query with retry logic.""" + """Generate embedding for a single query. + + Retry behaviour (two independent concerns, non-overlapping): + - ConnectError / TimeoutException: up to MAX_RETRIES=3 attempts with + 0.5/1/2 s exponential backoff (existing behaviour, unchanged). + - HTTP 429 (rate limit): outer loop up to RATE_LIMIT_MAX_RETRIES=5, + respecting Retry-After header. A 429 response does NOT trigger the + ConnectError path. + """ _ = task_instruction # Check cache first @@ -313,79 +329,154 @@ class OllamaEmbedder: return cached client = self._get_sync_client() - last_error: Exception | None = None + last_rl_retry_after: float | None = None - for attempt in range(MAX_RETRIES): - try: - response = client.post( - f"{self.base_url}/api/embed", - json={"model": self.model, "input": query}, - ) - embeddings = self._handle_embed_response(response, input_count=1) - result = embeddings[0] - self._cache.put(query, result) - return result + for rl_attempt in range(RATE_LIMIT_MAX_RETRIES): + got_429 = False + last_error: Exception | None = None - except httpx.ConnectError as e: - last_error = OllamaConnectionError( - f"Cannot connect to Ollama at {self.base_url}: {e}" - ) - except httpx.TimeoutException as e: - last_error = OllamaConnectionError(f"Ollama request timed out: {e}") - except (OllamaModelError, OllamaEmbedderError): - raise - except Exception as e: - last_error = OllamaEmbedderError(f"Unexpected error: {e}") + # --- inner loop: ConnectError / Timeout (unchanged) --- + for attempt in range(MAX_RETRIES): + try: + response = client.post( + f"{self.base_url}/api/embed", + json={"model": self.model, "input": query}, + ) + # 429 check — must NOT enter the ConnectError path + if response.status_code == HTTP_TOO_MANY_REQUESTS: + retry_after = parse_retry_after_header(response) + last_rl_retry_after = retry_after + backoff = ( + retry_after + if retry_after is not None + else float(2**rl_attempt) + ) + logger.warning( + "Ollama rate limited (429), retrying", + provider="ollama", + attempt=rl_attempt + 1, + max_retries=RATE_LIMIT_MAX_RETRIES, + backoff_duration=backoff, + ) + got_429 = True + break # break inner loop; outer loop will sleep + retry - if attempt < MAX_RETRIES - 1: - delay = RETRY_DELAY_BASE * (2**attempt) - logger.warning( - "Ollama embed_query retry", - attempt=attempt + 1, - delay=delay, - error=str(last_error), - ) - time.sleep(delay) + embeddings = self._handle_embed_response(response, input_count=1) + result = embeddings[0] + self._cache.put(query, result) + return result - raise last_error or OllamaEmbedderError("Max retries exceeded") + except httpx.ConnectError as e: + last_error = OllamaConnectionError( + f"Cannot connect to Ollama at {self.base_url}: {e}" + ) + except httpx.TimeoutException as e: + last_error = OllamaConnectionError(f"Ollama request timed out: {e}") + except (OllamaModelError, OllamaEmbedderError): + raise + except Exception as e: + last_error = OllamaEmbedderError(f"Unexpected error: {e}") + + if attempt < MAX_RETRIES - 1: + delay = RETRY_DELAY_BASE * (2**attempt) + logger.warning( + "Ollama embed_query retry", + attempt=attempt + 1, + delay=delay, + error=str(last_error), + ) + time.sleep(delay) + # --- end inner loop --- + + if not got_429: + # ConnectError / Timeout exhausted — same behaviour as before + raise last_error or OllamaEmbedderError("Max retries exceeded") + + # 429: sleep and try again (outer loop) + backoff = ( + last_rl_retry_after + if last_rl_retry_after is not None + else float(2**rl_attempt) + ) + if rl_attempt < RATE_LIMIT_MAX_RETRIES - 1: + time.sleep(backoff) + + raise RateLimitError(provider="ollama", retry_after=last_rl_retry_after) def _embed_batch_sync( self, client: httpx.Client, batch: list[str], batch_index: int ) -> list[list[float]]: - """Embed a single batch synchronously with retry logic.""" - last_error: Exception | None = None + """Embed a single batch synchronously. - for attempt in range(MAX_RETRIES): - try: - response = client.post( - f"{self.base_url}/api/embed", - json={"model": self.model, "input": batch}, - ) - return self._handle_embed_response(response, input_count=len(batch)) + Same two-concern retry composition as :meth:`embed_query`: + inner ConnectError/Timeout loop (unchanged) + outer 429 loop. + """ + last_rl_retry_after: float | None = None - except httpx.ConnectError as e: - last_error = OllamaConnectionError( - f"Cannot connect to Ollama at {self.base_url}: {e}" - ) - except httpx.TimeoutException as e: - last_error = OllamaConnectionError(f"Ollama request timed out: {e}") - except (OllamaModelError, OllamaEmbedderError): - raise - except Exception as e: - last_error = OllamaEmbedderError(f"Unexpected error: {e}") + for rl_attempt in range(RATE_LIMIT_MAX_RETRIES): + got_429 = False + last_error: Exception | None = None - if attempt < MAX_RETRIES - 1: - delay = RETRY_DELAY_BASE * (2**attempt) - logger.warning( - "Ollama embed_documents retry", - attempt=attempt + 1, - batch_index=batch_index, - delay=delay, - error=str(last_error), - ) - time.sleep(delay) + for attempt in range(MAX_RETRIES): + try: + response = client.post( + f"{self.base_url}/api/embed", + json={"model": self.model, "input": batch}, + ) + if response.status_code == HTTP_TOO_MANY_REQUESTS: + retry_after = parse_retry_after_header(response) + last_rl_retry_after = retry_after + backoff = ( + retry_after + if retry_after is not None + else float(2**rl_attempt) + ) + logger.warning( + "Ollama rate limited (429), retrying", + provider="ollama", + attempt=rl_attempt + 1, + max_retries=RATE_LIMIT_MAX_RETRIES, + backoff_duration=backoff, + ) + got_429 = True + break - raise last_error or OllamaEmbedderError("Max retries exceeded") + return self._handle_embed_response(response, input_count=len(batch)) + + except httpx.ConnectError as e: + last_error = OllamaConnectionError( + f"Cannot connect to Ollama at {self.base_url}: {e}" + ) + except httpx.TimeoutException as e: + last_error = OllamaConnectionError(f"Ollama request timed out: {e}") + except (OllamaModelError, OllamaEmbedderError): + raise + except Exception as e: + last_error = OllamaEmbedderError(f"Unexpected error: {e}") + + if attempt < MAX_RETRIES - 1: + delay = RETRY_DELAY_BASE * (2**attempt) + logger.warning( + "Ollama embed_documents retry", + attempt=attempt + 1, + batch_index=batch_index, + delay=delay, + error=str(last_error), + ) + time.sleep(delay) + + if not got_429: + raise last_error or OllamaEmbedderError("Max retries exceeded") + + backoff = ( + last_rl_retry_after + if last_rl_retry_after is not None + else float(2**rl_attempt) + ) + if rl_attempt < RATE_LIMIT_MAX_RETRIES - 1: + time.sleep(backoff) + + raise RateLimitError(provider="ollama", retry_after=last_rl_retry_after) def _partition_cached_documents( self, documents: list[str] @@ -464,49 +555,90 @@ class OllamaEmbedder: batch: list[str], batch_index: int, ) -> list[list[float]]: - """Embed a single batch with semaphore-limited concurrency.""" + """Embed a single batch with semaphore-limited concurrency. + + Same two-concern retry composition as :meth:`embed_query`: + inner ConnectError/Timeout loop (unchanged) + outer 429 loop (async). + """ semaphore = self._get_semaphore() async with semaphore: - last_error: Exception | None = None + last_rl_retry_after: float | None = None - for attempt in range(MAX_RETRIES): - try: - logger.debug( - "Parallel embed batch", - batch_index=batch_index, - batch_size=len(batch), - attempt=attempt, - ) - response = await client.post( - f"{self.base_url}/api/embed", - json={"model": self.model, "input": batch}, - ) - return self._handle_embed_response(response, input_count=len(batch)) + for rl_attempt in range(RATE_LIMIT_MAX_RETRIES): + got_429 = False + last_error: Exception | None = None - except httpx.ConnectError as e: - last_error = OllamaConnectionError( - f"Cannot connect to Ollama at {self.base_url}: {e}" - ) - except httpx.TimeoutException as e: - last_error = OllamaConnectionError(f"Ollama request timed out: {e}") - except (OllamaModelError, OllamaEmbedderError): - raise - except Exception as e: - last_error = OllamaEmbedderError(f"Unexpected error: {e}") + for attempt in range(MAX_RETRIES): + try: + logger.debug( + "Parallel embed batch", + batch_index=batch_index, + batch_size=len(batch), + attempt=attempt, + ) + response = await client.post( + f"{self.base_url}/api/embed", + json={"model": self.model, "input": batch}, + ) + if response.status_code == HTTP_TOO_MANY_REQUESTS: + retry_after = parse_retry_after_header(response) + last_rl_retry_after = retry_after + backoff = ( + retry_after + if retry_after is not None + else float(2**rl_attempt) + ) + logger.warning( + "Ollama rate limited (429), retrying", + provider="ollama", + attempt=rl_attempt + 1, + max_retries=RATE_LIMIT_MAX_RETRIES, + backoff_duration=backoff, + ) + got_429 = True + break - if attempt < MAX_RETRIES - 1: - delay = RETRY_DELAY_BASE * (2**attempt) - logger.warning( - "Parallel embed batch retry", - batch_index=batch_index, - attempt=attempt + 1, - delay=delay, - error=str(last_error), - ) - await asyncio.sleep(delay) + return self._handle_embed_response( + response, input_count=len(batch) + ) - raise last_error or OllamaEmbedderError("Max retries exceeded") + except httpx.ConnectError as e: + last_error = OllamaConnectionError( + f"Cannot connect to Ollama at {self.base_url}: {e}" + ) + except httpx.TimeoutException as e: + last_error = OllamaConnectionError( + f"Ollama request timed out: {e}" + ) + except (OllamaModelError, OllamaEmbedderError): + raise + except Exception as e: + last_error = OllamaEmbedderError(f"Unexpected error: {e}") + + if attempt < MAX_RETRIES - 1: + delay = RETRY_DELAY_BASE * (2**attempt) + logger.warning( + "Parallel embed batch retry", + batch_index=batch_index, + attempt=attempt + 1, + delay=delay, + error=str(last_error), + ) + await asyncio.sleep(delay) + + if not got_429: + raise last_error or OllamaEmbedderError("Max retries exceeded") + + backoff = ( + last_rl_retry_after + if last_rl_retry_after is not None + else float(2**rl_attempt) + ) + if rl_attempt < RATE_LIMIT_MAX_RETRIES - 1: + await asyncio.sleep(backoff) + + raise RateLimitError(provider="ollama", retry_after=last_rl_retry_after) async def _run_parallel_batches( self, batches: list[list[str]] @@ -650,49 +782,93 @@ class OllamaEmbedder: return chunks async def aembed_query(self, query: str) -> list[float]: - """Async version of embed_query with retry logic and caching.""" + """Async version of embed_query with retry logic and caching. + + Same two-concern retry composition as :meth:`embed_query`: + inner ConnectError/Timeout loop (unchanged) + outer 429 loop (async). + """ # Check cache first cached = self._cache.get(query) if cached is not None: return cached - last_error: Exception | None = None + last_rl_retry_after: float | None = None - for attempt in range(MAX_RETRIES): - # Create fresh client each attempt to avoid event loop issues - async with self._create_async_client() as client: - try: - response = await client.post( - f"{self.base_url}/api/embed", - json={"model": self.model, "input": query}, + for rl_attempt in range(RATE_LIMIT_MAX_RETRIES): + got_429 = False + last_error: Exception | None = None + + for attempt in range(MAX_RETRIES): + # Create fresh client each attempt to avoid event loop issues + async with self._create_async_client() as client: + try: + response = await client.post( + f"{self.base_url}/api/embed", + json={"model": self.model, "input": query}, + ) + if response.status_code == HTTP_TOO_MANY_REQUESTS: + retry_after = parse_retry_after_header(response) + last_rl_retry_after = retry_after + backoff = ( + retry_after + if retry_after is not None + else float(2**rl_attempt) + ) + logger.warning( + "Ollama rate limited (429), retrying", + provider="ollama", + attempt=rl_attempt + 1, + max_retries=RATE_LIMIT_MAX_RETRIES, + backoff_duration=backoff, + ) + got_429 = True + break + + embeddings = self._handle_embed_response( + response, input_count=1 + ) + result = embeddings[0] + self._cache.put(query, result) + return result + + except httpx.ConnectError as e: + last_error = OllamaConnectionError( + f"Cannot connect to Ollama at {self.base_url}: {e}" + ) + except httpx.TimeoutException as e: + last_error = OllamaConnectionError( + f"Ollama request timed out: {e}" + ) + except (OllamaModelError, OllamaEmbedderError): + raise + except Exception as e: + last_error = OllamaEmbedderError(f"Unexpected error: {e}") + + if got_429: + break # exit inner loop cleanly + + if attempt < MAX_RETRIES - 1: + delay = RETRY_DELAY_BASE * (2**attempt) + logger.warning( + "Ollama aembed_query retry", + attempt=attempt + 1, + delay=delay, + error=str(last_error), ) - embeddings = self._handle_embed_response(response, input_count=1) - result = embeddings[0] - self._cache.put(query, result) - return result + await asyncio.sleep(delay) - except httpx.ConnectError as e: - last_error = OllamaConnectionError( - f"Cannot connect to Ollama at {self.base_url}: {e}" - ) - except httpx.TimeoutException as e: - last_error = OllamaConnectionError(f"Ollama request timed out: {e}") - except (OllamaModelError, OllamaEmbedderError): - raise - except Exception as e: - last_error = OllamaEmbedderError(f"Unexpected error: {e}") + if not got_429: + raise last_error or OllamaEmbedderError("Max retries exceeded") - if attempt < MAX_RETRIES - 1: - delay = RETRY_DELAY_BASE * (2**attempt) - logger.warning( - "Ollama aembed_query retry", - attempt=attempt + 1, - delay=delay, - error=str(last_error), - ) - await asyncio.sleep(delay) + backoff = ( + last_rl_retry_after + if last_rl_retry_after is not None + else float(2**rl_attempt) + ) + if rl_attempt < RATE_LIMIT_MAX_RETRIES - 1: + await asyncio.sleep(backoff) - raise last_error or OllamaEmbedderError("Max retries exceeded") + raise RateLimitError(provider="ollama", retry_after=last_rl_retry_after) async def aembed_documents( self, documents: list[str], batch_size: int = DEFAULT_BATCH_SIZE diff --git a/roboco/services/optimal_brain/validator.py b/roboco/services/optimal_brain/validator.py index e097b0df..3ce45349 100644 --- a/roboco/services/optimal_brain/validator.py +++ b/roboco/services/optimal_brain/validator.py @@ -19,6 +19,12 @@ import structlog from roboco.config import settings from roboco.models.optimal import SearchResult, ValidationResult +from roboco.services.exceptions import ( + HTTP_TOO_MANY_REQUESTS, + MAX_RATE_LIMIT_RETRIES, + RateLimitError, + parse_retry_after_header, +) logger = structlog.get_logger() @@ -492,6 +498,10 @@ class ValidatorService: """ Use LLM to validate context against standards. + Retries the Ollama call up to MAX_RATE_LIMIT_RETRIES times on HTTP 429, + respecting the Retry-After header. Each attempt is bounded by + LLM_TIMEOUT_SECONDS so a single slow call cannot block the retry loop. + Returns: Tuple of (violations, warnings) """ @@ -512,41 +522,63 @@ RELEVANT STANDARDS: Analyze the context and identify any violations of the standards above. Return your analysis as JSON.""" - try: - async with asyncio.timeout(LLM_TIMEOUT_SECONDS): - async with httpx.AsyncClient(timeout=LLM_TIMEOUT_SECONDS) as client: - response = await client.post( - f"{settings.local_llm_base_url}/chat/completions", - json={ - "model": settings.local_llm_model, - "messages": [ - {"role": "system", "content": VALIDATION_SYSTEM_PROMPT}, - {"role": "user", "content": user_prompt}, - ], - "max_tokens": 2048, - "temperature": 0.1, # Low temp for consistent analysis - "options": {"num_ctx": 8192}, - }, - ) + llm_url = f"{settings.local_llm_base_url}/chat/completions" + payload = { + "model": settings.local_llm_model, + "messages": [ + {"role": "system", "content": VALIDATION_SYSTEM_PROMPT}, + {"role": "user", "content": user_prompt}, + ], + "max_tokens": 2048, + "temperature": 0.1, # Low temp for consistent analysis + "options": {"num_ctx": 8192}, + } - if response.is_success: - data = response.json() - raw_response = data["choices"][0]["message"]["content"] - return self._parse_llm_response(raw_response) - else: - logger.warning( - "LLM validation call failed", - status=response.status_code, - error=response.text[:200], - ) - raise RuntimeError(f"LLM call failed: {response.status_code}") + last_rl_retry_after: float | None = None - except TimeoutError: - logger.warning("LLM validation timed out") - raise - except httpx.TimeoutException: - logger.warning("LLM validation HTTP timeout") - raise + for rl_attempt in range(MAX_RATE_LIMIT_RETRIES): + # Each attempt gets its own timeout — raises if the call hangs + try: + async with asyncio.timeout(LLM_TIMEOUT_SECONDS): + async with httpx.AsyncClient(timeout=LLM_TIMEOUT_SECONDS) as client: + response = await client.post(llm_url, json=payload) + except TimeoutError: + logger.warning("LLM validation timed out") + raise + except httpx.TimeoutException: + logger.warning("LLM validation HTTP timeout") + raise + + if response.status_code == HTTP_TOO_MANY_REQUESTS: + retry_after = parse_retry_after_header(response) + last_rl_retry_after = retry_after + backoff = ( + retry_after if retry_after is not None else float(2**rl_attempt) + ) + logger.warning( + "Ollama rate limited (429), retrying", + provider="ollama", + attempt=rl_attempt + 1, + max_retries=MAX_RATE_LIMIT_RETRIES, + backoff_duration=backoff, + ) + if rl_attempt < MAX_RATE_LIMIT_RETRIES - 1: + await asyncio.sleep(backoff) + continue + + if response.is_success: + data = response.json() + raw_response = data["choices"][0]["message"]["content"] + return self._parse_llm_response(raw_response) + else: + logger.warning( + "LLM validation call failed", + status=response.status_code, + error=response.text[:200], + ) + raise RuntimeError(f"LLM call failed: {response.status_code}") + + raise RateLimitError(provider="ollama", retry_after=last_rl_retry_after) def _build_standards_context(self, standards: list[SearchResult]) -> str: """Build a formatted string of standards for the LLM.""" diff --git a/tests/unit/gateway/test_i_am_blocked_rate_limited.py b/tests/unit/gateway/test_i_am_blocked_rate_limited.py new file mode 100644 index 00000000..61caa912 --- /dev/null +++ b/tests/unit/gateway/test_i_am_blocked_rate_limited.py @@ -0,0 +1,508 @@ +"""Unit tests for the rate-limited path in Choreographer.i_am_blocked. + +Acceptance criteria verified here: +- AC1: i_am_blocked(reason='rate_limited') calls RateLimitStateTracker.activate() + and stores affected agent IDs; all active agents on the rate-limited + provider are subsequently marked waiting-long. +- AC3: POST /v1/i_am_blocked with reason='rate_limited' does NOT transition + the task to 'blocked'; the task remains in its current status + (in_progress) and the calling agent is parked via + mark_waiting_long(waiting_for='rate_limit_lifted'). +- AC4: mark_waiting_long is called for every orchestrator-tracked active agent + sharing the affected provider — call count equals active agent count. +- AC5: A RATE_LIMIT_HIT event is published to the StreamEventBus with fields + provider, affectedAgents, retryAfterSeconds, and timestamp. +""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock, patch +from uuid import uuid4 + +from roboco.models.events import EventType +from roboco.services.gateway.choreographer import Choreographer, ChoreographerDeps + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +_ACTIVE_AGENTS = ["be-dev-1", "be-dev-2", "be-qa"] +_PROVIDER = "anthropic" + + +def _make_evidence_repo() -> AsyncMock: + repo = AsyncMock() + for method in ( + "list_unread_a2a", + "list_unread_mentions", + "list_pending_notifications", + "task_metadata_gaps", + "recent_team_activity", + "blockers_in_lane", + "journal_highlights_for_task", + ): + getattr(repo, method).return_value = [] + return repo + + +def _make_task_svc(agent_id: object, task_id: object) -> AsyncMock: + t = MagicMock( + id=task_id, + status="in_progress", + assigned_to=agent_id, + pre_block_state=None, + task_type="code", + team="backend", + # Avoid issues with spec iteration in claim guards + dependency_ids=[], + # acceptance_criteria needed by some paths + acceptance_criteria=[], + quick_context=None, + ) + task_svc = AsyncMock() + task_svc.session = MagicMock() + task_svc.session.begin_nested = MagicMock( + return_value=MagicMock( + __aenter__=AsyncMock(return_value=None), + __aexit__=AsyncMock(return_value=False), + ) + ) + task_svc.get.return_value = t + task_svc.agent_for.return_value = MagicMock( + id=agent_id, + role="developer", + team="backend", + slug="be-dev-1", # calling agent's slug + ) + return task_svc + + +def _make_orchestrator( + active_agents: list[str] | None = None, + provider: str = _PROVIDER, +) -> MagicMock: + """Build a synchronous/async orchestrator mock.""" + agents = active_agents if active_agents is not None else _ACTIVE_AGENTS + orch = MagicMock() + orch.get_provider_for_agent = MagicMock(return_value=provider) + orch.get_active_agent_slugs_for_provider = MagicMock(return_value=agents) + orch.mark_waiting_long = AsyncMock(return_value=None) + return orch + + +def _make_stream_bus() -> AsyncMock: + bus = AsyncMock() + bus.publish = AsyncMock(return_value="msg-id-1") + return bus + + +def _make_deps( + agent_id: object, + task_id: object, + orchestrator: MagicMock | None = None, + stream_bus: AsyncMock | None = None, +) -> ChoreographerDeps: + return ChoreographerDeps( + task=_make_task_svc(agent_id, task_id), + work_session=AsyncMock(), + git=AsyncMock(), + a2a=AsyncMock(), + journal=AsyncMock(), + audit=AsyncMock(), + evidence_repo=_make_evidence_repo(), + orchestrator=orchestrator, + stream_bus=stream_bus, + ) + + +# --------------------------------------------------------------------------- +# AC3: Task stays in in_progress, agent parked via mark_waiting_long +# --------------------------------------------------------------------------- + + +class TestRateLimitedDoesNotBlockTask: + async def test_task_status_remains_in_progress(self) -> None: + """reason='rate_limited' must NOT transition the task to 'blocked'.""" + agent_id = uuid4() + task_id = uuid4() + orch = _make_orchestrator() + deps = _make_deps(agent_id, task_id, orchestrator=orch) + c = Choreographer(deps) + + env = await c.i_am_blocked(agent_id, task_id, "rate_limited") + + assert env.error is None + assert env.status == "in_progress" + + async def test_verb_runner_block_action_not_called(self) -> None: + """The `block` action (task.escalate) must NOT run on rate_limited path.""" + agent_id = uuid4() + task_id = uuid4() + deps = _make_deps(agent_id, task_id) + c = Choreographer(deps) + + await c.i_am_blocked(agent_id, task_id, "rate_limited") + + # The VerbRunner calls task.escalate for the normal block path. + # In the rate-limited path this must NOT happen. + deps.task.escalate.assert_not_awaited() + + async def test_calling_agent_parked_via_mark_waiting_long(self) -> None: + """mark_waiting_long must be called with waiting_for='rate_limit_lifted'.""" + agent_id = uuid4() + task_id = uuid4() + orch = _make_orchestrator(active_agents=["be-dev-1"]) + deps = _make_deps(agent_id, task_id, orchestrator=orch) + c = Choreographer(deps) + + await c.i_am_blocked(agent_id, task_id, "rate_limited") + + # Verify that at least one mark_waiting_long call uses the right reason. + # The implementation calls mark_waiting_long(slug, waiting_for=..., ...) + # so waiting_for is always a keyword argument. + waiting_for_values = [ + c.kwargs.get("waiting_for") for c in orch.mark_waiting_long.call_args_list + ] + assert "rate_limit_lifted" in waiting_for_values + + async def test_case_insensitive_reason_match(self) -> None: + """reason='Rate_Limited' (any case) should trigger the special path.""" + agent_id = uuid4() + task_id = uuid4() + orch = _make_orchestrator(active_agents=["be-dev-1"]) + deps = _make_deps(agent_id, task_id, orchestrator=orch) + c = Choreographer(deps) + + env = await c.i_am_blocked(agent_id, task_id, "Rate_Limited") + + assert env.error is None + assert env.status == "in_progress" + + async def test_struggle_journal_still_written(self) -> None: + """journal.write_struggle must still be written on the rate_limited path.""" + agent_id = uuid4() + task_id = uuid4() + deps = _make_deps(agent_id, task_id) + c = Choreographer(deps) + + await c.i_am_blocked(agent_id, task_id, "rate_limited") + + deps.journal.write_struggle.assert_awaited_once() + + +# --------------------------------------------------------------------------- +# AC4: mark_waiting_long called for every active agent on affected provider +# --------------------------------------------------------------------------- + + +class TestMarkWaitingLongCallCount: + async def test_call_count_equals_active_agent_count(self) -> None: + """mark_waiting_long must be called once per active agent.""" + agent_id = uuid4() + task_id = uuid4() + active = ["be-dev-1", "be-dev-2", "be-dev-3"] + orch = _make_orchestrator(active_agents=active) + deps = _make_deps(agent_id, task_id, orchestrator=orch) + c = Choreographer(deps) + + await c.i_am_blocked(agent_id, task_id, "rate_limited") + + assert orch.mark_waiting_long.call_count == len(active) + + async def test_call_count_with_single_active_agent(self) -> None: + agent_id = uuid4() + task_id = uuid4() + orch = _make_orchestrator(active_agents=["be-dev-1"]) + deps = _make_deps(agent_id, task_id, orchestrator=orch) + c = Choreographer(deps) + + await c.i_am_blocked(agent_id, task_id, "rate_limited") + + assert orch.mark_waiting_long.call_count == 1 + + async def test_no_calls_when_no_active_agents(self) -> None: + agent_id = uuid4() + task_id = uuid4() + orch = _make_orchestrator(active_agents=[]) + deps = _make_deps(agent_id, task_id, orchestrator=orch) + c = Choreographer(deps) + + await c.i_am_blocked(agent_id, task_id, "rate_limited") + + assert orch.mark_waiting_long.call_count == 0 + + async def test_no_calls_when_orchestrator_is_none(self) -> None: + """When orchestrator is not wired in, no parking happens but no crash.""" + agent_id = uuid4() + task_id = uuid4() + deps = _make_deps(agent_id, task_id, orchestrator=None) + c = Choreographer(deps) + + env = await c.i_am_blocked(agent_id, task_id, "rate_limited") + + # Should still succeed; no orchestrator = no parking + assert env.error is None + assert env.status == "in_progress" + + async def test_mark_waiting_long_receives_waiting_for_arg(self) -> None: + """Every mark_waiting_long call must carry waiting_for='rate_limit_lifted'.""" + agent_id = uuid4() + task_id = uuid4() + active = ["be-dev-1", "be-qa"] + orch = _make_orchestrator(active_agents=active) + deps = _make_deps(agent_id, task_id, orchestrator=orch) + c = Choreographer(deps) + + await c.i_am_blocked(agent_id, task_id, "rate_limited") + + for c_args in orch.mark_waiting_long.call_args_list: + # mark_waiting_long(slug, waiting_for=..., ...) — waiting_for is a kwarg + assert c_args.kwargs.get("waiting_for") == "rate_limit_lifted" + + +# --------------------------------------------------------------------------- +# AC5: RATE_LIMIT_HIT event published with correct payload structure +# --------------------------------------------------------------------------- + + +class TestRateLimitHitEventPublished: + async def test_stream_bus_publish_called_once(self) -> None: + agent_id = uuid4() + task_id = uuid4() + orch = _make_orchestrator() + bus = _make_stream_bus() + deps = _make_deps(agent_id, task_id, orchestrator=orch, stream_bus=bus) + c = Choreographer(deps) + + await c.i_am_blocked(agent_id, task_id, "rate_limited") + + bus.publish.assert_awaited_once() + + async def test_event_type_is_rate_limit_hit(self) -> None: + agent_id = uuid4() + task_id = uuid4() + orch = _make_orchestrator() + bus = _make_stream_bus() + deps = _make_deps(agent_id, task_id, orchestrator=orch, stream_bus=bus) + c = Choreographer(deps) + + await c.i_am_blocked(agent_id, task_id, "rate_limited") + + event = bus.publish.call_args.args[0] + assert event.type == EventType.RATE_LIMIT_HIT + + async def test_event_data_has_provider_field(self) -> None: + agent_id = uuid4() + task_id = uuid4() + orch = _make_orchestrator(provider="anthropic") + bus = _make_stream_bus() + deps = _make_deps(agent_id, task_id, orchestrator=orch, stream_bus=bus) + c = Choreographer(deps) + + await c.i_am_blocked(agent_id, task_id, "rate_limited") + + event = bus.publish.call_args.args[0] + assert "provider" in event.data + assert event.data["provider"] == "anthropic" + + async def test_event_data_has_affected_agents_list(self) -> None: + agent_id = uuid4() + task_id = uuid4() + active = ["be-dev-1", "be-dev-2"] + orch = _make_orchestrator(active_agents=active) + bus = _make_stream_bus() + deps = _make_deps(agent_id, task_id, orchestrator=orch, stream_bus=bus) + c = Choreographer(deps) + + await c.i_am_blocked(agent_id, task_id, "rate_limited") + + event = bus.publish.call_args.args[0] + assert "affectedAgents" in event.data + assert isinstance(event.data["affectedAgents"], list) + assert event.data["affectedAgents"] == active + + async def test_event_data_has_retry_after_seconds_null_by_default(self) -> None: + agent_id = uuid4() + task_id = uuid4() + orch = _make_orchestrator() + bus = _make_stream_bus() + deps = _make_deps(agent_id, task_id, orchestrator=orch, stream_bus=bus) + c = Choreographer(deps) + + await c.i_am_blocked(agent_id, task_id, "rate_limited") + + event = bus.publish.call_args.args[0] + assert "retryAfterSeconds" in event.data + assert event.data["retryAfterSeconds"] is None + + async def test_event_data_retry_after_parsed_from_what_needed(self) -> None: + """If what_needed is a numeric string, it becomes retryAfterSeconds.""" + agent_id = uuid4() + task_id = uuid4() + orch = _make_orchestrator() + bus = _make_stream_bus() + deps = _make_deps(agent_id, task_id, orchestrator=orch, stream_bus=bus) + c = Choreographer(deps) + + await c.i_am_blocked(agent_id, task_id, "rate_limited", what_needed="30") + + event = bus.publish.call_args.args[0] + assert event.data["retryAfterSeconds"] == float("30") + + async def test_event_data_has_timestamp_iso_string(self) -> None: + agent_id = uuid4() + task_id = uuid4() + orch = _make_orchestrator() + bus = _make_stream_bus() + deps = _make_deps(agent_id, task_id, orchestrator=orch, stream_bus=bus) + c = Choreographer(deps) + + await c.i_am_blocked(agent_id, task_id, "rate_limited") + + event = bus.publish.call_args.args[0] + assert "timestamp" in event.data + # ISO string: must be a non-empty string + ts = event.data["timestamp"] + assert isinstance(ts, str) and len(ts) > 0 + + async def test_no_publish_when_stream_bus_is_none(self) -> None: + """When stream_bus is not wired in, no publish is attempted.""" + agent_id = uuid4() + task_id = uuid4() + orch = _make_orchestrator() + # stream_bus=None: no bus + deps = _make_deps(agent_id, task_id, orchestrator=orch, stream_bus=None) + c = Choreographer(deps) + + env = await c.i_am_blocked(agent_id, task_id, "rate_limited") + + # Should still succeed + assert env.error is None + assert env.status == "in_progress" + + +# --------------------------------------------------------------------------- +# AC1: RateLimitStateTracker.activate() called on rate_limited path +# --------------------------------------------------------------------------- + +_TRACKER_PATCH = "roboco.services.gateway.rate_limit_tracker.RateLimitStateTracker" + + +class TestRateLimitTrackerActivateOnParking: + """Verify that _handle_rate_limited_parking() calls activate().""" + + async def test_activate_called_when_provider_known(self) -> None: + """activate() must be called once when provider != 'unknown'.""" + agent_id = uuid4() + task_id = uuid4() + orch = _make_orchestrator(active_agents=["be-dev-1"], provider=_PROVIDER) + deps = _make_deps(agent_id, task_id, orchestrator=orch) + c = Choreographer(deps) + + mock_tracker = AsyncMock() + mock_tracker.activate = AsyncMock(return_value=None) + mock_tracker_cls = MagicMock(return_value=mock_tracker) + + with patch(_TRACKER_PATCH, mock_tracker_cls): + await c.i_am_blocked(agent_id, task_id, "rate_limited") + + mock_tracker_cls.assert_called_once_with(_PROVIDER) + mock_tracker.activate.assert_awaited_once() + + async def test_activate_receives_affected_agents(self) -> None: + """activate() must be called with the affected_agents list.""" + agent_id = uuid4() + task_id = uuid4() + active = ["be-dev-1", "be-dev-2"] + orch = _make_orchestrator(active_agents=active, provider=_PROVIDER) + deps = _make_deps(agent_id, task_id, orchestrator=orch) + c = Choreographer(deps) + + mock_tracker = AsyncMock() + mock_tracker.activate = AsyncMock(return_value=None) + mock_tracker_cls = MagicMock(return_value=mock_tracker) + + with patch(_TRACKER_PATCH, mock_tracker_cls): + await c.i_am_blocked(agent_id, task_id, "rate_limited") + + call_kwargs = mock_tracker.activate.call_args.kwargs + assert call_kwargs.get("affected_agents") == active + + async def test_activate_receives_retry_after_from_what_needed(self) -> None: + """activate() must receive retry_after parsed from what_needed.""" + agent_id = uuid4() + task_id = uuid4() + orch = _make_orchestrator(active_agents=["be-dev-1"], provider=_PROVIDER) + deps = _make_deps(agent_id, task_id, orchestrator=orch) + c = Choreographer(deps) + + mock_tracker = AsyncMock() + mock_tracker.activate = AsyncMock(return_value=None) + mock_tracker_cls = MagicMock(return_value=mock_tracker) + + with patch(_TRACKER_PATCH, mock_tracker_cls): + await c.i_am_blocked(agent_id, task_id, "rate_limited", what_needed="45") + + call_kwargs = mock_tracker.activate.call_args.kwargs + assert call_kwargs.get("retry_after") == float("45") + + async def test_activate_retry_after_none_when_what_needed_not_numeric(self) -> None: + """activate() must receive retry_after=None when what_needed is not a number.""" + agent_id = uuid4() + task_id = uuid4() + orch = _make_orchestrator(active_agents=["be-dev-1"], provider=_PROVIDER) + deps = _make_deps(agent_id, task_id, orchestrator=orch) + c = Choreographer(deps) + + mock_tracker = AsyncMock() + mock_tracker.activate = AsyncMock(return_value=None) + mock_tracker_cls = MagicMock(return_value=mock_tracker) + + with patch(_TRACKER_PATCH, mock_tracker_cls): + await c.i_am_blocked( + agent_id, task_id, "rate_limited", what_needed="retry soon" + ) + + call_kwargs = mock_tracker.activate.call_args.kwargs + assert call_kwargs.get("retry_after") is None + + async def test_activate_skipped_when_provider_unknown(self) -> None: + """activate() must NOT be called when provider resolves to 'unknown'.""" + agent_id = uuid4() + task_id = uuid4() + # get_provider_for_agent returns None → provider stays 'unknown' + orch = MagicMock() + orch.get_provider_for_agent = MagicMock(return_value=None) + orch.get_active_agent_slugs_for_provider = MagicMock(return_value=[]) + orch.mark_waiting_long = AsyncMock(return_value=None) + deps = _make_deps(agent_id, task_id, orchestrator=orch) + c = Choreographer(deps) + + mock_tracker = AsyncMock() + mock_tracker.activate = AsyncMock(return_value=None) + mock_tracker_cls = MagicMock(return_value=mock_tracker) + + with patch(_TRACKER_PATCH, mock_tracker_cls): + env = await c.i_am_blocked(agent_id, task_id, "rate_limited") + + # No crash, no activate call + assert env.error is None + mock_tracker.activate.assert_not_awaited() + + async def test_activate_failure_does_not_crash_path(self) -> None: + """If activate() raises, _handle_rate_limited_parking must still succeed.""" + agent_id = uuid4() + task_id = uuid4() + orch = _make_orchestrator(active_agents=["be-dev-1"], provider=_PROVIDER) + deps = _make_deps(agent_id, task_id, orchestrator=orch) + c = Choreographer(deps) + + mock_tracker = AsyncMock() + mock_tracker.activate = AsyncMock(side_effect=RuntimeError("redis down")) + mock_tracker_cls = MagicMock(return_value=mock_tracker) + + with patch(_TRACKER_PATCH, mock_tracker_cls): + env = await c.i_am_blocked(agent_id, task_id, "rate_limited") + + assert env.error is None + assert env.status == "in_progress" diff --git a/tests/unit/gateway/test_trigger_filter.py b/tests/unit/gateway/test_trigger_filter.py index 9102fa64..bf6f39a2 100644 --- a/tests/unit/gateway/test_trigger_filter.py +++ b/tests/unit/gateway/test_trigger_filter.py @@ -34,17 +34,21 @@ def _task( return t -def _trigger( +def _trigger( # noqa: PLR0913 kind: TriggerKind, skill: str | None = None, recent_spawns_for_task: int = 0, recent_spawns_for_role: int = 0, + provider: str | None = None, + provider_rate_limited: bool = False, ) -> TriggerContext: return TriggerContext( kind=kind, skill=skill, recent_spawns_for_task=recent_spawns_for_task, recent_spawns_for_role=recent_spawns_for_role, + provider=provider, + provider_rate_limited=provider_rate_limited, ) @@ -148,3 +152,101 @@ class TestCooldown: ) assert decision.outcome == SpawnDecision.QUEUE assert "rate" in decision.reason.lower() + + +class TestProviderRateLimitGate: + """Rule 2: provider rate-limit gate fires before claimant-lock/cooldown.""" + + def test_queues_when_provider_rate_limited(self) -> None: + """QUEUE outcome when provider_rate_limited=True.""" + t = _task(status="in_progress") + decision = decide_spawn( + task=t, + trigger=_trigger( + TriggerKind.NOTIFICATION, + provider="anthropic", + provider_rate_limited=True, + ), + config=_DEFAULT_CONFIG, + ) + assert decision.outcome == SpawnDecision.QUEUE + assert "provider anthropic rate-limited" in decision.reason + + def test_reason_contains_provider_name(self) -> None: + """Reason string must contain the provider name.""" + t = _task(status="pending") + decision = decide_spawn( + task=t, + trigger=_trigger( + TriggerKind.SCAN, + provider="ollama_cloud", + provider_rate_limited=True, + ), + config=_DEFAULT_CONFIG, + ) + assert "ollama_cloud" in decision.reason + + def test_reason_contains_unknown_when_no_provider_name(self) -> None: + """When provider is None, reason still contains 'unknown'.""" + t = _task(status="in_progress") + decision = decide_spawn( + task=t, + trigger=_trigger( + TriggerKind.NOTIFICATION, + provider=None, + provider_rate_limited=True, + ), + config=_DEFAULT_CONFIG, + ) + assert decision.outcome == SpawnDecision.QUEUE + assert "unknown" in decision.reason + + def test_no_queue_injection_when_not_rate_limited(self) -> None: + """SPAWN when provider_rate_limited=False and all other gates clear.""" + t = _task(status="in_progress") + decision = decide_spawn( + task=t, + trigger=_trigger( + TriggerKind.NOTIFICATION, + provider="anthropic", + provider_rate_limited=False, + ), + config=_DEFAULT_CONFIG, + ) + assert decision.outcome == SpawnDecision.SPAWN + + def test_stale_drop_fires_before_rate_limit_gate(self) -> None: + """Rule 1 (stale-drop) fires before rule 2 (rate-limit gate).""" + t = _task(status="completed") + decision = decide_spawn( + task=t, + trigger=_trigger( + TriggerKind.NOTIFICATION, + provider="anthropic", + provider_rate_limited=True, + ), + config=_DEFAULT_CONFIG, + ) + # Rule 1 fires first — outcome must be DROP, not QUEUE + assert decision.outcome == SpawnDecision.DROP + + def test_rate_limit_gate_fires_before_claimant_lock(self) -> None: + """Rule 2 (rate-limit gate) fires before rule 3 (single-claimant invariant).""" + recent = datetime.now(tz=UTC) + t = _task( + status="in_progress", + active_claimant_id=uuid4(), + last_heartbeat_at=recent, + ) + decision = decide_spawn( + task=t, + trigger=_trigger( + TriggerKind.NOTIFICATION, + provider="anthropic", + provider_rate_limited=True, + ), + config=_DEFAULT_CONFIG, + ) + # Both gates would QUEUE but reason must come from rate-limit (rule 2) + assert decision.outcome == SpawnDecision.QUEUE + assert "rate-limited" in decision.reason diff --git a/tests/unit/runtime/test_rate_limit_sweep.py b/tests/unit/runtime/test_rate_limit_sweep.py new file mode 100644 index 00000000..78e82f59 --- /dev/null +++ b/tests/unit/runtime/test_rate_limit_sweep.py @@ -0,0 +1,550 @@ +"""Unit tests for the rate-limit sweeper probe loop (AC4, AC8). + +Tests cover: +- probe-success path: tracker.clear() + resolve_wait + RATE_LIMIT_LIFTED event +- probe-failure path: increment_probe_failures is called +- CEO notification fires at threshold 10 exactly once per episode +- ``_do_probe`` / ``_make_tracker`` are injectable boundaries for mocking +""" + +from __future__ import annotations + +import fnmatch +import json +from datetime import UTC, datetime +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch +from uuid import uuid4 + +from httpx import ASGITransport, AsyncClient +from roboco.api.app import create_app +from roboco.models.events import EventType +from roboco.models.runtime import WaitingRecord +from roboco.runtime.orchestrator import AgentOrchestrator +from roboco.services.gateway.rate_limit_tracker import RateLimitStateTracker + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_redis_mock(initial_store: dict[str, Any] | None = None) -> AsyncMock: + """Fake redis.asyncio.Redis backed by a plain dict.""" + store: dict[str, Any] = initial_store if initial_store is not None else {} + + async def _get(key: str) -> bytes | None: + val = store.get(key) + if val is None: + return None + return str(val).encode() if not isinstance(val, bytes) else val + + async def _set(key: str, value: Any) -> None: + store[key] = value + + async def _delete(key: str) -> int: + return 1 if store.pop(key, None) is not None else 0 + + async def _scan( + _cursor: int, match: str = "*", count: int = 100 # noqa: ARG001 + ) -> tuple[int, list[bytes]]: + # Simple in-memory scan: return all matching keys in one shot + matches = [k.encode() for k in store if fnmatch.fnmatch(k, match)] + return (0, matches) + + async def _aclose() -> None: + pass + + mock = AsyncMock() + mock.get = AsyncMock(side_effect=_get) + mock.set = AsyncMock(side_effect=_set) + mock.delete = AsyncMock(side_effect=_delete) + mock.scan = AsyncMock(side_effect=_scan) + mock.aclose = AsyncMock(side_effect=_aclose) + mock._store = store + return mock + + +def _make_orchestrator() -> AgentOrchestrator: + """Build a minimal orchestrator via __new__ (no __init__ side-effects).""" + orch = AgentOrchestrator.__new__(AgentOrchestrator) + orch._running = True + orch._waiting_records: dict[str, WaitingRecord] = {} + orch._instances: dict[str, Any] = {} + orch._rate_limit_ceo_notified: set[str] = set() + return orch + + +def _make_tracker_mock(failure_return: int = 1) -> AsyncMock: + """Create an async mock RateLimitStateTracker instance.""" + mock = AsyncMock() + mock.clear = AsyncMock() + mock.increment_probe_failures = AsyncMock(return_value=failure_return) + mock.reset_probe_failures = AsyncMock() + return mock + + +def _make_active_state( + _provider: str = "anthropic", + retry_after: float | None = None, + probe_failures: int = 0, + activated_at: datetime | None = None, +) -> dict[str, Any]: + """Build a tracker state dict.""" + at = activated_at or datetime.now(UTC) + return { + "rate_limited": True, + "activated_at": at.isoformat(), + "retry_after": retry_after, + "affected_agents": ["be-dev-1"], + "probe_failures": probe_failures, + } + + +def _waiting_record( + agent_id: str, + provider: str = "anthropic", + task_id: str | None = None, +) -> WaitingRecord: + return WaitingRecord( + agent_id=agent_id, + task_id=task_id or str(uuid4()), + waiting_for="rate_limit_lifted", + waiting_since=datetime.now(UTC), + context={"provider": provider}, + ) + + +# --------------------------------------------------------------------------- +# Tests: probe-success path (AC4) +# --------------------------------------------------------------------------- + + +class TestProbeSuccessPath: + """When _do_probe returns True the rate limit should be cleared and + all parked agents resolved.""" + + async def test_tracker_clear_called_on_success(self) -> None: + """tracker.clear() is invoked when the probe succeeds.""" + orch = _make_orchestrator() + provider = "anthropic" + state = _make_active_state(provider, retry_after=None) + + tracker_mock = _make_tracker_mock() + orch._make_tracker = MagicMock(return_value=tracker_mock) # type: ignore[method-assign] + orch.resolve_wait = AsyncMock(return_value=None) + + async def fake_do_probe(_p: str) -> bool: + return True + + orch._do_probe = fake_do_probe # type: ignore[method-assign] + + with patch("roboco.events.get_event_bus") as mock_bus_fn: + bus_mock = AsyncMock() + bus_mock.publish = AsyncMock() + mock_bus_fn.return_value = bus_mock + + await orch._probe_one_provider(provider, state) + + tracker_mock.clear.assert_awaited_once() + + async def test_resolve_wait_called_for_parked_agents(self) -> None: + """resolve_wait is called for each agent waiting for rate_limit_lifted.""" + orch = _make_orchestrator() + provider = "anthropic" + state = _make_active_state(provider, retry_after=None) + + agent1 = "be-dev-1" + agent2 = "be-dev-2" + orch._waiting_records = { + agent1: _waiting_record(agent1, provider), + agent2: _waiting_record(agent2, provider), + "be-qa-1": _waiting_record( + "be-qa-1", "other-provider" + ), # different provider + } + + orch.resolve_wait = AsyncMock(return_value=None) + + tracker_mock = _make_tracker_mock() + orch._make_tracker = MagicMock(return_value=tracker_mock) # type: ignore[method-assign] + + async def fake_do_probe(_p: str) -> bool: + return True + + orch._do_probe = fake_do_probe # type: ignore[method-assign] + + with patch("roboco.events.get_event_bus") as mock_bus_fn: + bus_mock = AsyncMock() + bus_mock.publish = AsyncMock() + mock_bus_fn.return_value = bus_mock + + await orch._probe_one_provider(provider, state) + + # Only the two anthropic-parked agents should be resolved + assert orch.resolve_wait.await_count == 2 # noqa: PLR2004 + resolved_ids = {call.args[0] for call in orch.resolve_wait.call_args_list} + assert agent1 in resolved_ids + assert agent2 in resolved_ids + assert "be-qa-1" not in resolved_ids + + async def test_rate_limit_lifted_event_published(self) -> None: + """RATE_LIMIT_LIFTED event is published to the bus on probe success.""" + orch = _make_orchestrator() + provider = "anthropic" + state = _make_active_state(provider, retry_after=None) + + orch.resolve_wait = AsyncMock(return_value=None) + + tracker_mock = _make_tracker_mock() + orch._make_tracker = MagicMock(return_value=tracker_mock) # type: ignore[method-assign] + + published_events: list[Any] = [] + + async def fake_do_probe(_p: str) -> bool: + return True + + orch._do_probe = fake_do_probe # type: ignore[method-assign] + + with patch("roboco.events.get_event_bus") as mock_bus_fn: + bus_mock = AsyncMock() + bus_mock.publish = AsyncMock(side_effect=published_events.append) + mock_bus_fn.return_value = bus_mock + + await orch._probe_one_provider(provider, state) + + assert len(published_events) == 1 + event = published_events[0] + assert event.type == EventType.RATE_LIMIT_LIFTED + assert event.data["provider"] == provider + + async def test_ceo_notified_flag_cleared_on_success(self) -> None: + """_rate_limit_ceo_notified is cleared when probe succeeds.""" + orch = _make_orchestrator() + provider = "anthropic" + orch._rate_limit_ceo_notified.add(provider) # simulates prior episode + state = _make_active_state(provider, retry_after=None) + + orch.resolve_wait = AsyncMock(return_value=None) + + tracker_mock = _make_tracker_mock() + orch._make_tracker = MagicMock(return_value=tracker_mock) # type: ignore[method-assign] + + async def fake_do_probe(_p: str) -> bool: + return True + + orch._do_probe = fake_do_probe # type: ignore[method-assign] + + with patch("roboco.events.get_event_bus") as mock_bus_fn: + bus_mock = AsyncMock() + bus_mock.publish = AsyncMock() + mock_bus_fn.return_value = bus_mock + + await orch._probe_one_provider(provider, state) + + assert provider not in orch._rate_limit_ceo_notified + + async def test_probe_skipped_before_estimated_lift_at(self) -> None: + """When retry_after has not elapsed yet the probe is skipped entirely.""" + orch = _make_orchestrator() + provider = "anthropic" + # Set activated_at to now; retry_after = 300s → estimated lift in future + state = _make_active_state( + provider, + retry_after=300.0, + activated_at=datetime.now(UTC), + ) + + probe_called = [] + + async def fake_do_probe(_p: str) -> bool: + probe_called.append(_p) + return True + + orch._do_probe = fake_do_probe # type: ignore[method-assign] + + await orch._probe_one_provider(provider, state) + + assert probe_called == [] # probe was gated by time + + +# --------------------------------------------------------------------------- +# Tests: probe-failure path (AC4) +# --------------------------------------------------------------------------- + + +class TestProbeFailurePath: + """When _do_probe returns False the failure counter should be incremented.""" + + async def test_increment_probe_failures_called_on_failure(self) -> None: + """increment_probe_failures is called when the probe fails.""" + orch = _make_orchestrator() + provider = "anthropic" + state = _make_active_state(provider, retry_after=None) + + tracker_mock = _make_tracker_mock(failure_return=1) + orch._make_tracker = MagicMock(return_value=tracker_mock) # type: ignore[method-assign] + + async def fake_do_probe(_p: str) -> bool: + return False + + orch._do_probe = fake_do_probe # type: ignore[method-assign] + orch._notify_rate_limit_ceo = AsyncMock() + + await orch._probe_one_provider(provider, state) + + tracker_mock.increment_probe_failures.assert_awaited_once() + + async def test_clear_not_called_on_failure(self) -> None: + """tracker.clear() must NOT be called when the probe fails.""" + orch = _make_orchestrator() + provider = "anthropic" + state = _make_active_state(provider, retry_after=None) + + tracker_mock = _make_tracker_mock(failure_return=1) + orch._make_tracker = MagicMock(return_value=tracker_mock) # type: ignore[method-assign] + + async def fake_do_probe(_p: str) -> bool: + return False + + orch._do_probe = fake_do_probe # type: ignore[method-assign] + orch._notify_rate_limit_ceo = AsyncMock() + + await orch._probe_one_provider(provider, state) + + tracker_mock.clear.assert_not_awaited() + + +# --------------------------------------------------------------------------- +# Tests: CEO notification threshold (AC8) +# --------------------------------------------------------------------------- + + +class TestCEONotificationThreshold: + """CEO notification fires at count==10 exactly once per episode.""" + + async def test_notification_fires_at_exactly_10_failures(self) -> None: + """_notify_rate_limit_ceo is called when failure count hits 10.""" + orch = _make_orchestrator() + provider = "anthropic" + state = _make_active_state(provider, retry_after=None) + + # simulate already at 9 failures; next increment returns 10 + tracker_mock = _make_tracker_mock(failure_return=10) + orch._make_tracker = MagicMock(return_value=tracker_mock) # type: ignore[method-assign] + orch._notify_rate_limit_ceo = AsyncMock() + + async def fake_do_probe(_p: str) -> bool: + return False + + orch._do_probe = fake_do_probe # type: ignore[method-assign] + + await orch._probe_one_provider(provider, state) + + orch._notify_rate_limit_ceo.assert_awaited_once() + + async def test_notification_not_fired_before_threshold(self) -> None: + """No CEO notification below threshold 10.""" + orch = _make_orchestrator() + provider = "anthropic" + state = _make_active_state(provider, retry_after=None) + + tracker_mock = _make_tracker_mock(failure_return=9) + orch._make_tracker = MagicMock(return_value=tracker_mock) # type: ignore[method-assign] + orch._notify_rate_limit_ceo = AsyncMock() + + async def fake_do_probe(_p: str) -> bool: + return False + + orch._do_probe = fake_do_probe # type: ignore[method-assign] + + await orch._probe_one_provider(provider, state) + + orch._notify_rate_limit_ceo.assert_not_awaited() + + async def test_notification_sent_only_once_per_episode(self) -> None: + """Even if failures keep accumulating, the CEO is notified only once.""" + orch = _make_orchestrator() + provider = "anthropic" + state = _make_active_state(provider, retry_after=None) + + # Mark this episode as already notified + orch._rate_limit_ceo_notified.add(provider) + + tracker_mock = _make_tracker_mock(failure_return=15) + orch._make_tracker = MagicMock(return_value=tracker_mock) # type: ignore[method-assign] + orch._notify_rate_limit_ceo = AsyncMock() + + async def fake_do_probe(_p: str) -> bool: + return False + + orch._do_probe = fake_do_probe # type: ignore[method-assign] + + await orch._probe_one_provider(provider, state) + + orch._notify_rate_limit_ceo.assert_not_awaited() + + async def test_new_episode_allows_new_notification(self) -> None: + """After a rate-limit clears (success) a new episode starts fresh.""" + orch = _make_orchestrator() + provider = "anthropic" + # Episode 1: had a notification + orch._rate_limit_ceo_notified.add(provider) + + success_state = _make_active_state(provider, retry_after=None) + orch.resolve_wait = AsyncMock(return_value=None) + + tracker_mock = _make_tracker_mock(failure_return=10) + orch._make_tracker = MagicMock(return_value=tracker_mock) # type: ignore[method-assign] + notify_mock = AsyncMock() + orch._notify_rate_limit_ceo = notify_mock + + async def fake_do_probe_success(_p: str) -> bool: + return True + + orch._do_probe = fake_do_probe_success # type: ignore[method-assign] + + with patch("roboco.events.get_event_bus") as mock_bus_fn: + bus_mock = AsyncMock() + bus_mock.publish = AsyncMock() + mock_bus_fn.return_value = bus_mock + + # Success clears the episode flag + await orch._probe_one_provider(provider, success_state) + + assert provider not in orch._rate_limit_ceo_notified + + # Episode 2: simulate a new failure reaching threshold 10 + async def fake_do_probe_fail(_p: str) -> bool: + return False + + orch._do_probe = fake_do_probe_fail # type: ignore[method-assign] + + failure_state = _make_active_state(provider, retry_after=None) + await orch._probe_one_provider(provider, failure_state) + + # Notification SHOULD fire for the new episode + notify_mock.assert_awaited_once() + + +# --------------------------------------------------------------------------- +# Tests: list_rate_limited_providers +# --------------------------------------------------------------------------- + + +class TestListRateLimitedProviders: + """list_rate_limited_providers scans Redis for active rate-limit keys.""" + + async def test_returns_empty_when_no_keys(self) -> None: + redis_mock = _make_redis_mock() + with patch("redis.asyncio.from_url", return_value=redis_mock): + result = await RateLimitStateTracker.list_rate_limited_providers() + assert result == [] + + async def test_returns_active_provider(self) -> None: + state = { + "rate_limited": True, + "activated_at": datetime.now(UTC).isoformat(), + "retry_after": 60.0, + "affected_agents": ["be-dev-1"], + "probe_failures": 0, + } + store = {"roboco:rate_limit:anthropic:state": json.dumps(state).encode()} + redis_mock = _make_redis_mock(store) + + with patch("redis.asyncio.from_url", return_value=redis_mock): + result = await RateLimitStateTracker.list_rate_limited_providers() + + assert len(result) == 1 + provider, returned_state = result[0] + assert provider == "anthropic" + assert returned_state["rate_limited"] is True + + async def test_ignores_cleared_providers(self) -> None: + state = { + "rate_limited": False, + "activated_at": datetime.now(UTC).isoformat(), + "retry_after": 60.0, + "affected_agents": [], + "probe_failures": 2, + } + store = {"roboco:rate_limit:anthropic:state": json.dumps(state).encode()} + redis_mock = _make_redis_mock(store) + + with patch("redis.asyncio.from_url", return_value=redis_mock): + result = await RateLimitStateTracker.list_rate_limited_providers() + + assert result == [] + + +# --------------------------------------------------------------------------- +# Tests: GET /api/system/rate-limits endpoint schema (AC9) +# --------------------------------------------------------------------------- + + +class TestRateLimitsEndpoint: + """GET /api/system/rate-limits returns correct schema.""" + + async def test_returns_empty_list_when_no_rate_limits(self) -> None: + app = create_app() + + with patch( + "roboco.api.routes.system.RateLimitStateTracker" + ".list_rate_limited_providers", + new_callable=AsyncMock, + return_value=[], + ): + async with AsyncClient( + transport=ASGITransport(app=app), base_url="http://test" + ) as client: + resp = await client.get("/api/system/rate-limits") + + assert resp.status_code == 200 # noqa: PLR2004 + assert resp.json() == [] + + async def test_returns_provider_state_when_rate_limited(self) -> None: + app = create_app() + + state = { + "rate_limited": True, + "activated_at": "2026-06-11T00:00:00+00:00", + "retry_after": 60.0, + "affected_agents": ["be-dev-1"], + "probe_failures": 3, + } + + with patch( + "roboco.api.routes.system.RateLimitStateTracker" + ".list_rate_limited_providers", + new_callable=AsyncMock, + return_value=[("anthropic", state)], + ): + async with AsyncClient( + transport=ASGITransport(app=app), base_url="http://test" + ) as client: + resp = await client.get("/api/system/rate-limits") + + assert resp.status_code == 200 # noqa: PLR2004 + data = resp.json() + assert len(data) == 1 + entry = data[0] + assert entry["provider"] == "anthropic" + assert entry["rate_limited"] is True + assert entry["probe_failures"] == 3 # noqa: PLR2004 + assert entry["retry_after"] == 60.0 # noqa: PLR2004 + + async def test_endpoint_not_404(self) -> None: + """The endpoint must be registered in app.py — no 404.""" + app = create_app() + + with patch( + "roboco.api.routes.system.RateLimitStateTracker" + ".list_rate_limited_providers", + new_callable=AsyncMock, + return_value=[], + ): + async with AsyncClient( + transport=ASGITransport(app=app), base_url="http://test" + ) as client: + resp = await client.get("/api/system/rate-limits") + + assert resp.status_code != 404 # noqa: PLR2004 diff --git a/tests/unit/services/test_rate_limit_retry.py b/tests/unit/services/test_rate_limit_retry.py new file mode 100644 index 00000000..18c256c2 --- /dev/null +++ b/tests/unit/services/test_rate_limit_retry.py @@ -0,0 +1,691 @@ +""" +Unit tests for rate-limit retry behaviour across all LLM call sites. + +Covers acceptance criteria: +- 5-retry exhaustion raises RateLimitError +- Retry-After header drives the sleep duration +- Partial retries then success returns the correct result +- ConnectError / TimeoutException in OllamaEmbedder does NOT trigger the 429 + retry path (the two concerns are composed without double-retrying) +""" + +from __future__ import annotations + +import json +import sys +import types +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch +from uuid import uuid4 + +import anthropic as anthropic_mod +import httpx +import pytest +import pytest_asyncio # noqa: F401 - registers asyncio mode + +# --------------------------------------------------------------------------- +# Ensure piragi stubs are present before the optimal_brain modules are imported +# --------------------------------------------------------------------------- + +_PIRAGI_STUB_NAMES = ( + "piragi", + "piragi.types", + "piragi.stores", + "piragi.stores.postgres", + "piragi.chunking", + "piragi.semantic_chunking", +) + + +def _stub_piragi() -> None: + mock = MagicMock() + for name in _PIRAGI_STUB_NAMES: + if name not in sys.modules: + mod = types.ModuleType(name) + mod.__dict__.update( + { + "AsyncRagi": mock, + "Citation": mock, + "Document": mock, + "Chunk": mock, + "PostgresStore": mock, + } + ) + sys.modules[name] = mod + + +_stub_piragi() + +# --------------------------------------------------------------------------- +# Module imports (after stubs are injected) +# --------------------------------------------------------------------------- + +from roboco.models.extraction import ExtractionContext # noqa: E402 +from roboco.models.optimal import IndexType # noqa: E402 +from roboco.services.exceptions import ( # noqa: E402 + MAX_RATE_LIMIT_RETRIES, + RateLimitError, + parse_retry_after_header, +) +from roboco.services.extraction import ExtractionService # noqa: E402 +from roboco.services.optimal_brain.indexes.journals import ( # noqa: E402 + JournalsIndexPlugin, +) +from roboco.services.optimal_brain.mentor import MentorService # noqa: E402 +from roboco.services.optimal_brain.ollama_embedder import ( # noqa: E402 + MAX_RETRIES, + OllamaConnectionError, + OllamaEmbedder, +) +from roboco.services.optimal_brain.validator import ValidatorService # noqa: E402 + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +_RETRY_AFTER_FLOAT = 30.0 +_RETRY_AFTER_FLOAT_2 = 2.5 +_RETRY_AFTER_7 = 7.0 +_RETRY_AFTER_9 = 9.0 +_RETRY_AFTER_12 = 12.0 +_RETRY_AFTER_5 = 5.0 +_EMBED_DIM = 4 # zero-vector dimension in test responses +_CALLS_2RL_1_SUCCESS = 3 # 2 rate-limit errors then 1 success + +_EMBED_PATH = ( + "roboco.services.optimal_brain.ollama_embedder.OllamaEmbedder._create_async_client" +) + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_response( + status_code: int, + body: Any = None, + retry_after: str | None = None, +) -> httpx.Response: + """Build a minimal httpx.Response for use in mocks.""" + headers: dict[str, str] = {} + if retry_after is not None: + headers["retry-after"] = retry_after + content = json.dumps(body or {}).encode() + return httpx.Response( + status_code=status_code, + headers=headers, + content=content, + ) + + +def _success_embed_response(n: int = 1) -> httpx.Response: + """Return a valid Ollama /api/embed response with *n* zero-vectors.""" + body = {"embeddings": [[0.0] * _EMBED_DIM] * n} + return _make_response(200, body) + + +def _429_response(retry_after: str | None = None) -> httpx.Response: + return _make_response(429, {"error": "rate limited"}, retry_after=retry_after) + + +def _make_async_client_mock(post_return: Any = None) -> AsyncMock: + """Return an async context manager mock exposing `.post`.""" + client = AsyncMock() + client.__aenter__ = AsyncMock(return_value=client) + client.__aexit__ = AsyncMock(return_value=False) + if post_return is not None: + client.post = AsyncMock(return_value=post_return) + return client + + +def _make_extraction_context() -> ExtractionContext: + return ExtractionContext( + content="This is a test message that is long enough.", + agent_id=uuid4(), + channel_id=uuid4(), + session_id=uuid4(), + group_id=uuid4(), + ) + + +def _make_anthropic_rl_exc( + retry_after: str | None = None, +) -> anthropic_mod.RateLimitError: + """Build a minimal anthropic.RateLimitError.""" + fake_resp = MagicMock() + fake_resp.headers = {} if retry_after is None else {"retry-after": retry_after} + return anthropic_mod.RateLimitError( + message="rate limit", + response=fake_resp, + body={}, + ) + + +def _make_journal_plugin() -> JournalsIndexPlugin: + """Create a minimal JournalsIndexPlugin without initialising piragi.""" + plugin = JournalsIndexPlugin.__new__(JournalsIndexPlugin) + plugin._config = MagicMock() + plugin._config.llm_base_url = "http://ollama-test:11434/v1" + plugin._config.llm_model = "glm-5:cloud" + plugin._ragi = MagicMock() + plugin._initialized = True + return plugin + + +def _make_search_outcome(content: str = "context text") -> Any: + mock_outcome = MagicMock() + mock_outcome.success = True + mock_outcome.results = [ + MagicMock(content=content, source="src", score=0.9, index_type=None) + ] + return mock_outcome + + +def _make_sources() -> list[Any]: + return [ + MagicMock(content="ctx", source="s", score=0.9, index_type=IndexType.JOURNALS) + ] + + +def _make_standards() -> list[Any]: + s = MagicMock() + s.content = "### PY-001: Use Type Hints\nMust add return type annotations." + return [s] + + +# =========================================================================== +# 1. parse_retry_after_header +# =========================================================================== + + +class TestParseRetryAfterHeader: + def test_integer_seconds(self) -> None: + resp = _429_response(retry_after="30") + assert parse_retry_after_header(resp) == _RETRY_AFTER_FLOAT + + def test_float_seconds(self) -> None: + resp = _429_response(retry_after="2.5") + assert parse_retry_after_header(resp) == _RETRY_AFTER_FLOAT_2 + + def test_missing_header_returns_none(self) -> None: + resp = _make_response(429) + assert parse_retry_after_header(resp) is None + + def test_non_numeric_returns_none(self) -> None: + resp = _429_response(retry_after="Wed, 21 Oct 2015 07:28:00 GMT") + assert parse_retry_after_header(resp) is None + + +# =========================================================================== +# 2. RateLimitError class +# =========================================================================== + + +class TestRateLimitError: + def test_fields(self) -> None: + err = RateLimitError(provider="anthropic", retry_after=_RETRY_AFTER_FLOAT) + assert err.provider == "anthropic" + assert err.retry_after == _RETRY_AFTER_FLOAT + + def test_message_includes_provider(self) -> None: + err = RateLimitError(provider="ollama") + assert "ollama" in str(err) + + def test_message_includes_retry_after_when_set(self) -> None: + err = RateLimitError(provider="ollama", retry_after=15.0) + assert "15" in str(err) + + def test_none_retry_after(self) -> None: + err = RateLimitError(provider="anthropic", retry_after=None) + assert err.retry_after is None + + +# =========================================================================== +# 3. OllamaEmbedder - async path (aembed_query) +# =========================================================================== + + +class TestOllamaEmbedderAembed: + """Tests for the async aembed_query method.""" + + def _make_embedder(self) -> OllamaEmbedder: + return OllamaEmbedder(base_url="http://ollama-test:11434") + + async def test_five_consecutive_429s_raise_rate_limit_error(self) -> None: + """After MAX_RATE_LIMIT_RETRIES attempts all 429 -> RateLimitError.""" + embedder = self._make_embedder() + mock_c = _make_async_client_mock(post_return=_429_response()) + + with ( + patch(_EMBED_PATH, return_value=mock_c), + patch("asyncio.sleep", new_callable=AsyncMock), + pytest.raises(RateLimitError) as exc_info, + ): + await embedder.aembed_query("hello") + + assert exc_info.value.provider == "ollama" + assert mock_c.post.call_count == MAX_RATE_LIMIT_RETRIES + + async def test_retry_after_header_respected_as_sleep_duration(self) -> None: + """Retry-After: 7 -> asyncio.sleep(7.0) on each inter-attempt gap.""" + embedder = self._make_embedder() + sleep_calls: list[float] = [] + + async def _fake_sleep(secs: float) -> None: + sleep_calls.append(secs) + + mock_c = _make_async_client_mock(post_return=_429_response(retry_after="7")) + + with ( + patch(_EMBED_PATH, return_value=mock_c), + patch("asyncio.sleep", side_effect=_fake_sleep), + pytest.raises(RateLimitError), + ): + await embedder.aembed_query("hello") + + assert all(s == _RETRY_AFTER_7 for s in sleep_calls) + assert len(sleep_calls) == MAX_RATE_LIMIT_RETRIES - 1 + + async def test_partial_retries_then_success(self) -> None: + """Two 429s then a 200 -> returns the embedding list.""" + embedder = self._make_embedder() + mock_c = AsyncMock() + mock_c.__aenter__ = AsyncMock(return_value=mock_c) + mock_c.__aexit__ = AsyncMock(return_value=False) + mock_c.post = AsyncMock( + side_effect=[ + _429_response(), + _429_response(), + _success_embed_response(), + ] + ) + + with ( + patch(_EMBED_PATH, return_value=mock_c), + patch("asyncio.sleep", new_callable=AsyncMock), + ): + result = await embedder.aembed_query("hello") + + assert isinstance(result, list) + assert len(result) == _EMBED_DIM + + async def test_connect_error_does_not_trigger_429_retry_path( + self, + ) -> None: + """ConnectError -> OllamaConnectionError after MAX_RETRIES=3, not 5.""" + embedder = self._make_embedder() + mock_c = AsyncMock() + mock_c.__aenter__ = AsyncMock(return_value=mock_c) + mock_c.__aexit__ = AsyncMock(return_value=False) + mock_c.post = AsyncMock(side_effect=httpx.ConnectError("refused")) + + with ( + patch(_EMBED_PATH, return_value=mock_c), + patch("asyncio.sleep", new_callable=AsyncMock), + pytest.raises(OllamaConnectionError), + ): + await embedder.aembed_query("hello") + + assert mock_c.post.call_count == MAX_RETRIES + + async def test_timeout_does_not_trigger_429_retry_path(self) -> None: + """TimeoutException -> OllamaConnectionError after MAX_RETRIES=3.""" + embedder = self._make_embedder() + mock_c = AsyncMock() + mock_c.__aenter__ = AsyncMock(return_value=mock_c) + mock_c.__aexit__ = AsyncMock(return_value=False) + mock_c.post = AsyncMock(side_effect=httpx.TimeoutException("timed out")) + + with ( + patch(_EMBED_PATH, return_value=mock_c), + patch("asyncio.sleep", new_callable=AsyncMock), + pytest.raises(OllamaConnectionError), + ): + await embedder.aembed_query("hello") + + assert mock_c.post.call_count == MAX_RETRIES + + +# =========================================================================== +# 4. OllamaEmbedder - sync path (embed_query) +# =========================================================================== + + +class TestOllamaEmbedderSync: + """Tests for the synchronous embed_query method.""" + + def _make_embedder(self) -> OllamaEmbedder: + return OllamaEmbedder(base_url="http://ollama-test:11434") + + def test_five_consecutive_429s_raise_rate_limit_error(self) -> None: + embedder = self._make_embedder() + mock_client = MagicMock() + mock_client.post.return_value = _429_response() + + with ( + patch.object(embedder, "_get_sync_client", return_value=mock_client), + patch("time.sleep"), + pytest.raises(RateLimitError) as exc_info, + ): + embedder.embed_query("hello") + + assert exc_info.value.provider == "ollama" + assert mock_client.post.call_count == MAX_RATE_LIMIT_RETRIES + + def test_retry_after_header_respected_as_sleep_duration(self) -> None: + embedder = self._make_embedder() + sleep_calls: list[float] = [] + + mock_client = MagicMock() + mock_client.post.return_value = _429_response(retry_after="9") + + with ( + patch.object(embedder, "_get_sync_client", return_value=mock_client), + patch("time.sleep", side_effect=sleep_calls.append), + pytest.raises(RateLimitError), + ): + embedder.embed_query("hello") + + assert all(s == _RETRY_AFTER_9 for s in sleep_calls) + assert len(sleep_calls) == MAX_RATE_LIMIT_RETRIES - 1 + + def test_partial_retries_then_success(self) -> None: + embedder = self._make_embedder() + mock_client = MagicMock() + mock_client.post.side_effect = [_429_response(), _success_embed_response()] + + with ( + patch.object(embedder, "_get_sync_client", return_value=mock_client), + patch("time.sleep"), + ): + result = embedder.embed_query("hello") + + assert isinstance(result, list) + + def test_connect_error_does_not_trigger_rate_limit_retry(self) -> None: + embedder = self._make_embedder() + mock_client = MagicMock() + mock_client.post.side_effect = httpx.ConnectError("refused") + + with ( + patch.object(embedder, "_get_sync_client", return_value=mock_client), + patch("time.sleep"), + pytest.raises(OllamaConnectionError), + ): + embedder.embed_query("hello") + + assert mock_client.post.call_count == MAX_RETRIES + + +# =========================================================================== +# 5. OllamaEmbedder - _embed_batch_sync +# =========================================================================== + + +class TestOllamaEmbedBatchSync: + def _make_embedder(self) -> OllamaEmbedder: + return OllamaEmbedder(base_url="http://ollama-test:11434") + + def test_five_429s_raise_rate_limit_error(self) -> None: + embedder = self._make_embedder() + mock_client = MagicMock() + mock_client.post.return_value = _429_response() + + with patch("time.sleep"), pytest.raises(RateLimitError): + embedder._embed_batch_sync(mock_client, ["a", "b"], batch_index=0) + + assert mock_client.post.call_count == MAX_RATE_LIMIT_RETRIES + + def test_connect_error_raises_connection_error_not_rate_limit(self) -> None: + embedder = self._make_embedder() + mock_client = MagicMock() + mock_client.post.side_effect = httpx.ConnectError("refused") + + with patch("time.sleep"), pytest.raises(OllamaConnectionError): + embedder._embed_batch_sync(mock_client, ["a"], batch_index=0) + + assert mock_client.post.call_count == MAX_RETRIES + + +# =========================================================================== +# 6. extraction.py - Anthropic rate-limit retry +# =========================================================================== + + +class TestExtractionAnthropicRetry: + """Tests for ExtractionService.extract_with_llm Anthropic retry logic. + + AsyncAnthropic is imported *inside* extract_with_llm, so we patch at the + anthropic module level. The client is NOT used as a context manager there. + """ + + async def test_five_rate_limit_errors_raises_rate_limit_error(self) -> None: + """RateLimitError raised 5 times -> our RateLimitError propagated.""" + svc = ExtractionService() + api_exc = _make_anthropic_rl_exc(retry_after="5") + + mock_client = MagicMock() + mock_client.messages.create = AsyncMock(side_effect=api_exc) + + with ( + patch("anthropic.AsyncAnthropic") as mock_cls, + patch("roboco.config.settings") as mock_settings, + patch("asyncio.sleep", new_callable=AsyncMock), + pytest.raises(RateLimitError) as exc_info, + ): + mock_cls.return_value = mock_client + mock_settings.anthropic_api_key = "test-key" + await svc.extract_with_llm(_make_extraction_context()) + + assert exc_info.value.provider == "anthropic" + assert mock_client.messages.create.call_count == MAX_RATE_LIMIT_RETRIES + + async def test_retry_after_header_drives_sleep_duration(self) -> None: + """Retry-After: 12 -> asyncio.sleep(12) called on each gap.""" + svc = ExtractionService() + sleep_calls: list[float] = [] + api_exc = _make_anthropic_rl_exc(retry_after="12") + + mock_client = MagicMock() + mock_client.messages.create = AsyncMock(side_effect=api_exc) + + async def _fake_sleep(secs: float) -> None: + sleep_calls.append(secs) + + with ( + patch("anthropic.AsyncAnthropic") as mock_cls, + patch("roboco.config.settings") as mock_settings, + patch("asyncio.sleep", side_effect=_fake_sleep), + pytest.raises(RateLimitError), + ): + mock_cls.return_value = mock_client + mock_settings.anthropic_api_key = "test-key" + await svc.extract_with_llm(_make_extraction_context()) + + assert all(s == _RETRY_AFTER_12 for s in sleep_calls) + assert len(sleep_calls) == MAX_RATE_LIMIT_RETRIES - 1 + + async def test_partial_rate_limit_then_success_returns_result(self) -> None: + """Two RateLimitErrors then success -> result returned, no raise.""" + svc = ExtractionService() + api_exc = _make_anthropic_rl_exc() + + text_block = MagicMock() + text_block.text = "[N,]{type,content,confidence}:\nreasoning,Hello world,0.9" + success_response = MagicMock() + success_response.content = [text_block] + + mock_client = MagicMock() + mock_client.messages.create = AsyncMock( + side_effect=[api_exc, api_exc, success_response] + ) + + raised = False + with ( + patch("anthropic.AsyncAnthropic") as mock_cls, + patch("roboco.config.settings") as mock_settings, + patch("asyncio.sleep", new_callable=AsyncMock), + ): + mock_cls.return_value = mock_client + mock_settings.anthropic_api_key = "test-key" + try: + result = await svc.extract_with_llm(_make_extraction_context()) + assert result is not None + except RateLimitError: + raised = True + + assert not raised, "RateLimitError raised even though 3rd attempt succeeded" + assert mock_client.messages.create.call_count == _CALLS_2RL_1_SUCCESS + + +# =========================================================================== +# 7. indexes/base.py - BaseIndexPlugin.ask() LLM 429 retry +# =========================================================================== + + +class TestIndexAsk429Retry: + """Tests for the LLM call in BaseIndexPlugin.ask().""" + + async def test_ask_raises_rate_limit_after_five_429s(self) -> None: + plugin = _make_journal_plugin() + mock_c = _make_async_client_mock(post_return=_429_response()) + + with ( + patch.object( + plugin, "search", AsyncMock(return_value=_make_search_outcome()) + ), + patch("httpx.AsyncClient", return_value=mock_c), + patch("asyncio.sleep", new_callable=AsyncMock), + pytest.raises(RateLimitError) as exc_info, + ): + await plugin.ask("what is X") + + assert exc_info.value.provider == "ollama" + assert mock_c.post.call_count == MAX_RATE_LIMIT_RETRIES + + async def test_ask_partial_retry_then_success(self) -> None: + plugin = _make_journal_plugin() + success_body = {"choices": [{"message": {"content": "Here is the answer."}}]} + success_resp = _make_response(200, success_body) + mock_c = AsyncMock() + mock_c.__aenter__ = AsyncMock(return_value=mock_c) + mock_c.__aexit__ = AsyncMock(return_value=False) + mock_c.post = AsyncMock(side_effect=[_429_response(), success_resp]) + + with ( + patch.object( + plugin, "search", AsyncMock(return_value=_make_search_outcome()) + ), + patch("httpx.AsyncClient", return_value=mock_c), + patch("asyncio.sleep", new_callable=AsyncMock), + ): + answer, results = await plugin.ask("what is X") + + assert answer == "Here is the answer." + assert results is not None + + +# =========================================================================== +# 8. mentor.py - MentorService._synthesize_answer() 429 retry +# =========================================================================== + + +class TestMentorSynthesizeAnswer429: + async def test_raises_rate_limit_after_five_429s(self) -> None: + mentor = MentorService() + mentor._optimal_service = MagicMock() + mock_c = _make_async_client_mock(post_return=_429_response()) + + with ( + patch("httpx.AsyncClient", return_value=mock_c), + patch("asyncio.sleep", new_callable=AsyncMock), + pytest.raises(RateLimitError) as exc_info, + ): + await mentor._synthesize_answer( + question="test", + sources=_make_sources(), + conversation_context="", + agent_profile=None, + journal_context=[], + ) + + assert exc_info.value.provider == "ollama" + assert mock_c.post.call_count == MAX_RATE_LIMIT_RETRIES + + async def test_partial_retry_then_success(self) -> None: + mentor = MentorService() + mentor._optimal_service = MagicMock() + success_body = {"choices": [{"message": {"content": "Great answer."}}]} + mock_c = AsyncMock() + mock_c.__aenter__ = AsyncMock(return_value=mock_c) + mock_c.__aexit__ = AsyncMock(return_value=False) + mock_c.post = AsyncMock( + side_effect=[_429_response(), _make_response(200, success_body)] + ) + + with ( + patch("httpx.AsyncClient", return_value=mock_c), + patch("asyncio.sleep", new_callable=AsyncMock), + ): + answer = await mentor._synthesize_answer( + question="test", + sources=_make_sources(), + conversation_context="", + agent_profile=None, + journal_context=[], + ) + + assert answer == "Great answer." + + +# =========================================================================== +# 9. validator.py - ValidatorService._validate_with_llm() 429 retry +# =========================================================================== + + +class TestValidatorLLMRetry: + async def test_raises_rate_limit_after_five_429s(self) -> None: + validator = ValidatorService() + validator._optimal_service = MagicMock() + validator._llm_available = True + mock_c = _make_async_client_mock(post_return=_429_response()) + + with ( + patch("httpx.AsyncClient", return_value=mock_c), + patch("asyncio.sleep", new_callable=AsyncMock), + pytest.raises(RateLimitError) as exc_info, + ): + await validator._validate_with_llm( + action_type="create_endpoint", + context="def foo(): pass", + standards=_make_standards(), + ) + + assert exc_info.value.provider == "ollama" + assert mock_c.post.call_count == MAX_RATE_LIMIT_RETRIES + + async def test_partial_retry_then_success(self) -> None: + validator = ValidatorService() + validator._optimal_service = MagicMock() + validator._llm_available = True + success_content = '{"violations": [], "summary": "ok"}' + success_body = {"choices": [{"message": {"content": success_content}}]} + mock_c = AsyncMock() + mock_c.__aenter__ = AsyncMock(return_value=mock_c) + mock_c.__aexit__ = AsyncMock(return_value=False) + mock_c.post = AsyncMock( + side_effect=[_429_response(), _make_response(200, success_body)] + ) + + with ( + patch("httpx.AsyncClient", return_value=mock_c), + patch("asyncio.sleep", new_callable=AsyncMock), + ): + violations, warnings = await validator._validate_with_llm( + action_type="create_endpoint", + context="def foo(): pass", + standards=_make_standards(), + ) + + assert violations == [] + assert warnings == [] diff --git a/tests/unit/services/test_rate_limit_tracker.py b/tests/unit/services/test_rate_limit_tracker.py new file mode 100644 index 00000000..145d473f --- /dev/null +++ b/tests/unit/services/test_rate_limit_tracker.py @@ -0,0 +1,244 @@ +"""Unit tests for RateLimitStateTracker. + +These tests use mock Redis clients (no real Redis server required) to +verify the state-management logic. The cross-reconnection persistence +test constructs *two* RateLimitStateTracker instances that share the +same mock Redis store, proving that state written by one instance is +visible to a fresh instance. +""" + +from __future__ import annotations + +import json +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from roboco.services.gateway.rate_limit_tracker import RateLimitStateTracker + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_redis_mock(initial_store: dict[str, Any] | None = None) -> AsyncMock: + """Build an async Redis mock backed by a plain dict. + + The mock supports ``get``, ``set``, and ``delete`` with the same + semantics as the real redis.asyncio.Redis client. + """ + # Use the dict AS-IS (no copy) so that two mocks sharing the same + # dict object see each other's writes and deletes — this is what the + # cross-reconnection persistence tests rely on. + store: dict[str, Any] = initial_store if initial_store is not None else {} + + async def _get(key: str) -> bytes | None: + val = store.get(key) + if val is None: + return None + if isinstance(val, bytes): + return val + return str(val).encode() + + async def _set(key: str, value: Any) -> None: + store[key] = value + + async def _delete(key: str) -> int: + return 1 if store.pop(key, None) is not None else 0 + + mock = AsyncMock() + mock.get = AsyncMock(side_effect=_get) + mock.set = AsyncMock(side_effect=_set) + mock.delete = AsyncMock(side_effect=_delete) + # Stash the backing store so tests can inspect raw state + mock._store = store + return mock + + +def _make_tracker( + provider: str = "anthropic", + redis_mock: AsyncMock | None = None, +) -> RateLimitStateTracker: + """Build a tracker with an injected mock Redis client.""" + tracker = RateLimitStateTracker(provider=provider, redis_url="redis://unused") + if redis_mock is not None: + tracker._redis = redis_mock # type: ignore[assignment] + return tracker + + +# --------------------------------------------------------------------------- +# Tests: basic operations +# --------------------------------------------------------------------------- + + +class TestActivateAndRead: + async def test_is_rate_limited_false_by_default(self) -> None: + mock = _make_redis_mock() + tracker = _make_tracker(redis_mock=mock) + assert await tracker.is_rate_limited() is False + + async def test_get_state_empty_by_default(self) -> None: + mock = _make_redis_mock() + tracker = _make_tracker(redis_mock=mock) + assert await tracker.get_state() == {} + + async def test_activate_sets_rate_limited_true(self) -> None: + mock = _make_redis_mock() + tracker = _make_tracker(redis_mock=mock) + await tracker.activate() + assert await tracker.is_rate_limited() is True + + async def test_activate_stores_retry_after(self) -> None: + mock = _make_redis_mock() + tracker = _make_tracker(redis_mock=mock) + await tracker.activate(retry_after=30.0) + state = await tracker.get_state() + assert state["retry_after"] == 30.0 + + async def test_activate_stores_affected_agents(self) -> None: + mock = _make_redis_mock() + tracker = _make_tracker(redis_mock=mock) + await tracker.activate(affected_agents=["be-dev-1", "be-dev-2"]) + state = await tracker.get_state() + assert state["affected_agents"] == ["be-dev-1", "be-dev-2"] + + async def test_activate_initialises_probe_failures_zero(self) -> None: + mock = _make_redis_mock() + tracker = _make_tracker(redis_mock=mock) + await tracker.activate() + state = await tracker.get_state() + assert state["probe_failures"] == 0 + + async def test_clear_removes_state(self) -> None: + mock = _make_redis_mock() + tracker = _make_tracker(redis_mock=mock) + await tracker.activate() + await tracker.clear() + assert await tracker.is_rate_limited() is False + assert await tracker.get_state() == {} + + +class TestProbeFailures: + async def test_increment_starts_from_zero(self) -> None: + mock = _make_redis_mock() + tracker = _make_tracker(redis_mock=mock) + await tracker.activate() + count = await tracker.increment_probe_failures() + assert count == 1 + + async def test_increment_accumulates(self) -> None: + mock = _make_redis_mock() + tracker = _make_tracker(redis_mock=mock) + await tracker.activate() + await tracker.increment_probe_failures() + await tracker.increment_probe_failures() + count = await tracker.increment_probe_failures() + assert count == 3 + + async def test_reset_sets_zero(self) -> None: + mock = _make_redis_mock() + tracker = _make_tracker(redis_mock=mock) + await tracker.activate() + await tracker.increment_probe_failures() + await tracker.increment_probe_failures() + await tracker.reset_probe_failures() + state = await tracker.get_state() + assert state["probe_failures"] == 0 + + +# --------------------------------------------------------------------------- +# Tests: cross-reconnection persistence +# --------------------------------------------------------------------------- +# +# AC2: "State persists across client reconnection: a test writes state via +# activate(), creates a new RateLimitStateTracker instance pointing at the +# same Redis URL, calls is_rate_limited() and get_state() and gets back the +# same values — proving state survives a process restart." +# +# We simulate this by sharing the same backing dict between two mock Redis +# clients — one injected into the first tracker and one injected into the +# second. Both clients read from and write to the same dict, so the second +# tracker "sees" everything the first wrote. +# --------------------------------------------------------------------------- + + +class TestStatePersistsAcrossReconnection: + async def test_is_rate_limited_survives_reconnection(self) -> None: + shared_store: dict[str, Any] = {} + + # First "connection": write rate-limit state + mock_a = _make_redis_mock(initial_store=shared_store) + tracker_a = _make_tracker(provider="anthropic", redis_mock=mock_a) + await tracker_a.activate(retry_after=60.0, affected_agents=["be-dev-1"]) + + # The mock writes into shared_store directly (our _set stores raw). + # We need to seed the second mock from the same backing store. + # Because mock_a._store IS shared_store (same dict object), we only + # need to give mock_b access to the same dict. + mock_b = _make_redis_mock(initial_store=mock_a._store) + tracker_b = RateLimitStateTracker( + provider="anthropic", redis_url="redis://unused" + ) + tracker_b._redis = mock_b # type: ignore[assignment] + + assert await tracker_b.is_rate_limited() is True + + async def test_get_state_survives_reconnection(self) -> None: + shared_store: dict[str, Any] = {} + + mock_a = _make_redis_mock(initial_store=shared_store) + tracker_a = _make_tracker(provider="anthropic", redis_mock=mock_a) + await tracker_a.activate(retry_after=45.0, affected_agents=["be-dev-2"]) + + mock_b = _make_redis_mock(initial_store=mock_a._store) + tracker_b = RateLimitStateTracker( + provider="anthropic", redis_url="redis://unused" + ) + tracker_b._redis = mock_b # type: ignore[assignment] + + state = await tracker_b.get_state() + assert state["rate_limited"] is True + assert state["retry_after"] == 45.0 + assert state["affected_agents"] == ["be-dev-2"] + + async def test_clear_via_first_instance_visible_to_second(self) -> None: + shared_store: dict[str, Any] = {} + + mock_a = _make_redis_mock(initial_store=shared_store) + tracker_a = _make_tracker(provider="anthropic", redis_mock=mock_a) + await tracker_a.activate() + + # Second instance points at the same store + mock_b = _make_redis_mock(initial_store=mock_a._store) + tracker_b = RateLimitStateTracker( + provider="anthropic", redis_url="redis://unused" + ) + tracker_b._redis = mock_b # type: ignore[assignment] + + # Write clear via tracker_a + await tracker_a.clear() + + # tracker_b observes the cleared state + assert await tracker_b.is_rate_limited() is False + assert await tracker_b.get_state() == {} + + +# --------------------------------------------------------------------------- +# Tests: different providers are isolated +# --------------------------------------------------------------------------- + + +class TestProviderIsolation: + async def test_activating_one_provider_does_not_affect_another(self) -> None: + store: dict[str, Any] = {} + mock_a = _make_redis_mock(initial_store=store) + mock_b = _make_redis_mock(initial_store=store) + + tracker_anthropic = _make_tracker(provider="anthropic", redis_mock=mock_a) + tracker_ollama = _make_tracker(provider="ollama_cloud", redis_mock=mock_b) + + await tracker_anthropic.activate() + assert await tracker_anthropic.is_rate_limited() is True + assert await tracker_ollama.is_rate_limited() is False