# backend/services/token_manager.py """ OMEGA Token Limit Checkpoint & Resume System When any Gemini 3.5 Flash call hits a ResourceExhausted (429) token limit: 1. Saves the EXACT task state + all partial output generated so far to SQLite 2. Launches a background monitor that probes the API every 60s until tokens refresh 3. On refresh, resumes the task from the EXACT LAST WORD — not from the beginning All Gemini call sites (research, auto_upgrade, captcha, gate check) use gemini_call_with_checkpoint() instead of calling the API directly. """ import asyncio import logging import os import sqlite3 import time import uuid from typing import Optional import google.generativeai as genai from backend.services.usb_vault import KeyDomain, resolve_vault_key def get_runtime_location() -> str: return "cloud" if os.environ.get("SPACE_ID") else "pc" # Only the conversational path should be shaped in the assistant's voice. # Structured task types (blueprint JSON, code implementation, research digests, # captcha/osint payloads) must stay clean — a "You are JARVIS…" preamble there # corrupts JSON/code output. _PERSONA_TASK_TYPES = {"general", "gaming"} def _persona_system_prompt(persona: str, task_type: str = "general") -> str: """The identity prompt for the EXPLICITLY requested persona. The mobile app (and other callers) pass the persona they want per-turn, so we select on that string directly rather than the persisted global mode. Without this, conversational LLM replies carry no JARVIS/FRIDAY character — the identity file that is supposed to be the source of truth was being ignored on the primary /api/chat path. Returns "" for non-conversational task types so structured output is never polluted. """ if task_type not in _PERSONA_TASK_TYPES: return "" try: from modules.assistant_identity import ( JARVIS_PERSONALITY_PROMPT, FRIDAY_PERSONALITY_PROMPT, ) p = (persona or "jarvis").strip().lower() return JARVIS_PERSONALITY_PROMPT if p == "jarvis" else FRIDAY_PERSONALITY_PROMPT except Exception: return "" # Load all 15 keys into the token manager pool GOOGLE_API_KEYS = {} try: for domain in KeyDomain: try: key_val = resolve_vault_key(domain) if key_val: GOOGLE_API_KEYS[domain.value] = key_val except Exception: continue if get_runtime_location() == "cloud": GEMINI_API_KEY = GOOGLE_API_KEYS.get(KeyDomain.CLOUD.value, "") else: GEMINI_API_KEY = GOOGLE_API_KEYS.get(KeyDomain.CHAT.value, "") if GEMINI_API_KEY: genai.configure(api_key=GEMINI_API_KEY) else: logging.warning("TokenManager: Default API Key is empty.") logging.info(f"TokenManager: Successfully loaded {len(GOOGLE_API_KEYS)} Google/Gemini API keys into the rotation pool.") except Exception as e: logging.warning(f"TokenManager: API Key initialization failed: {e}") GEMINI_API_KEY = "" EXHAUSTED_GOOGLE_KEYS = set() # Measured live against a real key, 2026-07-21: # gemini-3.5-flash HTTP 503 "currently experiencing high demand" <- was set here # gemini-2.0-flash HTTP 429 free-tier request cap # gemini-2.5-flash-lite HTTP 404 retired # gemini-3-flash-preview OK but 35.5s # gemini-3.1-flash-lite OK in 1.1s, correct answer <- chosen # # The 503 is Google-side capacity, NOT quota consumption: these keys were unused # for weeks. Because the primary model always failed, every request fell through # to the NVIDIA tier — which is why chat worked at all, and why research paid for # a doomed Gemini attempt before doing any real work. MODEL = "gemini-3.1-flash-lite" # Google-side model cascade, mirroring the NVIDIA tier cascade in nvidia_vault. # # A single hardcoded model is fragile in exactly two ways this system has already # been bitten by: the model gets retired (gemini-2.5-flash-lite -> 404) or it hits # Google-side capacity (gemini-3.5-flash -> 503 "high demand"). Neither means the # KEY is bad, and neither should cost a drop to the fallback tier — another # Gemini model is usually fine right now. # # Ordered by what was measured live 2026-07-21 against a real key: # gemini-3.1-flash-lite OK, 1.1s, correct answer <- best right now # gemini-3.5-flash 503 today, but transient — capacity comes and goes # gemini-3-flash-preview OK but 35.5s # gemini-2.0-flash 429 free-tier cap today # The list is tried in order on any transient failure, and only when every entry # fails does the request drop to the NVIDIA tier. GEMINI_MODEL_CASCADE = [ MODEL, "gemini-3.5-flash", "gemini-3-flash-preview", "gemini-2.5-flash", "gemini-2.0-flash", ] def _is_transient_error(err_str: str) -> bool: """Transient upstream failure — worth trying another model/tier.""" return any(t in err_str for t in ( "resourceexhausted", "429", "quota", "rate limit", "too many requests", "503", "500", "high demand", "overloaded", "unavailable", "internal error", "deadline", "timeout", "404", "not found", "no longer available", )) # ───────────────────────────────────────────────────────────────────────────── # DB Helpers # ───────────────────────────────────────────────────────────────────────────── from backend.services.usb_monitor import get_db_path def _init_checkpoint_table(): """Ensure token_checkpoints table exists in memory.db.""" try: with sqlite3.connect(get_db_path()) as conn: conn.execute("PRAGMA journal_mode=WAL") conn.execute(""" CREATE TABLE IF NOT EXISTS token_checkpoints ( id TEXT PRIMARY KEY, task_type TEXT NOT NULL, original_prompt TEXT NOT NULL, partial_result TEXT DEFAULT '', last_word TEXT DEFAULT '', status TEXT DEFAULT 'pending_resume', created_at INTEGER, resumed_at INTEGER, completed_at INTEGER, retry_count INTEGER DEFAULT 0, persona TEXT DEFAULT 'jarvis' ) """) conn.commit() except Exception as e: logging.error(f"TokenManager: Failed to init checkpoint table: {e}") _init_checkpoint_table() def _save_checkpoint(task_id: str, task_type: str, prompt: str, partial: str, last_word: str, persona: str = "jarvis"): try: with sqlite3.connect(get_db_path()) as conn: conn.execute("PRAGMA journal_mode=WAL") conn.execute(""" INSERT OR REPLACE INTO token_checkpoints (id, task_type, original_prompt, partial_result, last_word, status, created_at, persona) VALUES (?,?,?,?,?,'pending_resume',?,?) """, (task_id, task_type, prompt, partial, last_word, int(time.time()), persona)) conn.commit() logging.warning(f"TokenManager: Checkpoint saved [{task_type}:{task_id}] last_word='{last_word}'") except Exception as e: logging.error(f"TokenManager: Failed to save checkpoint: {e}") def _mark_checkpoint_resumed(task_id: str): try: with sqlite3.connect(get_db_path()) as conn: conn.execute( "UPDATE token_checkpoints SET status='resumed', resumed_at=? WHERE id=?", (int(time.time()), task_id) ) conn.commit() except Exception as e: logging.error(f"TokenManager: Failed to mark resumed: {e}") def _mark_checkpoint_complete(task_id: str): try: with sqlite3.connect(get_db_path()) as conn: conn.execute( "UPDATE token_checkpoints SET status='complete', completed_at=? WHERE id=?", (int(time.time()), task_id) ) conn.commit() except Exception as e: logging.error(f"TokenManager: Failed to mark complete: {e}") def _get_pending_checkpoints() -> list: try: with sqlite3.connect(get_db_path()) as conn: conn.row_factory = sqlite3.Row rows = conn.execute( "SELECT * FROM token_checkpoints WHERE status='pending_resume' ORDER BY created_at ASC" ).fetchall() return [dict(r) for r in rows] except Exception as e: logging.error(f"TokenManager: Failed to fetch pending checkpoints: {e}") return [] # ───────────────────────────────────────────────────────────────────────────── # Core Gemini Call — with auto checkpoint on 429 # ───────────────────────────────────────────────────────────────────────────── def _extract_last_word(text: str) -> str: """Return the very last word of the partial output for precise resume anchoring.""" if not text: return "" words = text.strip().split() return words[-1] if words else "" def _build_resume_prompt(original_prompt: str, partial_result: str, last_word: str) -> str: """ Construct a continuation prompt so Gemini resumes from the EXACT last word. """ if not partial_result: return original_prompt return ( f"{original_prompt}\n\n" f"---OMEGA RESUME INSTRUCTION---\n" f"You were previously generating a response and hit a token limit mid-way.\n" f"The partial output generated so far ended with the word: '{last_word}'\n" f"Here is the partial output so far:\n\n{partial_result}\n\n" f"Continue EXACTLY from where you left off. " f"Do NOT repeat any of the partial output above. " f"Start your response from the word that comes AFTER '{last_word}'." ) def _get_key_for_task(task_type: str) -> str: from backend.services.usb_vault import KeyDomain, resolve_vault_key location = get_runtime_location() if task_type == "research": return resolve_vault_key(KeyDomain.RESEARCH_ENGINE) elif task_type == "captcha": return resolve_vault_key(KeyDomain.CAPTCHA_SOLVER_CLOUD if location == "cloud" else KeyDomain.CAPTCHA_SOLVER_PC) elif task_type == "implement": return resolve_vault_key(KeyDomain.AUTO_UPGRADE_CLOUD if location == "cloud" else KeyDomain.AUTO_UPGRADE_PC) elif task_type == "osint": return resolve_vault_key(KeyDomain.OSINT_PROTOCOL_CLOUD if location == "cloud" else KeyDomain.OSINT_PROTOCOL_PC) elif task_type == "image": return resolve_vault_key(KeyDomain.IMAGE_GENERATION_CLOUD if location == "cloud" else KeyDomain.IMAGE_GENERATION_PC) elif task_type == "gaming": return resolve_vault_key(KeyDomain.GAMING_COACH) elif task_type == "supervisor": return resolve_vault_key(KeyDomain.SUPERVISOR_HEAL_CLOUD if location == "cloud" else KeyDomain.SUPERVISOR_HEAL_PC) elif task_type == "cloud_sync": return resolve_vault_key(KeyDomain.CLOUD_SYNC_DOMAIN) else: return resolve_vault_key(KeyDomain.CLOUD if location == "cloud" else KeyDomain.CHAT) def _broadcast_live_key_status(ecosystem: str, model: str, status: str): """Safely fire WS broadcast from synchronous threads""" payload = { "event": "live_key_status", "payload": { "active_ecosystem": ecosystem, "active_model": model, "status": status } } try: from backend.ws.agent_ws import ws_manager loop = asyncio.get_running_loop() loop.create_task(ws_manager.broadcast(payload)) except RuntimeError: pass # No loop running in this thread def _nvidia_fallback_call_sync(prompt: str, task_id: str, task_type: str, partial_so_far: str, persona: str) -> str: """Fallback to NVIDIA Vault Models.""" _broadcast_live_key_status("NVIDIA", "glm-5.1", "fallback_mode") from backend.services.memory_service import GlobalOmniMemory omni_context = GlobalOmniMemory.get_global_context_stream() # Construct exact prompt with resume logic and hive mind context actual_prompt = _build_resume_prompt(prompt, partial_so_far, _extract_last_word(partial_so_far)) if not partial_so_far: # NVIDIA models take no separate system channel here, so the persona # prompt is prepended (same place the hive-mind context goes). _sys = _persona_system_prompt(persona, task_type) actual_prompt = (_sys + "\n\n" if _sys else "") + omni_context + actual_prompt # Delegate to the NVIDIA Vault API directly from backend.services.nvidia_vault import call_nvidia_model, NvidiaModel # GLM is the intended head of the heavy-compute fallback tier. call_nvidia_model # already cascades from here through the rest of that tier (MiniMax, DeepSeek, # Nemotron-Ultra) with up to 3 key-failovers per model, so a single slow or # unavailable model does not strand the request. # # Only the model ID changed: "z-ai/glm-5.1" was retired upstream and returned # HTTP 410 on all 15 keys, which wasted the first attempt of every heavy-tier # call. It is now glm-5.2. fallback_response = call_nvidia_model(actual_prompt, NvidiaModel.GLM_5_2) # Record what NVIDIA just did to the Hive Mind GlobalOmniMemory.record_action("NVIDIA", "glm-5.1", task_type, fallback_response[:100].replace('\n', ' ') + "...") combined = partial_so_far + fallback_response _mark_checkpoint_complete(task_id) return combined def _gemini_call_sync(prompt: str, task_id: str, task_type: str, partial_so_far: str = "", persona: str = "jarvis") -> str: key = _get_key_for_task(task_type) # 1. If we know it's exhausted, fallback immediately if key in EXHAUSTED_GOOGLE_KEYS: logging.info(f"TokenManager: Google Key exhausted. Routing {task_type} instantly to NVIDIA Fallback.") return _nvidia_fallback_call_sync(prompt, task_id, task_type, partial_so_far, persona) genai.configure(api_key=key) _broadcast_live_key_status("GOOGLE", MODEL, "primary_active") from backend.services.memory_service import GlobalOmniMemory omni_context = GlobalOmniMemory.get_global_context_stream() # Shape the reply in the active assistant's voice (JARVIS/FRIDAY). Skip on a # resume continuation so we don't restate the persona mid-sentence. _sys = _persona_system_prompt(persona, task_type) if not partial_so_far else None actual_prompt = _build_resume_prompt(prompt, partial_so_far, _extract_last_word(partial_so_far)) if not partial_so_far: actual_prompt = omni_context + actual_prompt # Walk the Gemini cascade before considering the NVIDIA tier. A model that is # retired or momentarily at capacity must not cost the whole Google side — # the next model in the list is usually healthy. last_exc = None for attempt_model in GEMINI_MODEL_CASCADE: try: model = (genai.GenerativeModel(attempt_model, system_instruction=_sys) if _sys else genai.GenerativeModel(attempt_model)) response = model.generate_content(actual_prompt) new_text = response.text or "" if attempt_model != MODEL: logging.warning("TokenManager: primary model unavailable; served by %s", attempt_model) _broadcast_live_key_status("GOOGLE", attempt_model, "primary_active") # Record what Google just did to the Hive Mind GlobalOmniMemory.record_action("GOOGLE", attempt_model, task_type, new_text[:100].replace('\n', ' ') + "...") combined = partial_so_far + new_text _mark_checkpoint_complete(task_id) return combined except Exception as exc: last_exc = exc if _is_transient_error(str(exc).lower()): logging.info("TokenManager: %s unavailable (%s); trying next Gemini model.", attempt_model, str(exc)[:70]) continue raise # Every Gemini model failed — fall through to the original handling, which # decides between key-exhaustion bookkeeping and the NVIDIA tier. try: raise last_exc except Exception as e: err_str = str(e).lower() # Any TRANSIENT upstream failure should fall back to NVIDIA, not just a # quota error. This previously matched 429/quota only, so a Google-side # capacity 503 ("This model is currently experiencing high demand" — which # gemini-3.5-flash returns constantly) fell through to `raise` and killed # the request outright, never reaching the NVIDIA tier that exists for # exactly this situation. is_rate_limit = ( "resourceexhausted" in err_str or "429" in err_str or "quota" in err_str or "rate limit" in err_str or "too many requests" in err_str # Google-side capacity / availability, not our usage: or "503" in err_str or "500" in err_str or "high demand" in err_str or "overloaded" in err_str or "unavailable" in err_str or "internal error" in err_str or "deadline" in err_str or "timeout" in err_str ) # A capacity/availability failure is the MODEL's problem, not the key's. # Blacklisting the key for a 503 would permanently route a perfectly # healthy credential to the fallback tier — so only real quota errors # mark a key exhausted. Both still fall back to NVIDIA. is_key_exhausted = ( "resourceexhausted" in err_str or "429" in err_str or "quota" in err_str or "rate limit" in err_str or "too many requests" in err_str ) if is_rate_limit: logging.warning( "TokenManager: transient upstream failure for [%s:%s] (%s). Triggering NVIDIA Fallback.", task_type, task_id, "quota" if is_key_exhausted else "capacity/503", ) if not is_key_exhausted: # Model-side outage: keep the key healthy, just serve this request # from the fallback tier. return _nvidia_fallback_call_sync(prompt, task_id, task_type, partial_so_far, persona) # 2. Add to exhausted pool if key not in EXHAUSTED_GOOGLE_KEYS: EXHAUSTED_GOOGLE_KEYS.add(key) # We can't spawn an async task cleanly from a sync thread without a loop, # so we will raise a special exception that the async wrapper catches to spawn the pinger. raise _Google429Trigger(key=key, prompt=prompt, task_id=task_id, task_type=task_type, partial_so_far=partial_so_far, persona=persona) return _nvidia_fallback_call_sync(prompt, task_id, task_type, partial_so_far, persona) raise class _Google429Trigger(Exception): def __init__(self, key, prompt, task_id, task_type, partial_so_far, persona): self.key = key self.prompt = prompt self.task_id = task_id self.task_type = task_type self.partial_so_far = partial_so_far self.persona = persona TokenLimitHit = _Google429Trigger async def gemini_call_with_checkpoint( prompt: str, task_type: str = "general", persona: str = "jarvis", task_id: Optional[str] = None, partial_so_far: str = "" ) -> str: """ Universal Gemini 3.5 Flash call with automatic checkpoint & resume. """ if task_id is None: task_id = str(uuid.uuid4()) try: return await asyncio.to_thread( _gemini_call_sync, prompt, task_id, task_type, partial_so_far, persona ) except _Google429Trigger as e: # 1. Start the token refresh monitor loop asyncio.create_task(start_token_refresh_monitor()) # 2. Immediately complete THIS task using NVIDIA fallback return await asyncio.to_thread( _nvidia_fallback_call_sync, e.prompt, e.task_id, e.task_type, e.partial_so_far, e.persona ) # For backward compatibility where 'run_task' is imported from token_manager async def run_task(prompt: str, task_type: str = "general", persona: str = "jarvis") -> str: return await gemini_call_with_checkpoint(prompt, task_type, persona) # ───────────────────────────────────────────────────────────────────────────── # Background Cloud Monitor — polls until tokens refresh then resumes # ───────────────────────────────────────────────────────────────────────────── _monitor_running = False async def start_token_refresh_monitor(): """ Background loop that: 1. Checks for pending_resume checkpoints every 60s 2. Probes Gemini with a minimal token-cost ping 3. When the probe succeeds, resumes ALL pending tasks from their exact last word 4. Broadcasts omega:task_resumed WS event for each resumed task """ global _monitor_running if _monitor_running: return _monitor_running = True logging.info("TokenManager: Token refresh monitor STARTED.") while _monitor_running: try: pending = _get_pending_checkpoints() # We must probe if there are pending checkpoints OR if we have exhausted keys that need refreshing if pending or EXHAUSTED_GOOGLE_KEYS: if pending: logging.info(f"TokenManager: {len(pending)} checkpoint(s) pending resume. Probing API...") else: logging.info("TokenManager: No checkpoints, but EXHAUSTED_GOOGLE_KEYS has entries. Probing API...") # Lightweight probe — minimal cost # Check each exhausted key independently keys_to_probe = list(EXHAUSTED_GOOGLE_KEYS) if EXHAUSTED_GOOGLE_KEYS else [_get_key_for_task("general")] recovered_keys = [] for key in keys_to_probe: try: genai.configure(api_key=key) model = genai.GenerativeModel(MODEL) model.generate_content("ping") recovered_keys.append(key) except Exception as probe_err: err_str = str(probe_err).lower() if "resourceexhausted" in err_str or "429" in err_str or "quota" in err_str: pass # Still exhausted else: logging.error(f"TokenManager: Probe error on key: {probe_err}") if not recovered_keys and keys_to_probe: logging.warning("TokenManager: API still exhausted. Will retry in 60s.") if recovered_keys: logging.info(f"TokenManager: API Probe OK for {len(recovered_keys)} keys. Clearing them from EXHAUSTED_GOOGLE_KEYS.") for k in recovered_keys: if k in EXHAUSTED_GOOGLE_KEYS: EXHAUSTED_GOOGLE_KEYS.remove(k) if pending: # Resume all pending checkpoints now that some keys are refreshed for checkpoint in pending: asyncio.create_task(_resume_checkpoint(checkpoint)) except Exception as e: logging.error(f"TokenManager: Monitor loop error: {e}") await asyncio.sleep(60) logging.info("TokenManager: Token refresh monitor STOPPED.") async def _resume_checkpoint(checkpoint: dict): """Resume a single checkpoint from its exact last word.""" task_id = checkpoint["id"] task_type = checkpoint["task_type"] original_prompt = checkpoint["original_prompt"] partial = checkpoint["partial_result"] or "" last_word = checkpoint["last_word"] or "" persona = checkpoint.get("persona", "jarvis") logging.info(f"TokenManager: Resuming [{task_type}:{task_id}] from last_word='{last_word}'") _mark_checkpoint_resumed(task_id) try: # Broadcast WS notification from backend.ws.agent_ws import ws_manager await ws_manager.broadcast({ "event": "omega:task_resuming", "payload": { "task_id": task_id, "task_type": task_type, "last_word": last_word, "message": ( f"Token limit has refreshed, sir. Resuming {task_type} task from exactly '{last_word}'..." if persona == "jarvis" else f"We're back boss! Tokens refreshed — resuming {task_type} from exactly '{last_word}'!" ) } }) # Re-run the Gemini call from exact checkpoint result = await gemini_call_with_checkpoint( prompt=original_prompt, task_type=task_type, persona=persona, task_id=task_id, partial_so_far=partial ) _mark_checkpoint_complete(task_id) # Route result back to its originating system await _dispatch_resumed_result(task_type, task_id, result, persona, checkpoint) await ws_manager.broadcast({ "event": "omega:task_resumed", "payload": { "task_id": task_id, "task_type": task_type, "message": ( f"{task_type.title()} task completed after token refresh, sir." if persona == "jarvis" else f"Done boss! {task_type.title()} task finished after the token refresh!" ) } }) except TokenLimitHit: # Still throttled — will retry next monitor cycle logging.warning(f"TokenManager: Still throttled on resume for [{task_id}]. Reverting to pending_resume.") try: with sqlite3.connect(get_db_path()) as conn: conn.execute( "UPDATE token_checkpoints SET status='pending_resume', retry_count=retry_count+1 WHERE id=?", (task_id,) ) conn.commit() except Exception as e: # No local `import logging` — it would shadow the module import for # the whole of _resume_checkpoint() and break the logging calls above. logging.getLogger(__name__).error(f"Swallowed exception: {e}") except Exception as e: logging.error(f"TokenManager: Resume error for [{task_id}]: {e}") async def _dispatch_resumed_result(task_type: str, task_id: str, result: str, persona: str, checkpoint: dict): """Route the completed result back to the correct subsystem.""" try: if task_type == "research": # Re-parse the JSON research note and save it import json from backend.omega.research_engine import ResearchNote, ResearchCategory, _save_research_note raw = result.strip().strip("```json").strip("```").strip() data = json.loads(raw) raw_category = data.get("category", "Research Notes") valid_categories = {c.value: c for c in ResearchCategory} category = valid_categories.get(raw_category, ResearchCategory.RESEARCH_NOTES) note = ResearchNote( title=data.get("title", "Resumed Research"), summary=data.get("summary", result[:500]), importance=data.get("importance", "MEDIUM"), category=category, source=data.get("source", "Gemini 3.5 Flash (resumed)"), recommended_action=data.get("recommended_action", "") ) _save_research_note(note) logging.info(f"TokenManager: Research note saved after resume: {note.title}") elif task_type == "implement": # Re-parse code blocks and hot-reload them from backend.omega.auto_upgrade import parse_code_blocks, write_and_hot_reload, log_upgrade_to_db code_blocks = parse_code_blocks(result) files_modified = [] for block in code_blocks: await write_and_hot_reload(block.filepath, block.code) files_modified.append(block.filepath) log_upgrade_to_db( checkpoint.get("original_prompt", "Resumed task")[:60], files_modified, persona, "success_after_resume" ) logging.info(f"TokenManager: Implementation resumed and hot-reloaded: {files_modified}") else: # For gate/captcha/monitor tasks, just log the full result logging.info(f"TokenManager: [{task_type}] resumed result stored (length={len(result)})") except Exception as e: logging.error(f"TokenManager: dispatch_resumed_result error for [{task_type}]: {e}") def stop_token_refresh_monitor(): global _monitor_running _monitor_running = False logging.info("TokenManager: Token refresh monitor signalled to stop.")