import asyncio from datetime import datetime, timezone from app.config import ( ROOT_FOLDER_ID, POLL_SECONDS, HEARTBEAT_FILENAME, GOOGLE_TOKEN_JSON, SERVICE_NAME, VERSION, ) from app.drive_client import build_drive_service, list_children, upsert_json_file def _utc_now_iso() -> str: return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z") def _make_hb_payload() -> dict: return { "service": SERVICE_NAME, "status": "alive", "last_seen": _utc_now_iso(), "version": VERSION, "root_folder_id": ROOT_FOLDER_ID, } async def worker_loop(): print("[STARTUP] Worker loop starting") if not GOOGLE_TOKEN_JSON or not GOOGLE_TOKEN_JSON.strip(): print("[ERROR] Missing GOOGLE_TOKEN_JSON (HF Secret).") return if not ROOT_FOLDER_ID or not ROOT_FOLDER_ID.strip(): print("[ERROR] Missing ROOT_FOLDER_ID (HF Variable).") return # Build Drive client once; we’ll rebuild on transient failures service = build_drive_service(GOOGLE_TOKEN_JSON) # ---------------------------- # 1) RUN ONCE IMMEDIATELY # ---------------------------- try: files = list_children(service, ROOT_FOLDER_ID, page_size=20) print(f"[DRIVE] Found {len(files)} files") if files: print("[DRIVE] Sample:", ", ".join([f["name"] for f in files[:5]])) upsert_json_file(service, ROOT_FOLDER_ID, HEARTBEAT_FILENAME, _make_hb_payload()) print(f"[HEARTBEAT] Wrote {HEARTBEAT_FILENAME} to Drive root (folder={ROOT_FOLDER_ID})") except Exception as e: print(f"[ERROR] Startup drive/heartbeat failed: {e}") # backoff so we don't thrash on startup await asyncio.sleep(10) # ---------------------------- # 2) POLL FOREVER # ---------------------------- while True: try: files = list_children(service, ROOT_FOLDER_ID, page_size=20) print(f"[DRIVE] Found {len(files)} files") if files: print("[DRIVE] Sample:", ", ".join([f["name"] for f in files[:5]])) upsert_json_file(service, ROOT_FOLDER_ID, HEARTBEAT_FILENAME, _make_hb_payload()) print(f"[HEARTBEAT] Updated {HEARTBEAT_FILENAME} (every {POLL_SECONDS}s)") except Exception as e: # Transient HF/network issues show up here (Broken pipe, EOF, etc.) print(f"[ERROR] Drive/heartbeat poll failed: {e}") # Retry strategy: # 1) Rebuild drive client (fresh TLS session) # 2) Retry heartbeat write once try: service = build_drive_service(GOOGLE_TOKEN_JSON) upsert_json_file(service, ROOT_FOLDER_ID, HEARTBEAT_FILENAME, _make_hb_payload()) print("[HEARTBEAT] Recovered after rebuild/retry") except Exception as e2: print(f"[ERROR] Recovery retry failed: {e2}") # small backoff so we don't spam on repeated failures await asyncio.sleep(10) # normal poll interval await asyncio.sleep(POLL_SECONDS)