Update antigravity.py

Author: su-kaka

src/api/antigravity.py

@@ -56,7 +56,13 @@ def _extract_first_user_text(request_payload: Dict[str, Any]) -> str:
 
 
 def _generate_request_id() -> str:
-    return f"agent/{uuid.uuid4()}"
+    """生成完整格式的 requestId，对齐参考实现:
+    agent/{uuid}/{毫秒时间戳}/{trajectory_id}/{step}
+    """
+    trajectory_id = str(uuid.uuid4())
+    step = 1
+    ms = int(datetime.now(timezone.utc).timestamp() * 1000)
+    return f"agent/{uuid.uuid4()}/{ms}/{trajectory_id}/{step}"
 
 
 def _build_labels(model: str, trajectory_id: str, step: int) -> Dict[str, str]:
@@ -102,6 +108,7 @@ async def wrap_cli_request(
     gemini_request: Dict[str, Any],
     model: str,
     project_id: str,
+    enable_credit: bool = False,
 ) -> Tuple[Dict[str, Any], str]:
     """
     将 Gemini 格式请求包装成 Antigravity CLI 格式。
@@ -143,14 +150,19 @@ async def wrap_cli_request(
         "model": model,
         "userAgent": "antigravity",
         "requestType": "agent",
-        "enabledCreditTypes": ["GOOGLE_ONE_AI"],
     }
+    if enable_credit:
+        payload["enabledCreditTypes"] = ["GOOGLE_ONE_AI"]
     return payload, request_id
 
 
 # ==================== 辅助函数 ====================
 
-def build_antigravity_headers(access_token: str, extra_headers: Optional[Dict[str, str]] = None) -> Dict[str, str]:
+def build_antigravity_headers(
+    access_token: str,
+    extra_headers: Optional[Dict[str, str]] = None,
+    model_name: str = "",
+) -> Dict[str, str]:
     """构建 Antigravity CLI API 请求头。"""
     headers = {
         "User-Agent": ANTIGRAVITY_USER_AGENT,
@@ -159,11 +171,19 @@ def build_antigravity_headers(access_token: str, extra_headers: Optional[Dict[st
         "Accept": "*/*",
         "Accept-Encoding": "gzip",
         "Connection": "close",
+        "requestId": f"req-{uuid.uuid4()}",
     }
 
     for key, value in _sanitize_antigravity_headers(extra_headers).items():
         headers.setdefault(key, value)
 
+    # 根据模型名称判断 request_type
+    if model_name:
+        if "image" in model_name.lower():
+            headers["requestType"] = "image_gen"
+        else:
+            headers["requestType"] = "agent"
+
     return headers
 
 
@@ -237,6 +257,7 @@ async def stream_request(
     current_file, credential_data = cred_result
     access_token = credential_data.get("access_token") or credential_data.get("token")
     project_id = credential_data.get("project_id", "")
+    enable_credit = bool(credential_data.get("enable_credit", False))
 
     if not access_token:
         log.error(f"[ANTIGRAVITY STREAM] No access token in credential: {current_file}")
@@ -251,11 +272,11 @@ async def stream_request(
     antigravity_url = await get_antigravity_api_url()
     target_url = f"{antigravity_url}/v1internal:streamGenerateContent?alt=sse"
 
-    auth_headers = build_antigravity_headers(access_token, headers)
+    auth_headers = build_antigravity_headers(access_token, headers, model_name)
 
     # 构建 CLI 格式请求体
     inner_request = body.get("request", body)
-    final_payload, _ = await wrap_cli_request(inner_request, model_name, project_id)
+    final_payload, _ = await wrap_cli_request(inner_request, model_name, project_id, enable_credit)
 
     # 3. 调用stream_post_async进行请求
     retry_config = await get_retry_config()
@@ -513,6 +534,7 @@ async def non_stream_request(
     current_file, credential_data = cred_result
     access_token = credential_data.get("access_token") or credential_data.get("token")
     project_id = credential_data.get("project_id", "")
+    enable_credit = bool(credential_data.get("enable_credit", False))
 
     if not access_token:
         log.error(f"[ANTIGRAVITY] No access token in credential: {current_file}")
@@ -526,11 +548,11 @@ async def non_stream_request(
     antigravity_url = await get_antigravity_api_url()
     target_url = f"{antigravity_url}/v1internal:generateContent"
 
-    auth_headers = build_antigravity_headers(access_token, headers)
+    auth_headers = build_antigravity_headers(access_token, headers, model_name)
 
     # 构建 CLI 格式请求体
     inner_request = body.get("request", body)
-    final_payload, _ = await wrap_cli_request(inner_request, model_name, project_id)
+    final_payload, _ = await wrap_cli_request(inner_request, model_name, project_id, enable_credit)
 
     # 3. 调用post_async进行请求
     retry_config = await get_retry_config()
@@ -756,7 +778,7 @@ async def fetch_available_models() -> List[Dict[str, Any]]:
         return []
 
     # 构建请求头
-    headers = build_antigravity_headers(access_token)
+    headers = build_antigravity_headers(access_token, model_name="agent")
 
     try:
         # 使用 POST 请求获取模型列表
@@ -840,7 +862,7 @@ async def fetch_quota_info(access_token: str) -> Dict[str, Any]:
         }
     """
 
-    headers = build_antigravity_headers(access_token)
+    headers = build_antigravity_headers(access_token, model_name="agent")
 
     try:
         antigravity_url = await get_antigravity_api_url()