
    )jA                       d Z ddlmZ ddlZddlZddlZddlZddlZddl	Z	ddl
mZmZ ddlmZmZmZmZ  ej        e          ZdZddgZd	Zd
ZdZdZdZdZdZdZ G d de          Z  G d de           Z!ddd:dZ"d;dZ#eddd<d$Z$d=d&Z%e G d' d(                      Z&ddd)d>d+Z'd?d-Z(ddd)d@d/Z)e G d0 d1                      Z*ddd)dAd3Z+e G d4 d5                      Z,dddd6dBd9Z-dS )CuU  Google Code Assist API client — project discovery, onboarding, quota.

The Code Assist API powers Google's official gemini-cli. It sits at
``cloudcode-pa.googleapis.com`` and provides:

- Free tier access (generous daily quota) for personal Google accounts
- Paid tier access via GCP projects with billing / Workspace / Standard / Enterprise

This module handles the control-plane dance needed before inference:

1. ``load_code_assist()`` — probe the user's account to learn what tier they're on
   and whether a ``cloudaicompanionProject`` is already assigned.
2. ``onboard_user()`` — if the user hasn't been onboarded yet (new account, fresh
   free tier, etc.), call this with the chosen tier + project id. Supports LRO
   polling for slow provisioning.
3. ``retrieve_user_quota()`` — fetch the ``buckets[]`` array showing remaining
   quota per model, used by the ``/gquota`` slash command.

VPC-SC handling: enterprise accounts under a VPC Service Controls perimeter
will get ``SECURITY_POLICY_VIOLATED`` on ``load_code_assist``. We catch this
and force the account to ``standard-tier`` so the call chain still succeeds.

Derived from opencode-gemini-auth (MIT) and clawdbot/extensions/google. The
request/response shapes are specific to Google's internal Code Assist API,
documented nowhere public — we copy them from the reference implementations.
    )annotationsN)	dataclassfield)AnyDictListOptionalz#https://cloudcode-pa.googleapis.comz1https://daily-cloudcode-pa.sandbox.googleapis.comz4https://autopush-cloudcode-pa.sandbox.googleapis.comz	free-tierzlegacy-tierzstandard-tierz&google-api-nodejs-client/9.15.1 (gzip)zgl-node/24.0.0g      >@   g      @c                  2     e Zd ZdZddddddd fdZ xZS )CodeAssistErrora  Exception raised by the Code Assist (``cloudcode-pa``) integration.

    Carries HTTP status / response / retry-after metadata so the agent's
    ``error_classifier._extract_status_code`` and the main loop's Retry-After
    handling (which walks ``error.response.headers``) pick up the right
    signals.  Without these, 429s from the OAuth path look like opaque
    ``RuntimeError`` and skip the rate-limit path.
    code_assist_errorN)codestatus_coderesponseretry_afterdetailsmessagestrr   r   Optional[int]r   r   r   Optional[float]r   Optional[Dict[str, Any]]returnNonec                   t                                          |           || _        || _        || _        || _        |pi | _        d S )N)super__init__r   r   r   r   r   )selfr   r   r   r   r   r   	__class__s          ?/home/wildlama/.hermes/hermes-agent/agent/google_code_assist.pyr   zCodeAssistError.__init__N   sR     	!!!	
 '
 ! ' }"    )r   r   r   r   r   r   r   r   r   r   r   r   r   r   )__name__
__module____qualname____doc__r   __classcell__r   s   @r   r   r   D   sb          (%)'+,0% % % % % % % % % % % %r    r   c                  "     e Zd Zdd fdZ xZS )	ProjectIdRequiredError%GCP project id required for this tierr   r   r   r   c                N    t                                          |d           d S )Ncode_assist_project_id_requiredr   )r   r   )r   r   r   s     r   r   zProjectIdRequiredError.__init__o   s'    'HIIIIIr    )r)   )r   r   r   r   )r!   r"   r#   r   r%   r&   s   @r   r(   r(   n   sM        J J J J J J J J J J Jr    r(    user_agent_modelaccess_tokenr   r/   r   Dict[str, str]c          	         t           }|r| d| }ddd|  |t          t          t          j                              dS )Nz model/zapplication/jsonzBearer )zContent-TypeAcceptAuthorizationz
User-AgentzX-Goog-Api-Clientzx-activity-request-id)_GEMINI_CLI_USER_AGENT_X_GOOG_API_CLIENTr   uuiduuid4)r0   r/   uas      r   _build_headersr:   w   sZ    	B .--+--*$1<11/!$TZ\\!2!2  r    c                     ddddS )uL   Match Google's gemini-cli exactly — unrecognized metadata may be rejected.IDE_UNSPECIFIEDPLATFORM_UNSPECIFIEDGEMINI)ideTypeplatform
pluginType rB   r    r   _client_metadatarC      s     %*  r    )timeoutr/   urlbodyDict[str, Any]rD   floatc          	     X   t          j        |                              d          }t          j                            | |dt          ||                    }	 t          j                            ||          5 }|                                	                    dd          }|rt          j
        |          ni cd d d            S # 1 swxY w Y   d S # t          j        j        $ r}	d}
	 |	                                	                    dd          }
n# t          $ r Y nw xY wt          |
          rt          d	|
 d
          |	t          d|	j         d|
p|	j         d|	j                   |	d }	~	wt          j        j        $ r}	t          d|	 d          |	d }	~	ww xY w)Nzutf-8POSTr.   )datamethodheaders)rD   replace)errorsr-   zVPC-SC policy violation: code_assist_vpc_scr,   zCode Assist HTTP z: code_assist_http_zCode Assist request failed: code_assist_network_error)jsondumpsencodeurllibrequestRequestr:   urlopenreaddecodeloadserror	HTTPError	Exception_is_vpc_sc_violationr   r   reasonURLError)rE   rF   r0   rD   r/   rK   rW   r   rawexcdetails              r   
_post_jsonrf      sF    :d""7++Dn$$$v|>NOOO %  G^##GW#== 	2--//(((CCC&)14:c???r	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 <!   	XXZZ&&wy&AAFF 	 	 	D	  '' 	!4F44)    BBBF,@cjBB/SX//
 
 
 	 <    0300,
 
 
 	sn   !C <AC
=C 
CC CC F)+E8.)DE8
D%"E8$D%%AE88F)F$$F)boolc                   | sdS 	 t          j        |           }n # t           j        t          f$ r d| v cY S w xY wt	          |t
                    r|                    d          nd}t	          |t
                    sdS |                    d          pg }t	          |t                    r:|D ]7}t	          |t
                    r |                    d          pd}|dk    r dS 8t          |                    d	d                    }d|v S )
z=Detect a VPC Service Controls violation from a response body.FSECURITY_POLICY_VIOLATEDr]   Nr   ra   r-   Tr   )	rS   r\   JSONDecodeError
ValueError
isinstancedictgetlistr   )rF   parsedr]   r   itemra   msgs          r   r`   r`      s0    u2D!! *- 2 2 2)T11112 $.fd#;#;EFJJwEeT"" uii	""(bG'4     	  	 D$%%  (++1r77744
eii	2&&
'
'C%,,s    88c                  t    e Zd ZU dZdZded<   dZded<    ee          Z	ded<    ee
          Zd	ed
<   dS )CodeAssistProjectInfoz!Result from ``load_code_assist``.r-   r   current_tier_idcloudaicompanion_projectdefault_factoryz	List[str]allowed_tiersrG   rc   N)r!   r"   r#   r$   ru   __annotations__rv   r   ro   ry   rm   rc   rB   r    r   rt   rt      sw         ++O$&&&&&$uT:::M::::%555C555555r    rt   )
project_idr/   r{   c                  dd|it                      i}|r||d<   t          gt          z   }d}|D ]}| d}	 t          ||| |          }t	          |          c S # t
          $ rk}	|	j        dk    r8t                              d|           t          t          |	          cY d}	~	c S |	}t                              d
||	           Y d}	~	d}	~	ww xY w|r|t                      S )u   Call ``POST /v1internal:loadCodeAssist`` with prod → sandbox fallback.

    Returns whatever tier + project info Google reports. On VPC-SC violations,
    returns a synthetic ``standard-tier`` result so the chain can continue.
    metadataduetProjectcloudaicompanionProjectNz/v1internal:loadCodeAssistr.   rP   u6   VPC-SC violation on %s — defaulting to standard-tier)ru   rv   zloadCodeAssist failed on %s: %s)rC   CODE_ASSIST_ENDPOINTFALLBACK_ENDPOINTSrf   _parse_load_responser   r   loggerinfort   STANDARD_TIER_IDwarning)
r0   r{   r/   rF   	endpointslast_errendpointrE   resprd   s
             r   load_code_assistr      sb    	:
  
D  5*4&'%&);;I$(H  555	c4HXYYYD'----- 		 		 		x///TV^___,$4-7           HNN<hLLLHHHH		   """s#   !A
C&;C!C)CCr   c                   |                      d          pi }t          |t                    r$t          |                     d          pd          nd}t          |                      d          pd          }|                      d          pg }g }t          |t                    rU|D ]R}t          |t                    r;t          |                     d          pd          }|r|                    |           St          ||||           S )NcurrentTieridr-   r   allowedTiers)ru   rv   ry   rc   )rn   rl   rm   r   ro   appendrt   )r   current_tiertier_idprojectallowedallowed_idsttids           r   r   r     s   88M**0bL3=lD3Q3QYc,""4((.B///WYG$((455;<<Ghh~&&,"GK'4   , 	, 	,A!T"" ,!%%+++,, ,&&s+++ !(!	   r    r   c                  |t           k    r |t          k    r|st          d|d          |t                      d}|r||d<   t          }| d}t          ||| |          }|                    d          s|                    dd	          }|s|S t          t                    D ]}	t          j
        t                     | d
| }
	 t          |
i | |          }n6# t          $ r)}t                              d|	dz   |           Y d}~ed}~ww xY w|                    d          r|c S t                              dt                     |S )u  Call ``POST /v1internal:onboardUser`` to provision the user.

    For paid tiers, ``project_id`` is REQUIRED (raises ProjectIdRequiredError).
    For free tiers, ``project_id`` is optional — Google will assign one.

    Returns the final operation response. Polls ``/v1internal/<name>`` for up
    to ``_ONBOARDING_POLL_ATTEMPTS`` × ``_ONBOARDING_POLL_INTERVAL_SECONDS``
    (default: 12 × 5s = 1 min).
    zTier zQ requires a GCP project id. Set HERMES_GEMINI_PROJECT_ID or GOOGLE_CLOUD_PROJECT.)tierIdr}   r   z/v1internal:onboardUserr.   donenamer-   z/v1internal/z%Onboarding poll attempt %d failed: %s   Nz.Onboarding did not complete within %d attempts)FREE_TIER_IDLEGACY_TIER_IDr(   rC   r   rf   rn   range_ONBOARDING_POLL_ATTEMPTStimesleep!_ONBOARDING_POLL_INTERVAL_SECONDSr   r   r   )r0   r   r{   r/   rF   r   rE   r   op_nameattemptpoll_url	poll_resprd   s                r   onboard_userr     s     ,7n#<#<Z#<$DG D D D
 
 	
 $&& D  5*4&'#H
.
.
.Cc4@PQQQD 88F d((62&& 	K677 		! 		!GJ8999"9999H&x\Tdeee		"   FRSUXYYY }}V$$ !    !GIbcccKs   	C
D'DDc                  f    e Zd ZU ded<   dZded<   dZded<   dZded<    ee	          Z	d
ed<   dS )QuotaBucketr   model_idr-   
token_type        rH   remaining_fractionreset_time_isorw   rG   rc   N)
r!   r"   r#   rz   r   r   r   r   rm   rc   rB   r    r   r   r   T  sm         MMMJ #####N%555C555555r    r   List[QuotaBucket]c               F   i }|r||d<   t            d}t          ||| |          }|                    d          pg }g }t          |t                    s|S |D ]}t          |t
                    s|                    t          t          |                    d          pd          t          |                    d          pd          t          |                    d          pd	          t          |                    d
          pd          |                     |S )zDCall ``POST /v1internal:retrieveUserQuota`` and parse ``buckets[]``.r   z/v1internal:retrieveUserQuotar.   bucketsmodelIdr-   	tokenTyperemainingFractionr   	resetTime)r   r   r   r   rc   )
r   rf   rn   rl   ro   rm   r   r   r   rH   )	r0   r{   r/   rF   rE   r   raw_bucketsr   bs	            r   retrieve_user_quotar   ]  s<    D %$Y!
@
@
@Cc4@PQQQD((9%%+K!#Gk4((  	 	!T"" 	{y))/R00155--344$QUU+>%?%?%F3GGquu[117R88
 
 
 	 	 	 	 Nr    c                  L    e Zd ZU dZdZded<   dZded<   dZded<   dZded<   dS )	ProjectContextz)Resolved state for a given OAuth session.r-   r   r{   managed_project_idr   sourceN)	r!   r"   r#   r$   r{   rz   r   r   r   rB   r    r   r   r   ~  s]         33J     GFr    r   )configured_project_idenv_project_idr/   r   r   c                  |rt          |t          d          S |rt          |t          d          S t          | |          }|j        }|j        }|stt          | t          d|          }|                    d          pi }t          |t                    r&|p#t          |                    d          pd          }t          }d	}	nd
}	t          ||t          k    r|nd||	          S )aR  Figure out what project id + tier to use for requests.

    Priority:
      1. If configured_project_id or env_project_id is set, use that directly
         and short-circuit (no discovery needed).
      2. Otherwise call loadCodeAssist to see what Google says.
      3. If no tier assigned yet, onboard the user (free tier default).
    config)r{   r   r   envr.   r-   )r   r{   r/   r   r   	onboarded
discovered)r{   r   r   r   )r   r   r   rv   ru   r   r   rn   rl   rm   r   )
r0   r   r   r/   r   effective_projecttieronboard_respresponse_bodyr   s
             r   resolve_project_contextr     sC      
,$
 
 
 	

  
%$
 
 
 	
 L;KLLLD5D # -	
 
 
 %((44:mT** 	! K}(()BCCIrJJ  $040D0D,,"	   r    )r0   r   r/   r   r   r1   )r   r1   )rE   r   rF   rG   r0   r   rD   rH   r/   r   r   rG   )rF   r   r   rg   )r0   r   r{   r   r/   r   r   rt   )r   rG   r   rt   )
r0   r   r   r   r{   r   r/   r   r   rG   )r0   r   r{   r   r/   r   r   r   )
r0   r   r   r   r   r   r/   r   r   r   ).r$   
__future__r   rS   loggingr   urllib.errorrV   urllib.requestr7   dataclassesr   r   typingr   r   r   r	   	getLoggerr!   r   r   r   r   r   r   r5   r6   _DEFAULT_REQUEST_TIMEOUTr   r   RuntimeErrorr   r(   r:   rC   rf   r`   rt   r   r   r   r   r   r   r   rB   r    r   <module>r      s!   6 # " " " " "             ( ( ( ( ( ( ( ( , , , , , , , , , , , ,		8	$	$ =  8:  "  B %   $' !'% '% '% '% '%l '% '% '%TJ J J J J_ J J J BD          .% % % % % %P- - - -6 6 6 6 6 6 6 6 6 	'# '# '# '# '# '#T   8 1 1 1 1 1 1p 6 6 6 6 6 6 6 6 	     B         "$< < < < < < < <r    