@@ -534,8 +534,9 @@ async def select_default_project(projects: List[Dict[str, Any]]) -> Optional[str
534534async def fetch_project_id_and_tier (
535535 access_token : str ,
536536 user_agent : str ,
537- api_base_url : str
538- ) -> Tuple [Optional [str ], Optional [str ]]:
537+ api_base_url : str ,
538+ include_credits : bool = False ,
539+ ) -> Tuple [Optional [str ], Optional [str ]] | Tuple [Optional [str ], Optional [str ], Optional [int ]]:
539540 """
540541 从 API 获取 project_id 和订阅等级
541542
@@ -545,8 +546,10 @@ async def fetch_project_id_and_tier(
545546 api_base_url: API base URL (e.g., antigravity or code assist endpoint)
546547
547548 Returns:
548- (project_id, subscription_tier) 元组
549+ 默认返回 (project_id, subscription_tier)
550+ 当 include_credits=True 时返回 (project_id, subscription_tier, credit_amount)
549551 subscription_tier 可能是 "FREE", "PRO", "ULTRA" 或 None
552+ credit_amount 为积分数量（整数）或 None
550553 """
551554 headers = {
552555 'User-Agent' : user_agent ,
@@ -572,18 +575,32 @@ def _map_raw_tier(raw_tier: Optional[str]) -> Optional[str]:
572575 return tier_mapping .get (raw_tier .lower (), "pro" )
573576
574577 subscription_tier = None
578+ credit_amount : Optional [int ] = None
575579
576580 # 步骤 1: 尝试 loadCodeAssist
577581 try :
578- project_id , raw_tier = await _try_load_code_assist (api_base_url , headers )
582+ project_id , raw_tier , raw_credit_amount = await _try_load_code_assist (api_base_url , headers )
579583 subscription_tier = _map_raw_tier (raw_tier )
580584
585+ if raw_credit_amount is not None :
586+ try :
587+ credit_amount = int (raw_credit_amount )
588+ log .info (
589+ f"[fetch_project_id_and_tier] Found credit_amount: { credit_amount }
590+ )
591+ except (TypeError , ValueError ):
592+ log .warning (
593+ f"[fetch_project_id_and_tier] Invalid credit_amount: { raw_credit_amount }
594+ )
595+
581596 if raw_tier :
582597 log .info (
583598 f"[fetch_project_id_and_tier] Raw tier '{ raw_tier } { subscription_tier }
584599 )
585600
586601 if project_id :
602+ if include_credits :
603+ return project_id , subscription_tier , credit_amount
587604 return project_id , subscription_tier
588605
589606 log .warning ("[fetch_project_id_and_tier] loadCodeAssist did not return project_id, falling back to onboardUser" )
@@ -596,28 +613,35 @@ def _map_raw_tier(raw_tier: Optional[str]) -> Optional[str]:
596613 try :
597614 project_id = await _try_onboard_user (api_base_url , headers )
598615 if project_id :
616+ if include_credits :
617+ return project_id , subscription_tier , credit_amount
599618 return project_id , subscription_tier
600619
601620 log .error ("[fetch_project_id_and_tier] Failed to get project_id from both loadCodeAssist and onboardUser" )
621+ if include_credits :
622+ return None , subscription_tier , credit_amount
602623 return None , subscription_tier
603624
604625 except Exception as e :
605626 log .error (f"[fetch_project_id_and_tier] onboardUser failed: { type (e ).__name__ } { e } )
606627 import traceback
607628 log .debug (f"[fetch_project_id_and_tier] Traceback: { traceback .format_exc ()} )
629+ if include_credits :
630+ return None , subscription_tier , credit_amount
608631 return None , subscription_tier
609632
610633
611634async def _try_load_code_assist (
612635 api_base_url : str ,
613636 headers : dict
614- ) -> Tuple [Optional [str ], Optional [str ]]:
637+ ) -> Tuple [Optional [str ], Optional [str ], Optional [ str ] ]:
615638 """
616639 尝试通过 loadCodeAssist 获取 project_id 和订阅等级
617640
618641 Returns:
619- (project_id, subscription_tier) 元组
642+ (project_id, subscription_tier, credit_amount ) 元组
620643 subscription_tier 可能是 "FREE", "PRO", "ULTRA" 或 None
644+ credit_amount 为字符串格式积分或 None
621645 """
622646 request_url = f"{ api_base_url .rstrip ('/' )}
623647 request_body = {
@@ -648,6 +672,7 @@ async def _try_load_code_assist(
648672 # 提取订阅等级 - 优先使用 paidTier（更准确反映实际权益）
649673 paid_tier = data .get ("paidTier" , {})
650674 current_tier = data .get ("currentTier" , {})
675+ available_credits = paid_tier .get ("availableCredits" , []) if isinstance (paid_tier , dict ) else []
651676
652677 # paidTier.id 优先，然后是 currentTier.id
653678 subscription_tier = None
@@ -658,6 +683,15 @@ async def _try_load_code_assist(
658683 subscription_tier = current_tier .get ("id" )
659684 log .info (f"[loadCodeAssist] Found currentTier: { subscription_tier } )
660685
686+ # 提取积分数量（如果返回了 availableCredits）
687+ credit_amount = None
688+ if isinstance (available_credits , list ) and available_credits :
689+ first_credit = available_credits [0 ]
690+ if isinstance (first_credit , dict ):
691+ credit_amount = first_credit .get ("creditAmount" )
692+ if credit_amount is not None :
693+ log .info (f"[loadCodeAssist] Found creditAmount: { credit_amount } )
694+
661695 # 检查是否有 currentTier（表示用户已激活）
662696 if current_tier :
663697 log .info ("[loadCodeAssist] User is already activated" )
@@ -666,13 +700,13 @@ async def _try_load_code_assist(
666700 project_id = data .get ("cloudaicompanionProject" )
667701 if project_id :
668702 log .info (f"[loadCodeAssist] Successfully fetched project_id: { project_id } { subscription_tier } )
669- return project_id , subscription_tier
703+ return project_id , subscription_tier , credit_amount
670704
671705 log .warning ("[loadCodeAssist] No project_id in response" )
672- return None , subscription_tier
706+ return None , subscription_tier , credit_amount
673707 else :
674708 log .info ("[loadCodeAssist] User not activated yet (no currentTier)" )
675- return None , None
709+ return None , None , credit_amount
676710 else :
677711 log .warning (f"[loadCodeAssist] Failed: HTTP { response .status_code } )
678712 log .warning (f"[loadCodeAssist] Response body: { response .text [:500 ]} )
0 commit comments