From 2a5032651754e7e442bdac6fdff1a048f69cbbc3 Mon Sep 17 00:00:00 2001
From: piexian <64474352+piexian@users.noreply.github.com>
Date: Wed, 4 Feb 2026 04:57:16 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E4=BF=AE=E5=A4=8D=20Token=20=E5=A4=B1?=
=?UTF-8?q?=E6=95=88=E6=A0=87=E8=AE=B0=E5=8F=8A=E6=97=A5=E5=BF=97=E7=9B=B8?=
=?UTF-8?q?=E5=85=B3=E9=97=AE=E9=A2=98?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
Token 失效标记修复:
- record_fail() 同时接受 401 和 403 状态码
- consume() 不再清零 fail_count,只有 record_success() 才清零
- _wrap_stream() 只在成功完成时调用 consume(),失败时不扣费
其他修复:
- 优化日志输出和错误处理
- 改进重试机制
- 完善配置默认值
---
app/api/v1/chat.py | 51 +++++-
app/api/v1/image.py | 47 ++---
app/services/grok/assets.py | 33 +++-
app/services/grok/chat.py | 31 ++--
app/services/grok/media.py | 30 ++--
app/services/grok/processor.py | 306 ++++++++++++++++++++++++++++++---
app/services/grok/retry.py | 200 ++++++++++++++++-----
app/services/token/manager.py | 6 +-
app/services/token/models.py | 12 +-
config.defaults.toml | 12 +-
10 files changed, 593 insertions(+), 135 deletions(-)
diff --git a/app/api/v1/chat.py b/app/api/v1/chat.py
index 781ea91..7069670 100644
--- a/app/api/v1/chat.py
+++ b/app/api/v1/chat.py
@@ -16,7 +16,9 @@ from app.core.exceptions import ValidationException
router = APIRouter(tags=["Chat"])
-VALID_ROLES = ["developer", "system", "user", "assistant"]
+VALID_ROLES = ["developer", "system", "user", "assistant", "tool", "function"]
+# 角色别名映射 (OpenAI 兼容)
+ROLE_ALIASES = {"function": "tool"}
USER_CONTENT_TYPES = ["text", "image_url", "input_audio", "file"]
@@ -24,13 +26,19 @@ class MessageItem(BaseModel):
"""消息项"""
role: str
content: Union[str, List[Dict[str, Any]]]
-
+ tool_call_id: Optional[str] = None # tool 角色需要的字段
+ name: Optional[str] = None # function 角色的函数名
+
@field_validator("role")
@classmethod
def validate_role(cls, v):
- if v not in VALID_ROLES:
+ # 大小写归一化
+ v_lower = v.lower() if isinstance(v, str) else v
+ # 别名映射
+ v_normalized = ROLE_ALIASES.get(v_lower, v_lower)
+ if v_normalized not in VALID_ROLES:
raise ValueError(f"role must be one of {VALID_ROLES}")
- return v
+ return v_normalized
class VideoConfig(BaseModel):
@@ -98,10 +106,26 @@ class ChatCompletionRequest(BaseModel):
messages: List[MessageItem] = Field(..., description="消息数组")
stream: Optional[bool] = Field(None, description="是否流式输出")
thinking: Optional[str] = Field(None, description="思考模式: enabled/disabled/None")
-
+
# 视频生成配置
video_config: Optional[VideoConfig] = Field(None, description="视频生成参数")
-
+
+ @field_validator("stream", mode="before")
+ @classmethod
+ def validate_stream(cls, v):
+ """确保 stream 参数被正确解析为布尔值"""
+ if v is None:
+ return None
+ if isinstance(v, bool):
+ return v
+ if isinstance(v, str):
+ if v.lower() in ("true", "1", "yes"):
+ return True
+ if v.lower() in ("false", "0", "no"):
+ return False
+ # 其他情况尝试转换
+ return bool(v)
+
model_config = {
"extra": "ignore"
}
@@ -174,6 +198,14 @@ def validate_request(request: ChatCompletionRequest):
param=f"messages.{idx}.content.{block_idx}.type",
code="invalid_type"
)
+ elif msg.role in ("tool", "function"):
+ # tool/function 角色只支持 text 类型,但内容可以是 JSON 字符串
+ if block_type != "text":
+ raise ValidationException(
+ message=f"The `{msg.role}` role only supports 'text' type, got '{block_type}'",
+ param=f"messages.{idx}.content.{block_idx}.type",
+ code="invalid_type"
+ )
elif block_type != "text":
raise ValidationException(
message=f"The `{msg.role}` role only supports 'text' type, got '{block_type}'",
@@ -203,10 +235,13 @@ def validate_request(request: ChatCompletionRequest):
@router.post("/chat/completions")
async def chat_completions(request: ChatCompletionRequest):
"""Chat Completions API - 兼容 OpenAI"""
-
+ from app.core.logger import logger
+
# 参数验证
validate_request(request)
-
+
+ logger.debug(f"Chat request: model={request.model}, stream={request.stream}")
+
# 检测视频模型
model_info = ModelService.get(request.model)
if model_info and model_info.is_video:
diff --git a/app/api/v1/image.py b/app/api/v1/image.py
index 15fb9e1..28f8ec9 100644
--- a/app/api/v1/image.py
+++ b/app/api/v1/image.py
@@ -73,6 +73,7 @@ async def call_grok(token_mgr, token: str, prompt: str, model_info) -> List[str]
调用 Grok 获取图片,返回 base64 列表
"""
chat_service = GrokChatService()
+ success = False
try:
response = await chat_service.chat(
@@ -87,22 +88,24 @@ async def call_grok(token_mgr, token: str, prompt: str, model_info) -> List[str]
# 收集图片
processor = ImageCollectProcessor(model_info.model_id, token)
images = await processor.process(response)
+ success = True
return images
except Exception as e:
logger.error(f"Grok image call failed: {e}")
return []
finally:
- # 记录使用
- try:
- effort = (
- EffortType.HIGH
- if (model_info and model_info.cost.value == "high")
- else EffortType.LOW
- )
- await token_mgr.consume(token, effort)
- except Exception as e:
- logger.warning(f"Failed to consume token: {e}")
+ # 只在成功时记录使用,失败时不扣费(避免清零 fail_count)
+ if success:
+ try:
+ effort = (
+ EffortType.HIGH
+ if (model_info and model_info.cost.value == "high")
+ else EffortType.LOW
+ )
+ await token_mgr.consume(token, effort)
+ except Exception as e:
+ logger.warning(f"Failed to consume token: {e}")
@router.post("/images/generations")
@@ -163,21 +166,25 @@ async def create_image(request: ImageGenerationRequest):
processor = ImageStreamProcessor(model_info.model_id, token, n=request.n)
- # 包装流式响应,在完成时记录使用
+ # 包装流式响应,在成功完成时记录使用
async def _wrap_stream(stream):
+ success = False
try:
async for chunk in stream:
yield chunk
+ success = True
finally:
- try:
- effort = (
- EffortType.HIGH
- if (model_info and model_info.cost.value == "high")
- else EffortType.LOW
- )
- await token_mgr.consume(token, effort)
- except Exception as e:
- logger.warning(f"Failed to consume token: {e}")
+ # 只在成功完成时扣费
+ if success:
+ try:
+ effort = (
+ EffortType.HIGH
+ if (model_info and model_info.cost.value == "high")
+ else EffortType.LOW
+ )
+ await token_mgr.consume(token, effort)
+ except Exception as e:
+ logger.warning(f"Failed to consume token: {e}")
return StreamingResponse(
_wrap_stream(processor.process(response)),
diff --git a/app/services/grok/assets.py b/app/services/grok/assets.py
index 30af070..718e9a5 100644
--- a/app/services/grok/assets.py
+++ b/app/services/grok/assets.py
@@ -24,11 +24,12 @@ from curl_cffi.requests import AsyncSession
from app.core.logger import logger
from app.core.config import get_config
from app.core.exceptions import (
- AppException,
- UpstreamException,
+ AppException,
+ UpstreamException,
ValidationException
)
from app.services.grok.statsig import StatsigService
+from app.services.token.service import TokenService
# ==================== 常量 ====================
@@ -335,9 +336,33 @@ class UploadService(BaseService):
file_uri = result.get("fileUri", "")
logger.info(f"Upload success: {filename} -> {file_id}", extra={"file_id": file_id})
return file_id, file_uri
-
+
+ # 401/403 表示 token 失效,记录失败并标记冷却
+ if response.status_code in (401, 403):
+ logger.warning(
+ f"Upload token invalid: {response.status_code}, recording failure",
+ extra={"response": response.text[:200]}
+ )
+ try:
+ await TokenService.record_fail(
+ token,
+ status_code=response.status_code,
+ reason=f"upload_auth_failed"
+ )
+ except Exception as mark_err:
+ logger.error(f"Failed to record token failure: {mark_err}")
+
+ raise UpstreamException(
+ message=f"Upload authentication failed: {response.status_code}",
+ details={
+ "status": response.status_code,
+ "response": response.text[:200],
+ "token_invalidated": True
+ }
+ )
+
logger.error(
- f"Upload failed: {filename} - {response.status_code}",
+ f"Upload failed: {filename} - {response.status_code}",
extra={"response": response.text[:200]}
)
raise UpstreamException(
diff --git a/app/services/grok/chat.py b/app/services/grok/chat.py
index cca5c8b..add405a 100644
--- a/app/services/grok/chat.py
+++ b/app/services/grok/chat.py
@@ -475,24 +475,27 @@ class ChatService:
token: Token 字符串
model: 模型名称
"""
+ success = False
try:
async for chunk in stream:
yield chunk
+ success = True
finally:
- # 无论成功完成还是客户端断开,都记录使用
- try:
- model_info = ModelService.get(model)
- effort = (
- EffortType.HIGH
- if (model_info and model_info.cost.value == "high")
- else EffortType.LOW
- )
- await token_mgr.consume(token, effort)
- logger.debug(
- f"Stream completed, recorded usage for token {token[:10]}... (effort={effort.value})"
- )
- except Exception as e:
- logger.warning(f"Failed to record stream usage: {e}")
+ # 只在成功完成时记录使用,失败/异常时不扣费
+ if success:
+ try:
+ model_info = ModelService.get(model)
+ effort = (
+ EffortType.HIGH
+ if (model_info and model_info.cost.value == "high")
+ else EffortType.LOW
+ )
+ await token_mgr.consume(token, effort)
+ logger.debug(
+ f"Stream completed, recorded usage for token {token[:10]}... (effort={effort.value})"
+ )
+ except Exception as e:
+ logger.warning(f"Failed to record stream usage: {e}")
@staticmethod
async def completions(
diff --git a/app/services/grok/media.py b/app/services/grok/media.py
index 004e033..079ca4f 100644
--- a/app/services/grok/media.py
+++ b/app/services/grok/media.py
@@ -388,23 +388,27 @@ class VideoService:
token: Token 字符串
model: 模型名称
"""
+ success = False
try:
async for chunk in stream:
yield chunk
+ success = True
finally:
- try:
- model_info = ModelService.get(model)
- effort = (
- EffortType.HIGH
- if (model_info and model_info.cost.value == "high")
- else EffortType.LOW
- )
- await token_mgr.consume(token, effort)
- logger.debug(
- f"Video stream completed, recorded usage for token {token[:10]}... (effort={effort.value})"
- )
- except Exception as e:
- logger.warning(f"Failed to record video stream usage: {e}")
+ # 只在成功完成时记录使用,失败/异常时不扣费
+ if success:
+ try:
+ model_info = ModelService.get(model)
+ effort = (
+ EffortType.HIGH
+ if (model_info and model_info.cost.value == "high")
+ else EffortType.LOW
+ )
+ await token_mgr.consume(token, effort)
+ logger.debug(
+ f"Video stream completed, recorded usage for token {token[:10]}... (effort={effort.value})"
+ )
+ except Exception as e:
+ logger.warning(f"Failed to record video stream usage: {e}")
@staticmethod
async def completions(
diff --git a/app/services/grok/processor.py b/app/services/grok/processor.py
index 166ac96..4b3d42f 100644
--- a/app/services/grok/processor.py
+++ b/app/services/grok/processor.py
@@ -1,17 +1,79 @@
"""
OpenAI 响应格式处理器
"""
+import asyncio
import time
import uuid
import random
import orjson
-from typing import Any, AsyncGenerator, Optional, AsyncIterable, List
+from typing import Any, AsyncGenerator, Optional, AsyncIterable, List, TypeVar
+
+from curl_cffi.requests.errors import RequestsError
from app.core.config import get_config
from app.core.logger import logger
+from app.core.exceptions import UpstreamException
from app.services.grok.assets import DownloadService
+def _is_http2_stream_error(e: Exception) -> bool:
+ """检查是否为 HTTP/2 流错误"""
+ err_str = str(e).lower()
+ return "http/2" in err_str or "curl: (92)" in err_str or "stream" in err_str
+
+
+T = TypeVar("T")
+
+
+class StreamIdleTimeoutError(Exception):
+ """流空闲超时错误"""
+ def __init__(self, idle_seconds: float):
+ self.idle_seconds = idle_seconds
+ super().__init__(f"Stream idle timeout after {idle_seconds}s")
+
+
+async def _with_idle_timeout(
+ iterable: AsyncIterable[T],
+ idle_timeout: float,
+ model: str = ""
+) -> AsyncGenerator[T, None]:
+ """
+ 包装异步迭代器,添加空闲超时检测
+
+ Args:
+ iterable: 原始异步迭代器
+ idle_timeout: 空闲超时时间(秒),0 表示禁用
+ model: 模型名称(用于日志)
+
+ Yields:
+ 原始迭代器的元素
+
+ Raises:
+ StreamIdleTimeoutError: 当空闲超时时
+ """
+ if idle_timeout <= 0:
+ async for item in iterable:
+ yield item
+ return
+
+ iterator = iterable.__aiter__()
+ while True:
+ try:
+ item = await asyncio.wait_for(
+ iterator.__anext__(),
+ timeout=idle_timeout
+ )
+ yield item
+ except asyncio.TimeoutError:
+ logger.warning(
+ f"Stream idle timeout after {idle_timeout}s",
+ extra={"model": model, "idle_timeout": idle_timeout}
+ )
+ raise StreamIdleTimeoutError(idle_timeout)
+ except StopAsyncIteration:
+ break
+
+
ASSET_URL = "https://assets.grok.com/"
@@ -81,7 +143,7 @@ class BaseProcessor:
class StreamProcessor(BaseProcessor):
"""流式响应处理器"""
-
+
def __init__(self, model: str, token: str = "", think: bool = None):
super().__init__(model, token)
self.response_id: Optional[str] = None
@@ -90,16 +152,78 @@ class StreamProcessor(BaseProcessor):
self.role_sent: bool = False
self.filter_tags = get_config("grok.filter_tags", [])
self.image_format = get_config("app.image_format", "url")
-
+ # 用于过滤跨 token 的标签
+ self._tag_buffer: str = ""
+ self._in_filter_tag: bool = False
+
if think is None:
self.show_think = get_config("grok.thinking", False)
else:
self.show_think = think
+
+ def _filter_token(self, token: str) -> str:
+ """
+ 过滤 token 中的特殊标签(如 ...)
+ 支持跨 token 的标签过滤
+ """
+ if not self.filter_tags:
+ return token
+
+ result = []
+ i = 0
+ while i < len(token):
+ char = token[i]
+
+ # 如果在过滤标签内
+ if self._in_filter_tag:
+ self._tag_buffer += char
+ # 检查是否到达结束标签
+ if char == '>':
+ # 检查是否是自闭合标签或结束标签
+ if '/>' in self._tag_buffer or '' in self._tag_buffer:
+ self._in_filter_tag = False
+ self._tag_buffer = ""
+ # 检查是否是开始标签结束(非自闭合)
+ elif self._tag_buffer.startswith('' in self._tag_buffer:
+ # 继续等待结束标签
+ pass
+ i += 1
+ continue
+
+ # 检查是否开始一个过滤标签
+ if char == '<':
+ # 查看后续字符
+ remaining = token[i:]
+ tag_started = False
+ for tag in self.filter_tags:
+ if remaining.startswith(f'<{tag}'):
+ tag_started = True
+ break
+ # 部分匹配(可能跨 token)
+ if len(remaining) < len(tag) + 1:
+ for j in range(1, len(remaining) + 1):
+ if f'<{tag}'.startswith(remaining[:j]):
+ tag_started = True
+ break
+
+ if tag_started:
+ self._in_filter_tag = True
+ self._tag_buffer = char
+ i += 1
+ continue
+
+ result.append(char)
+ i += 1
+
+ return ''.join(result)
async def process(self, response: AsyncIterable[bytes]) -> AsyncGenerator[str, None]:
"""处理流式响应"""
+ # 获取空闲超时配置
+ idle_timeout = get_config("grok.stream_idle_timeout", 45.0)
+
try:
- async for line in response:
+ async for line in _with_idle_timeout(response, idle_timeout, self.model):
if not line:
continue
try:
@@ -162,15 +286,42 @@ class StreamProcessor(BaseProcessor):
# 普通 token
if (token := resp.get("token")) is not None:
- if token and not (self.filter_tags and any(t in token for t in self.filter_tags)):
- yield self._sse(token)
+ if token:
+ filtered = self._filter_token(token)
+ if filtered:
+ yield self._sse(filtered)
if self.think_opened:
yield self._sse("\n")
yield self._sse(finish="stop")
yield "data: [DONE]\n\n"
+ except asyncio.CancelledError:
+ # 客户端断开连接,静默处理
+ logger.debug("Stream cancelled by client", extra={"model": self.model})
+ except StreamIdleTimeoutError as e:
+ # 流空闲超时
+ raise UpstreamException(
+ message=f"Stream idle timeout after {e.idle_seconds}s",
+ status_code=504,
+ details={"error": str(e), "type": "stream_idle_timeout", "idle_seconds": e.idle_seconds}
+ )
+ except RequestsError as e:
+ # HTTP/2 流错误转换为 UpstreamException
+ if _is_http2_stream_error(e):
+ logger.warning(f"HTTP/2 stream error: {e}", extra={"model": self.model})
+ raise UpstreamException(
+ message="Upstream connection closed unexpectedly",
+ status_code=502,
+ details={"error": str(e), "type": "http2_stream_error"}
+ )
+ logger.error(f"Stream request error: {e}", extra={"model": self.model})
+ raise UpstreamException(
+ message=f"Upstream request failed: {e}",
+ status_code=502,
+ details={"error": str(e)}
+ )
except Exception as e:
- logger.error(f"Stream processing error: {e}", extra={"model": self.model})
+ logger.error(f"Stream processing error: {e}", extra={"model": self.model, "error_type": type(e).__name__})
raise
finally:
await self.close()
@@ -178,19 +329,42 @@ class StreamProcessor(BaseProcessor):
class CollectProcessor(BaseProcessor):
"""非流式响应处理器"""
-
+
+ # 需要过滤的标签
+ FILTER_TAGS = ["grok:render", "xaiartifact", "xai:tool_usage_card"]
+
def __init__(self, model: str, token: str = ""):
super().__init__(model, token)
self.image_format = get_config("app.image_format", "url")
+ self.filter_tags = get_config("grok.filter_tags", self.FILTER_TAGS)
+
+ def _filter_content(self, content: str) -> str:
+ """过滤内容中的特殊标签"""
+ import re
+ if not content or not self.filter_tags:
+ return content
+
+ result = content
+ for tag in self.filter_tags:
+ # 匹配 ... 或
+ pattern = rf'<{re.escape(tag)}[^>]*>.*?{re.escape(tag)}>|<{re.escape(tag)}[^>]*/>'
+ result = re.sub(pattern, '', result, flags=re.DOTALL)
+ # 匹配 ... 跨行
+ pattern2 = rf'<{re.escape(tag)}[^>]*>[\s\S]*?{re.escape(tag)}>'
+ result = re.sub(pattern2, '', result, flags=re.DOTALL)
+
+ return result
async def process(self, response: AsyncIterable[bytes]) -> dict[str, Any]:
"""处理并收集完整响应"""
response_id = ""
fingerprint = ""
content = ""
-
+ # 获取空闲超时配置
+ idle_timeout = get_config("grok.stream_idle_timeout", 45.0)
+
try:
- async for line in response:
+ async for line in _with_idle_timeout(response, idle_timeout, self.model):
if not line:
continue
try:
@@ -227,12 +401,25 @@ class CollectProcessor(BaseProcessor):
if (meta := mr.get("metadata", {})).get("llm_info", {}).get("modelHash"):
fingerprint = meta["llm_info"]["modelHash"]
-
+
+ except asyncio.CancelledError:
+ logger.debug("Collect cancelled by client", extra={"model": self.model})
+ except StreamIdleTimeoutError as e:
+ logger.warning(f"Collect idle timeout: {e}", extra={"model": self.model})
+ # 非流式模式下,超时后返回已收集的内容
+ except RequestsError as e:
+ if _is_http2_stream_error(e):
+ logger.warning(f"HTTP/2 stream error in collect: {e}", extra={"model": self.model})
+ else:
+ logger.error(f"Collect request error: {e}", extra={"model": self.model})
except Exception as e:
- logger.error(f"Collect processing error: {e}", extra={"model": self.model})
+ logger.error(f"Collect processing error: {e}", extra={"model": self.model, "error_type": type(e).__name__})
finally:
await self.close()
-
+
+ # 过滤特殊标签
+ content = self._filter_content(content)
+
return {
"id": response_id,
"object": "chat.completion",
@@ -276,8 +463,11 @@ class VideoStreamProcessor(BaseProcessor):
async def process(self, response: AsyncIterable[bytes]) -> AsyncGenerator[str, None]:
"""处理视频流式响应"""
+ # 视频生成使用更长的空闲超时
+ idle_timeout = get_config("grok.video_idle_timeout", 90.0)
+
try:
- async for line in response:
+ async for line in _with_idle_timeout(response, idle_timeout, self.model):
if not line:
continue
try:
@@ -329,8 +519,30 @@ class VideoStreamProcessor(BaseProcessor):
yield self._sse("\n")
yield self._sse(finish="stop")
yield "data: [DONE]\n\n"
+ except asyncio.CancelledError:
+ logger.debug("Video stream cancelled by client", extra={"model": self.model})
+ except StreamIdleTimeoutError as e:
+ raise UpstreamException(
+ message=f"Video stream idle timeout after {e.idle_seconds}s",
+ status_code=504,
+ details={"error": str(e), "type": "stream_idle_timeout", "idle_seconds": e.idle_seconds}
+ )
+ except RequestsError as e:
+ if _is_http2_stream_error(e):
+ logger.warning(f"HTTP/2 stream error in video: {e}", extra={"model": self.model})
+ raise UpstreamException(
+ message="Upstream connection closed unexpectedly",
+ status_code=502,
+ details={"error": str(e), "type": "http2_stream_error"}
+ )
+ logger.error(f"Video stream request error: {e}", extra={"model": self.model})
+ raise UpstreamException(
+ message=f"Upstream request failed: {e}",
+ status_code=502,
+ details={"error": str(e)}
+ )
except Exception as e:
- logger.error(f"Video stream processing error: {e}", extra={"model": self.model})
+ logger.error(f"Video stream processing error: {e}", extra={"model": self.model, "error_type": type(e).__name__})
finally:
await self.close()
@@ -352,9 +564,11 @@ class VideoCollectProcessor(BaseProcessor):
"""处理并收集视频响应"""
response_id = ""
content = ""
-
+ # 视频生成使用更长的空闲超时
+ idle_timeout = get_config("grok.video_idle_timeout", 90.0)
+
try:
- async for line in response:
+ async for line in _with_idle_timeout(response, idle_timeout, self.model):
if not line:
continue
try:
@@ -379,8 +593,17 @@ class VideoCollectProcessor(BaseProcessor):
content = self._build_video_html(final_video_url, final_thumbnail_url)
logger.info(f"Video generated: {video_url}")
+ except asyncio.CancelledError:
+ logger.debug("Video collect cancelled by client", extra={"model": self.model})
+ except StreamIdleTimeoutError as e:
+ logger.warning(f"Video collect idle timeout: {e}", extra={"model": self.model})
+ except RequestsError as e:
+ if _is_http2_stream_error(e):
+ logger.warning(f"HTTP/2 stream error in video collect: {e}", extra={"model": self.model})
+ else:
+ logger.error(f"Video collect request error: {e}", extra={"model": self.model})
except Exception as e:
- logger.error(f"Video collect processing error: {e}", extra={"model": self.model})
+ logger.error(f"Video collect processing error: {e}", extra={"model": self.model, "error_type": type(e).__name__})
finally:
await self.close()
@@ -414,9 +637,11 @@ class ImageStreamProcessor(BaseProcessor):
async def process(self, response: AsyncIterable[bytes]) -> AsyncGenerator[str, None]:
"""处理流式响应"""
final_images = []
-
+ # 图片生成使用标准空闲超时
+ idle_timeout = get_config("grok.stream_idle_timeout", 45.0)
+
try:
- async for line in response:
+ async for line in _with_idle_timeout(response, idle_timeout, self.model):
if not line:
continue
try:
@@ -477,8 +702,30 @@ class ImageStreamProcessor(BaseProcessor):
"input_tokens_details": {"text_tokens": 5, "image_tokens": 20}
}
})
+ except asyncio.CancelledError:
+ logger.debug("Image stream cancelled by client")
+ except StreamIdleTimeoutError as e:
+ raise UpstreamException(
+ message=f"Image stream idle timeout after {e.idle_seconds}s",
+ status_code=504,
+ details={"error": str(e), "type": "stream_idle_timeout", "idle_seconds": e.idle_seconds}
+ )
+ except RequestsError as e:
+ if _is_http2_stream_error(e):
+ logger.warning(f"HTTP/2 stream error in image: {e}")
+ raise UpstreamException(
+ message="Upstream connection closed unexpectedly",
+ status_code=502,
+ details={"error": str(e), "type": "http2_stream_error"}
+ )
+ logger.error(f"Image stream request error: {e}")
+ raise UpstreamException(
+ message=f"Upstream request failed: {e}",
+ status_code=502,
+ details={"error": str(e)}
+ )
except Exception as e:
- logger.error(f"Image stream processing error: {e}")
+ logger.error(f"Image stream processing error: {e}", extra={"error_type": type(e).__name__})
raise
finally:
await self.close()
@@ -493,9 +740,11 @@ class ImageCollectProcessor(BaseProcessor):
async def process(self, response: AsyncIterable[bytes]) -> List[str]:
"""处理并收集图片"""
images = []
-
+ # 图片生成使用标准空闲超时
+ idle_timeout = get_config("grok.stream_idle_timeout", 45.0)
+
try:
- async for line in response:
+ async for line in _with_idle_timeout(response, idle_timeout, self.model):
if not line:
continue
try:
@@ -517,8 +766,17 @@ class ImageCollectProcessor(BaseProcessor):
b64 = base64_data
images.append(b64)
+ except asyncio.CancelledError:
+ logger.debug("Image collect cancelled by client")
+ except StreamIdleTimeoutError as e:
+ logger.warning(f"Image collect idle timeout: {e}")
+ except RequestsError as e:
+ if _is_http2_stream_error(e):
+ logger.warning(f"HTTP/2 stream error in image collect: {e}")
+ else:
+ logger.error(f"Image collect request error: {e}")
except Exception as e:
- logger.error(f"Image collect processing error: {e}")
+ logger.error(f"Image collect processing error: {e}", extra={"error_type": type(e).__name__})
finally:
await self.close()
diff --git a/app/services/grok/retry.py b/app/services/grok/retry.py
index cabfe05..c828e7c 100644
--- a/app/services/grok/retry.py
+++ b/app/services/grok/retry.py
@@ -2,12 +2,14 @@
Grok API 重试工具
提供可配置的重试机制,支持:
-- 可配置的重试次数
-- 可配置的重试状态码
-- 仅记录最后一次失败
+- 指数退避 + decorrelated jitter
+- Retry-After header 支持
+- 429 专用退避策略
+- 重试预算控制
"""
import asyncio
+import random
from typing import Callable, Any, Optional, List
from functools import wraps
@@ -18,139 +20,254 @@ from app.core.exceptions import UpstreamException
class RetryConfig:
"""重试配置"""
-
+
@staticmethod
def get_max_retry() -> int:
"""获取最大重试次数"""
- return get_config("grok.max_retry", 1)
-
+ return get_config("grok.max_retry", 3)
+
@staticmethod
def get_retry_codes() -> List[int]:
"""获取可重试的状态码"""
return get_config("grok.retry_status_codes", [401, 429, 403])
+ @staticmethod
+ def get_backoff_base() -> float:
+ """获取退避基础时间(秒)"""
+ return get_config("grok.retry_backoff_base", 0.5)
+
+ @staticmethod
+ def get_backoff_factor() -> float:
+ """获取退避因子"""
+ return get_config("grok.retry_backoff_factor", 2.0)
+
+ @staticmethod
+ def get_backoff_max() -> float:
+ """获取最大退避时间(秒)"""
+ return get_config("grok.retry_backoff_max", 30.0)
+
+ @staticmethod
+ def get_retry_budget() -> float:
+ """获取重试总预算时间(秒)"""
+ return get_config("grok.retry_budget", 90.0)
+
class RetryContext:
"""重试上下文"""
-
+
def __init__(self):
self.attempt = 0
self.max_retry = RetryConfig.get_max_retry()
self.retry_codes = RetryConfig.get_retry_codes()
self.last_error = None
self.last_status = None
-
+ self.total_delay = 0.0
+ self.retry_budget = RetryConfig.get_retry_budget()
+
+ # 退避参数
+ self.backoff_base = RetryConfig.get_backoff_base()
+ self.backoff_factor = RetryConfig.get_backoff_factor()
+ self.backoff_max = RetryConfig.get_backoff_max()
+
+ # decorrelated jitter 状态
+ self._last_delay = self.backoff_base
+
def should_retry(self, status_code: int) -> bool:
"""判断是否重试"""
- return (
- self.attempt < self.max_retry and
- status_code in self.retry_codes
- )
-
+ if self.attempt >= self.max_retry:
+ return False
+ if status_code not in self.retry_codes:
+ return False
+ if self.total_delay >= self.retry_budget:
+ return False
+ return True
+
def record_error(self, status_code: int, error: Exception):
"""记录错误信息"""
self.last_status = status_code
self.last_error = error
self.attempt += 1
+ def calculate_delay(self, status_code: int, retry_after: Optional[float] = None) -> float:
+ """
+ 计算退避延迟时间
+
+ Args:
+ status_code: HTTP 状态码
+ retry_after: Retry-After header 值(秒)
+
+ Returns:
+ 延迟时间(秒)
+ """
+ # 优先使用 Retry-After
+ if retry_after is not None and retry_after > 0:
+ delay = min(retry_after, self.backoff_max)
+ self._last_delay = delay
+ return delay
+
+ # 429 使用 decorrelated jitter
+ if status_code == 429:
+ # decorrelated jitter: delay = random(base, last_delay * 3)
+ delay = random.uniform(self.backoff_base, self._last_delay * 3)
+ delay = min(delay, self.backoff_max)
+ self._last_delay = delay
+ return delay
+
+ # 其他状态码使用指数退避 + full jitter
+ exp_delay = self.backoff_base * (self.backoff_factor ** self.attempt)
+ delay = random.uniform(0, min(exp_delay, self.backoff_max))
+ return delay
+
+ def record_delay(self, delay: float):
+ """记录延迟时间"""
+ self.total_delay += delay
+
+
+def extract_retry_after(error: Exception) -> Optional[float]:
+ """
+ 从异常中提取 Retry-After 值
+
+ Args:
+ error: 异常对象
+
+ Returns:
+ Retry-After 秒数,或 None
+ """
+ if not isinstance(error, UpstreamException):
+ return None
+
+ details = error.details or {}
+
+ # 尝试从 details 中获取
+ retry_after = details.get("retry_after")
+ if retry_after is not None:
+ try:
+ return float(retry_after)
+ except (ValueError, TypeError):
+ pass
+
+ # 尝试从 headers 中获取
+ headers = details.get("headers", {})
+ if isinstance(headers, dict):
+ retry_after = headers.get("Retry-After") or headers.get("retry-after")
+ if retry_after is not None:
+ try:
+ return float(retry_after)
+ except (ValueError, TypeError):
+ pass
+
+ return None
+
async def retry_on_status(
func: Callable,
*args,
extract_status: Callable[[Exception], Optional[int]] = None,
- on_retry: Callable[[int, int, Exception], None] = None,
+ on_retry: Callable[[int, int, Exception, float], None] = None,
**kwargs
) -> Any:
"""
通用重试函数
-
+
Args:
func: 重试的异步函数
*args: 函数参数
extract_status: 异常提取状态码的函数
- on_retry: 重试时的回调函数
+ on_retry: 重试时的回调函数 (attempt, status_code, error, delay)
**kwargs: 函数关键字参数
-
+
Returns:
函数执行结果
-
+
Raises:
最后一次失败的异常
"""
ctx = RetryContext()
-
+
# 状态码提取器
if extract_status is None:
def extract_status(e: Exception) -> Optional[int]:
if isinstance(e, UpstreamException):
return e.details.get("status") if e.details else None
return None
-
+
while ctx.attempt <= ctx.max_retry:
try:
result = await func(*args, **kwargs)
-
+
# 记录日志
if ctx.attempt > 0:
logger.info(
- f"Retry succeeded after {ctx.attempt} attempts"
+ f"Retry succeeded after {ctx.attempt} attempts, "
+ f"total delay: {ctx.total_delay:.2f}s"
)
-
+
return result
-
+
except Exception as e:
# 提取状态码
status_code = extract_status(e)
-
+
if status_code is None:
# 错误无法识别
logger.error(f"Non-retryable error: {e}")
raise
-
+
# 记录错误
ctx.record_error(status_code, e)
-
+
# 判断是否重试
if ctx.should_retry(status_code):
- delay = 0.5 * (ctx.attempt + 1) # 渐进延迟
+ # 提取 Retry-After
+ retry_after = extract_retry_after(e)
+
+ # 计算延迟
+ delay = ctx.calculate_delay(status_code, retry_after)
+
+ # 检查是否超出预算
+ if ctx.total_delay + delay > ctx.retry_budget:
+ logger.warning(
+ f"Retry budget exhausted: {ctx.total_delay:.2f}s + {delay:.2f}s > {ctx.retry_budget}s"
+ )
+ raise
+
+ ctx.record_delay(delay)
+
logger.warning(
f"Retry {ctx.attempt}/{ctx.max_retry} for status {status_code}, "
- f"waiting {delay}s"
+ f"waiting {delay:.2f}s (total: {ctx.total_delay:.2f}s)"
+ + (f", Retry-After: {retry_after}s" if retry_after else "")
)
-
+
# 回调
if on_retry:
- on_retry(ctx.attempt, status_code, e)
-
+ on_retry(ctx.attempt, status_code, e, delay)
+
await asyncio.sleep(delay)
continue
else:
# 不可重试或重试次数耗尽
if status_code in ctx.retry_codes:
- # 打印当前尝试次数(包括最后一次)
- logger.warning(
- f"Retry {ctx.attempt}/{ctx.max_retry} for status {status_code}, failed"
- )
logger.error(
- f"Retry exhausted after {ctx.max_retry} attempts, "
- f"last status: {status_code}"
+ f"Retry exhausted after {ctx.attempt} attempts, "
+ f"last status: {status_code}, total delay: {ctx.total_delay:.2f}s"
)
else:
logger.error(
f"Non-retryable status code: {status_code}"
)
-
+
# 抛出最后一次的错误
raise
def with_retry(
extract_status: Callable[[Exception], Optional[int]] = None,
- on_retry: Callable[[int, int, Exception], None] = None
+ on_retry: Callable[[int, int, Exception, float], None] = None
):
"""
重试装饰器
-
+
Usage:
@with_retry()
async def my_api_call():
@@ -175,4 +292,5 @@ __all__ = [
"RetryContext",
"retry_on_status",
"with_retry",
+ "extract_retry_after",
]
diff --git a/app/services/token/manager.py b/app/services/token/manager.py
index 4e8dd75..e266fe5 100644
--- a/app/services/token/manager.py
+++ b/app/services/token/manager.py
@@ -307,15 +307,15 @@ class TokenManager:
for pool in self.pools.values():
token = pool.get(raw_token)
if token:
- if status_code == 401:
+ if status_code in (401, 403):
token.record_fail(status_code, reason)
logger.warning(
- f"Token {raw_token[:10]}...: recorded 401 failure "
+ f"Token {raw_token[:10]}...: recorded {status_code} failure "
f"({token.fail_count}/{FAIL_THRESHOLD}) - {reason}"
)
else:
logger.info(
- f"Token {raw_token[:10]}...: non-401 error ({status_code}) - {reason} (not counted)"
+ f"Token {raw_token[:10]}...: non-auth error ({status_code}) - {reason} (not counted)"
)
self._schedule_save()
return True
diff --git a/app/services/token/models.py b/app/services/token/models.py
index 6c06299..137f509 100644
--- a/app/services/token/models.py
+++ b/app/services/token/models.py
@@ -90,13 +90,13 @@ class TokenInfo(BaseModel):
self.use_count += actual_cost # 使用 actual_cost 避免配额不足时过度计数
self.quota = max(0, self.quota - actual_cost)
- # 成功消耗后清空失败计数
- self.fail_count = 0
- self.last_fail_reason = None
+ # 注意:不在这里清零 fail_count,只有 record_success() 才清零
+ # 这样可以避免失败后调用 consume 导致失败计数被重置
if self.quota == 0:
self.status = TokenStatus.COOLING
- elif self.status in [TokenStatus.COOLING, TokenStatus.EXPIRED]:
+ elif self.status == TokenStatus.COOLING:
+ # 只从 COOLING 恢复,不从 EXPIRED 恢复
self.status = TokenStatus.ACTIVE
return actual_cost
@@ -127,8 +127,8 @@ class TokenInfo(BaseModel):
def record_fail(self, status_code: int = 401, reason: str = ""):
"""记录失败,达到阈值后自动标记为 expired"""
- # 仅 401 错误才计入失败
- if status_code != 401:
+ # 401/403 错误计入失败(都表示认证/授权失败)
+ if status_code not in (401, 403):
return
self.fail_count += 1
diff --git a/config.defaults.toml b/config.defaults.toml
index da8be8e..a6ddbea 100644
--- a/config.defaults.toml
+++ b/config.defaults.toml
@@ -1,7 +1,7 @@
[grok]
temporary = true
-stream = true
-thinking = true
+stream = false
+thinking = false
dynamic_statsig = true
filter_tags = ["xaiartifact","xai:tool_usage_card","grok:render"]
timeout = 120
@@ -10,6 +10,14 @@ asset_proxy_url = ""
cf_clearance = ""
max_retry = 3
retry_status_codes = [401,429,403]
+# 重试退避配置
+retry_backoff_base = 0.5
+retry_backoff_factor = 2.0
+retry_backoff_max = 30.0
+retry_budget = 90.0
+# 流空闲超时配置 (秒)
+stream_idle_timeout = 45.0
+video_idle_timeout = 90.0
[app]
app_url = "http://127.0.0.1:8000"