feat: retry gaia task once after cooldown

src/lilith_agent/runner.py

@@ -4,6 +4,7 @@ import json
 import logging
 import os
 import re
+import time
 from pathlib import Path
 from typing import Any
 
@@ -159,11 +160,19 @@ def run_agent_on_questions(graph: Any, questions: list[dict], checkpoint_dir: st
     answers: list[dict] = []
 
     from lilith_agent.config import Config
-    from lilith_agent.models import get_cheap_model
+    from lilith_agent.models import RateLimitCooldownError, get_cheap_model, rate_limit_question_scope
     cfg = Config.from_env()
     cheap_model = get_cheap_model(cfg)
 
     total = len(questions)
+
+    def _invoke_task_once(task_state: dict, task_id: str):
+        from lilith_agent.memory import ephemeral_memory
+
+        with rate_limit_question_scope():
+            with ephemeral_memory():
+                return graph.invoke(task_state, {"configurable": {"thread_id": task_id}})
+
     for idx, question in enumerate(questions, start=1):
         reset_vision_state()
         task_id = question.get("task_id")
@@ -205,10 +214,23 @@ def run_agent_on_questions(graph: Any, questions: list[dict], checkpoint_dir: st
             "iterations": 0
         }
 
-        from lilith_agent.memory import ephemeral_memory
         try:
-            with ephemeral_memory():
-                result = graph.invoke(state, {"configurable": {"thread_id": task_id}})
+            try:
+                result = _invoke_task_once(state, task_id)
+            except RateLimitCooldownError as exc:
+                log_runner.warning(
+                    "[runner] task=%s rate limited provider=%s model=%s cooldown=%s",
+                    task_id,
+                    exc.provider,
+                    exc.model,
+                    exc.cooldown_seconds,
+                )
+                time.sleep(exc.cooldown_seconds)
+                result = _invoke_task_once(state, task_id)
+        except RateLimitCooldownError as exc:
+            log_runner.warning("[runner] task=%s rate limited after retry: %s", task_id, exc)
+            answers.append({"task_id": task_id, "submitted_answer": "AGENT ERROR: RATE LIMITED"})
+            continue
         except Exception as exc:
             log_runner.warning("[runner] task=%s agent error: %s", task_id, exc)
             answers.append(

tests/test_runner.py

@@ -5,7 +5,11 @@ from pathlib import Path
 
 import pytest
 
-from lilith_agent.runner import _wrap_user_question, _write_checkpoint_atomic
+from langchain_core.messages import AIMessage
+
+from lilith_agent.config import Config
+from lilith_agent.models import RateLimitCooldownError
+from lilith_agent.runner import run_agent_on_questions, _wrap_user_question, _write_checkpoint_atomic
 
 
 def test_wrap_escapes_closing_tag_to_prevent_injection():
@@ -57,3 +61,71 @@ def test_atomic_write_does_not_corrupt_existing_file_on_serialization_failure(tm
     data = json.loads(dest.read_text())
     assert data["submitted_answer"] == "good"
     assert list(tmp_path.glob("*.tmp")) == []
+
+
+@pytest.fixture
+def runner_test_config() -> Config:
+    return Config(
+        cheap_provider="google",
+        cheap_model="gemini-3-flash-preview",
+        strong_provider="google",
+        strong_model="gemini-3.1-pro",
+        extra_strong_provider="google",
+        extra_strong_model="gemini-3.1-pro",
+        vision_provider="fal",
+        vision_model="gemini-3-flash-preview",
+        fal_vision_api_key="",
+        api_url="",
+        checkpoint_dir="",
+        whisper_model="base",
+        anthropic_api_key="",
+        google_api_key="",
+        huggingface_api_key="",
+        tavily_api_key="",
+        lmstudio_base_url="",
+        max_tokens=1024,
+        llm_formatter_enabled=True,
+    )
+
+
+@pytest.fixture(autouse=True)
+def _isolate_runner_model_setup(monkeypatch, runner_test_config):
+    monkeypatch.setattr(Config, "from_env", classmethod(lambda cls: runner_test_config))
+    monkeypatch.setattr("lilith_agent.models.get_cheap_model", lambda cfg: object())
+
+
+class _GraphFailsOnceWithCooldown:
+    def __init__(self):
+        self.calls = 0
+        self.thread_ids = []
+
+    def invoke(self, state, config):
+        self.calls += 1
+        self.thread_ids.append(config["configurable"]["thread_id"])
+        if self.calls == 1:
+            raise RateLimitCooldownError(
+                provider="google",
+                model="gemini-3.1-pro",
+                cooldown_seconds=12,
+                original_error="429",
+            )
+        return {"messages": [AIMessage(content="Final Answer: 42")]}
+
+
+def test_runner_retries_same_question_once_after_cooldown(monkeypatch, tmp_path: Path):
+    monkeypatch.setattr("lilith_agent.runner._final_formatting_cleanup", lambda model, question, raw, llm_formatter_enabled=True: raw)
+    sleeps = []
+    monkeypatch.setattr("lilith_agent.runner.time.sleep", sleeps.append)
+    graph = _GraphFailsOnceWithCooldown()
+
+    answers = run_agent_on_questions(
+        graph,
+        [{"task_id": "task-1", "question": "What is 6*7?"}],
+        tmp_path,
+    )
+
+    assert graph.calls == 2
+    assert graph.thread_ids == ["task-1", "task-1"]
+    assert sleeps == [12]
+    assert answers == [{"task_id": "task-1", "submitted_answer": "Final Answer: 42"}]
+    assert (tmp_path / "task-1.json").exists()