Ma/sora2api
1
0
Fork 0
mirror of https://github.com/TheSmallHanCat/sora2api.git synced 2026-02-04 10:14:41 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: 优化CF429错误处理及超时日志记录

Browse Source
This commit is contained in:
TheSmallHanCat
2026-01-11 13:04:19 +08:00
parent ab6aee0dc5
commit fb0569c298
1 changed files with 150 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

168
src/services/generation_handler.py
View File

@@ -533,7 +533,7 @@ class GenerationHandler:
await self.token_manager.record_usage(token_obj.id, is_video=is_video)
# Poll for results with timeout
async for chunk in self._poll_task_result(task_id, token_obj.token, is_video, stream, prompt, token_obj.id):
async for chunk in self._poll_task_result(task_id, token_obj.token, is_video, stream, prompt, token_obj.id, log_id, start_time):
yield chunk
# Record success
@@ -591,12 +591,6 @@ class GenerationHandler:
if is_video and token_obj and self.concurrency_manager:
await self.concurrency_manager.release_video(token_obj.id)
# Record error (check if it's an overload error)
if token_obj:
error_str = str(e).lower()
is_overload = "heavy_load" in error_str or "under heavy load" in error_str
await self.token_manager.record_error(token_obj.id, is_overload=is_overload)
# Parse error message to check if it's a structured error (JSON)
error_response = None
try:
@@ -604,15 +598,31 @@ class GenerationHandler:
except:
pass
# Check for CF shield/429 error
is_cf_or_429 = False
if error_response and isinstance(error_response, dict):
error_info = error_response.get("error", {})
if error_info.get("code") == "cf_shield_429":
is_cf_or_429 = True
# Record error (check if it's an overload error or CF/429 error)
if token_obj:
error_str = str(e).lower()
is_overload = "heavy_load" in error_str or "under heavy load" in error_str
# Don't record error for CF shield/429 (not token's fault)
if not is_cf_or_429:
await self.token_manager.record_error(token_obj.id, is_overload=is_overload)
# Update log entry with error data
duration = time.time() - start_time
if log_id:
if error_response:
# Structured error (e.g., unsupported_country_code)
# Structured error (e.g., unsupported_country_code, cf_shield_429)
status_code = 429 if is_cf_or_429 else 400
await self.db.update_request_log(
log_id,
response_body=json.dumps(error_response),
status_code=400,
status_code=status_code,
duration=duration
)
else:
@@ -626,7 +636,8 @@ class GenerationHandler:
raise e
async def _poll_task_result(self, task_id: str, token: str, is_video: bool,
stream: bool, prompt: str, token_id: int = None) -> AsyncGenerator[str, None]:
stream: bool, prompt: str, token_id: int = None,
log_id: int = None, start_time: float = None) -> AsyncGenerator[str, None]:
"""Poll for task result with timeout"""
# Get timeout from config
timeout = config.video_timeout if is_video else config.image_timeout
@@ -669,7 +680,19 @@ class GenerationHandler:
await self.concurrency_manager.release_video(token_id)
debug_logger.log_info(f"Released concurrency slot for token {token_id} due to timeout")
# Update task status to failed
await self.db.update_task(task_id, "failed", 0, error_message=f"Generation timeout after {elapsed_time:.1f} seconds")
# Update request log with timeout error
if log_id and start_time:
duration = time.time() - start_time
await self.db.update_request_log(
log_id,
response_body=json.dumps({"error": f"Generation timeout after {elapsed_time:.1f} seconds"}),
status_code=408,
duration=duration
)
raise Exception(f"Upstream API timeout: Generation exceeded {timeout} seconds limit")
@@ -1057,6 +1080,61 @@ class GenerationHandler:
)
except Exception as e:
# Check for CF shield/429 error - don't retry these
error_str = str(e)
is_cf_or_429 = False
try:
error_response = json.loads(error_str)
if isinstance(error_response, dict):
error_info = error_response.get("error", {})
if error_info.get("code") == "cf_shield_429":
is_cf_or_429 = True
except (json.JSONDecodeError, ValueError):
pass
# CF shield/429 detected - fail immediately
if is_cf_or_429:
debug_logger.log_error(
error_message="CF Shield/429 detected during polling, failing task immediately",
status_code=429,
response_text=error_str
)
# Update task status to failed
await self.db.update_task(task_id, "failed", 0, error_message="Cloudflare challenge or rate limit (429) triggered")
# Update request log with CF/429 error
if log_id and start_time:
duration = time.time() - start_time
await self.db.update_request_log(
log_id,
response_body=json.dumps({"error": "Cloudflare challenge or rate limit (429) triggered"}),
status_code=429,
duration=duration
)
# Release resources
if not is_video and token_id:
await self.load_balancer.token_lock.release_lock(token_id)
if self.concurrency_manager:
await self.concurrency_manager.release_image(token_id)
if is_video and token_id and self.concurrency_manager:
await self.concurrency_manager.release_video(token_id)
# Send error message to client if streaming
if stream:
yield self._format_stream_chunk(
reasoning_content="**CF Shield/429 Error**\\n\\nCloudflare challenge or rate limit (429) triggered\\n"
)
yield self._format_stream_chunk(
content="❌ Generation failed: Cloudflare challenge or rate limit (429) triggered. Please change proxy or reduce request frequency.",
finish_reason="STOP"
)
yield "data: [DONE]\\n\\n"
# Exit polling immediately
return
# For other errors, retry if not last attempt
if attempt >= max_attempts - 1:
raise e
continue
@@ -1315,6 +1393,20 @@ class GenerationHandler:
yield "data: [DONE]\n\n"
except Exception as e:
# Parse error to check for CF shield/429
error_response = None
try:
error_response = json.loads(str(e))
except:
pass
# Check for CF shield/429 error
is_cf_or_429 = False
if error_response and isinstance(error_response, dict):
error_info = error_response.get("error", {})
if error_info.get("code") == "cf_shield_429":
is_cf_or_429 = True
# Log failed character creation
duration = time.time() - start_time
await self._log_request(
@@ -1328,13 +1420,21 @@ class GenerationHandler:
"success": False,
"error": str(e)
},
status_code=500,
status_code=429 if is_cf_or_429 else 500,
duration=duration
)
# Record error (check if it's an overload error or CF/429 error)
if token_obj:
error_str = str(e).lower()
is_overload = "heavy_load" in error_str or "under heavy load" in error_str
# Don't record error for CF shield/429 (not token's fault)
if not is_cf_or_429:
await self.token_manager.record_error(token_obj.id, is_overload=is_overload)
debug_logger.log_error(
error_message=f"Character creation failed: {str(e)}",
status_code=500,
status_code=429 if is_cf_or_429 else 500,
response_text=str(e)
)
raise
@@ -1531,14 +1631,30 @@ class GenerationHandler:
duration=duration
)
# Record error (check if it's an overload error)
# Parse error to check for CF shield/429
error_response = None
try:
error_response = json.loads(str(e))
except:
pass
# Check for CF shield/429 error
is_cf_or_429 = False
if error_response and isinstance(error_response, dict):
error_info = error_response.get("error", {})
if error_info.get("code") == "cf_shield_429":
is_cf_or_429 = True
# Record error (check if it's an overload error or CF/429 error)
if token_obj:
error_str = str(e).lower()
is_overload = "heavy_load" in error_str or "under heavy load" in error_str
await self.token_manager.record_error(token_obj.id, is_overload=is_overload)
# Don't record error for CF shield/429 (not token's fault)
if not is_cf_or_429:
await self.token_manager.record_error(token_obj.id, is_overload=is_overload)
debug_logger.log_error(
error_message=f"Character and video generation failed: {str(e)}",
status_code=500,
status_code=429 if is_cf_or_429 else 500,
response_text=str(e)
)
raise
@@ -1624,14 +1740,30 @@ class GenerationHandler:
await self.token_manager.record_success(token_obj.id, is_video=True)
except Exception as e:
# Record error (check if it's an overload error)
# Parse error to check for CF shield/429
error_response = None
try:
error_response = json.loads(str(e))
except:
pass
# Check for CF shield/429 error
is_cf_or_429 = False
if error_response and isinstance(error_response, dict):
error_info = error_response.get("error", {})
if error_info.get("code") == "cf_shield_429":
is_cf_or_429 = True
# Record error (check if it's an overload error or CF/429 error)
if token_obj:
error_str = str(e).lower()
is_overload = "heavy_load" in error_str or "under heavy load" in error_str
await self.token_manager.record_error(token_obj.id, is_overload=is_overload)
# Don't record error for CF shield/429 (not token's fault)
if not is_cf_or_429:
await self.token_manager.record_error(token_obj.id, is_overload=is_overload)
debug_logger.log_error(
error_message=f"Remix generation failed: {str(e)}",
status_code=500,
status_code=429 if is_cf_or_429 else 500,
response_text=str(e)
)
raise