Upload managed_auth.py

managed_auth.py

@@ -26,6 +26,8 @@ EDGE_ASSERTION_COOKIE_NAME = "smartpaper_edge_assertion"
 EDGE_ASSERTION_AUDIENCE = "smartpaper-streamlit"
 ASSERTION_CLOCK_SKEW_SECONDS = 30
 REQUEST_TIMEOUT = (5.0, 15.0)
+TASK_START_TIMEOUT = (8.0, 20.0)
+TASK_START_RETRY_DELAY_SECONDS = 0.5
 
 _HTTP_SESSION = requests.Session()
 _HTTP_SESSION.headers.update({
@@ -34,6 +36,16 @@ _HTTP_SESSION.headers.update({
     "User-Agent": "SmartPaper-HuggingFace/edge-first-v1",
 })
 
+# /task/start uses a direct connection first, then one environment-aware
+# retry. Both attempts reuse the exact same idempotency_key and request body.
+_TASK_DIRECT_SESSION = requests.Session()
+_TASK_DIRECT_SESSION.trust_env = False
+_TASK_DIRECT_SESSION.headers.update({
+    "Accept": "application/json",
+    "Content-Type": "application/json",
+    "User-Agent": "SmartPaper-HuggingFace/edge-first-v1",
+})
+
 
 def _runtime_value(name: str) -> str:
     value = os.environ.get(name, "")
@@ -200,8 +212,63 @@ def _authorized_post(path: str, payload: dict[str, Any] | None = None) -> dict[s
     return _post(path, body)
 
 
+def _post_task_start(payload: dict[str, Any]) -> dict[str, Any]:
+    """Start one task with one bounded retry and the same idempotency key."""
+    url = f"{_api_base()}/task/start"
+    headers = {"Authorization": f"Bearer {_service_key()}"}
+    transports = (_TASK_DIRECT_SESSION, _HTTP_SESSION)
+    last_error: Exception | None = None
+
+    for attempt, session in enumerate(transports):
+        try:
+            response = session.post(
+                url,
+                headers=headers,
+                json=payload,
+                timeout=TASK_START_TIMEOUT,
+                allow_redirects=False,
+            )
+        except (requests.exceptions.Timeout, requests.exceptions.ConnectionError) as exc:
+            last_error = exc
+            if attempt == 0:
+                time.sleep(TASK_START_RETRY_DELAY_SECONDS)
+                continue
+            raise
+
+        try:
+            data = response.json()
+        except ValueError as exc:
+            raise RuntimeError(
+                f"Authentication service returned non-JSON content (HTTP {response.status_code})."
+            ) from exc
+
+        if response.status_code >= 400:
+            message = (
+                data.get("message")
+                or data.get("reason")
+                or data.get("error")
+                or f"HTTP {response.status_code}"
+            )
+            raise RuntimeError(str(message))
+
+        return data
+
+    if last_error is not None:
+        raise last_error
+    raise RuntimeError("task_start_request_failed")
+
+
 def start_generation_task(*, idempotency_key: str, mode: str, interface_language: str) -> dict[str, Any]:
-    return _authorized_post("/task/start", {"idempotency_key": idempotency_key, "mode": mode, "interface_language": interface_language})
+    session_token = str(st.session_state.get("sp_session_token") or "").strip()
+    if not session_token:
+        raise RuntimeError("session_token_required")
+    payload = {
+        "idempotency_key": idempotency_key,
+        "mode": mode,
+        "interface_language": interface_language,
+        "session_token": session_token,
+    }
+    return _post_task_start(payload)
 
 
 def recover_generation_task() -> dict[str, Any]: