
    /j2                    P   S SK Jr  S SKrS SKrS SKJrJrJrJr  S SK	J
r
  S SKJrJr  S SKrSSKJrJrJr  SSKJr  S	rS
rSrSSS.r " S S\5      r " S S\5      r S   SS jjr SSSSSSSS.               SS jjjr S SSS.       S!S jjjr " S S5      rg)"    )annotationsN)AnyCallable	TypedDictcast)Path)LiteralNotRequired   )
OAuthErrorOpenAIErrorSubjectTokenProviderError)	to_threadz/urn:ietf:params:oauth:grant-type:token-exchangez#https://auth.openai.com/oauth/tokeni  z$urn:ietf:params:oauth:token-type:jwtz)urn:ietf:params:oauth:token-type:id_token)jwtidc                  *    \ rS rSr% S\S'   S\S'   Srg)SubjectTokenProvider   zLiteral['jwt', 'id']
token_typezCallable[[], str]	get_token N)__name__
__module____qualname____firstlineno____annotations____static_attributes__r       F/app/agent/.venv/lib/python3.13/site-packages/openai/auth/_workload.pyr   r      s    $$  r   r   c                  T    \ rS rSr% SrS\S'    S\S'    S\S'    S\S'    S	\S
'   Srg)WorkloadIdentity   z+A unique string that identifies the client.str	client_ididentity_provider_idservice_account_idr   providerzNotRequired[float]refresh_buffer_secondsr   N)r   r   r   r   __doc__r   r   r   r   r   r!   r!      s.    5N2GE""..r   r!   c                   ^  SU 4S jjnSUS.$ )a7  
Get a subject token provider for Kubernetes clusters with Workload Identity configured.

Cloud providers typically mount the subject token as a file in the container.

Args:
    token_file_path: path to the mounted service account token file. Defaults to `/var/run/secrets/kubernetes.io/serviceaccount/token`.
c                   >  [        TS5       n U R                  5       R                  5       nU(       d  [        ST S35      eUsS S S 5        $ ! , (       d  f       g = f! [         a  n[        ST SU 35      UeS nAff = f)NrzThe token file at z
 is empty.z!Failed to read the token file at z: )openreadstripr   	Exception)ftokenetoken_file_paths      r   r   5k8s_service_account_token_provider.<locals>.get_token;   s    	oos+q(36HHYYc4dee	 ,++
  	o+.OP_O``bcdbe,fgmnn	os4   A  6A	A  
AA  A   
B*A<<Br   r   r   returnr#   r   )r4   r   s   ` r   "k8s_service_account_token_providerr9   /   s    o  i88r   z
2018-02-01      $@)	object_idr$   
msi_res_idapi_versiontimeouthttp_clientc               6   ^ ^^^^^^ SUUUUUU U4S jjnSUS.$ )ab  
Get a subject token provider for Azure Managed Identities.

See: https://learn.microsoft.com/en-us/entra/identity/managed-identities-azure-resources/how-to-use-vm-token#get-a-token-using-http

Args:
    resource: the resource URI to request a token for. Defaults to `https://management.azure.com/` (Azure Resource Manager).
    object_id: the object ID of the managed identity to use, when multiple are assigned.
    client_id: the client ID of the managed identity to use, when multiple are assigned.
    msi_res_id: the ARM resource ID of the managed identity to use, when multiple are assigned.
    api_version: the Azure IMDS API version. Defaults to `2018-02-01`.
    timeout: the request timeout in seconds. Defaults to 10.0.
    http_client: optional httpx.Client instance to use for requests. If not provided, a new client will be created for each request.
c                   >  Sn TTS.nTb  TUS'   Tb  TUS'   T
b  T
US'   T	b  T	R                  XSS0TS9nO1[        R                  " 5        nUR                  XSS0TS9nS S S 5        WR                  (       a  [	        S	UR
                   3US
9eUR                  5       nUR                  S5      nU(       d
  [	        SUS
9e[        [        U5      $ ! , (       d  f       Nx= f! [         a  n[	        SU 35      UeS nAff = f)Nz5http://169.254.169.254/metadata/identity/oauth2/token)zapi-versionresourcer;   r$   r<   Metadatatrueparamsheadersr>   z4Failed to fetch Azure subject token from IMDS: HTTP responseaccess_tokenz3Azure IMDS response did not include an access_tokenz/Failed to fetch Azure subject token from IMDS: )
gethttpxClientis_errorr   status_codejsonr   r#   r0   )urlrF   rI   clientdatar2   r3   r=   r$   r?   r<   r;   rB   r>   s          r   r   8azure_managed_identity_token_provider.<locals>.get_tokena   s3   	jIC5@h%WF$&/{#$&/{#%'1|$&&??3
TZG[el?m\\^v%zz#zSYFZdkzlH $   /J8K_K_J`a%  ==?DHH^,E/IT\  U## $^  	j+.]^_]`,abhii	js1   AC% C"A1C% 
C"C% %
D/C>>Dr   r6   r7   r   )rB   r;   r$   r<   r=   r>   r?   r   s   ``````` r   %azure_managed_identity_token_providerrU   H   s    2j j@  i88r   )r>   r?   c               &   ^ ^^ SU UU4S jjnSUS.$ )a  
Get a subject token provider for GCP VM instances using the instance metadata server.

See: https://cloud.google.com/compute/docs/instances/verifying-instance-identity

Args:
    audience: the unique URI agreed upon by both the instance and the system verifying
        the instance's identity. Defaults to `https://api.openai.com/v1`.
    timeout: the request timeout in seconds. Defaults to 10.0.
    http_client: optional httpx.Client instance to use for requests. If not provided, a new client will be created for each request.
c                   >  Sn ST0nTb  TR                  XSS0TS9nO1[        R                  " 5        nUR                  XSS0TS9nS S S 5        WR                  (       a  [	        SUR
                   3US9eUR                  R                  5       nU(       d
  [	        SUS9eU$ ! , (       d  f       Nc= f! [         a  n[	        S	U 35      UeS nAff = f)
Nz]http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identityaudiencezMetadata-FlavorGooglerE   z=Failed to fetch GCP subject token from metadata server: HTTP rH   z+GCP metadata server returned an empty tokenz8Failed to fetch GCP subject token from metadata server: )	rK   rL   rM   rN   r   rO   textr/   r0   )	rQ   rF   rI   rR   r2   r3   rX   r?   r>   s	         r   r   (gcp_id_token_provider.<locals>.get_token   s    	sqC (+F&&??3HY[cGdnu?v\\^v%zz#GXZbFcmtzuH $   /ST\ThThSij%  MM'')E/0]hpqqL $^  	s+.fghfi,jkqrr	ss/   2B7 B&	AB7 &
B40B7 7
CCCr   r6   r7   r   )rX   r>   r?   r   s   ``` r   gcp_id_token_providerr\      s    $s s. Y77r   c                      \ rS rSr\S.   SS jjrSS jrSS jrSS jrSS jr	SS jr
SS	 jrSS
 jrSS jrSS jrSS jrSS jrSrg)WorkloadIdentityAuth   )token_exchange_urlc                   Xl         X l        S U l        S U l        S U l        SU l        [        R                  " 5       U l        [        R                  " U R                  5      U l
        g NF)workload_identityr`   _cached_token"_cached_token_expires_at_monotonic"_cached_token_refresh_at_monotonic_refreshing	threadingLock_lock	Condition
_condition)selfrc   r`   s      r   __init__WorkloadIdentityAuth.__init__   sT     "3"4)-@D/@D/!&^^%
#--djj9r   c                   U R                      U R                  (       aW  U R                  5       (       aB  U R                  R	                  5         U R                  (       a  U R                  5       (       a  MB  U R                  5       (       d8  U R                  5       (       d#  [        [        U R                  5      sS S S 5        $ U R                  (       a  U R                  (       a-  U R                  R	                  5         U R                  (       a  M-  U R                  nU R                  5       (       a  [        S5      e[        [        U5      sS S S 5        $ SU l        S S S 5         U R                  5         U R                      U R                  5       (       a  [        S5      e[        [        U R                  5      sS S S 5        U R                      SU l        U R                  R                  5         S S S 5        $ ! , (       d  f       N= f! , (       d  f       $ = f! , (       d  f       O= f U R                      SU l        U R                  R                  5         S S S 5        g ! , (       d  f       g = f! U R                      SU l        U R                  R                  5         S S S 5        f ! , (       d  f       f = f= f)Nz)Token is unusable after refresh completedTF)rj   rg   _token_unusablerl   wait_needs_refreshr   r#   rd   RuntimeError_perform_refresh
notify_all)rm   r2   s     r   r   WorkloadIdentityAuth.get_token   s   ZZ""t';';'='=$$& ""t';';'='= ''))$2E2E2G2GC!3!34 Z &&OO((* &&&**''))&'RSSC' Z  $D "		-!!#''))&'RSSC!3!34 
 #( **, 1 Z0  
 #( **, #( **, sz   A'G=6AG=AG=;G=G='I9 :H=	I9 "H=
H
H	
H-)I9 ="I((
I69K"J1(	K1
J?;Kc                H   #    [        U R                  5      I S h  vN $  N7fN)r   r   rm   s    r   get_token_async$WorkloadIdentityAuth.get_token_async   s     t~~....s   " "c                z    U R                      S U l        S U l        S U l        S S S 5        g ! , (       d  f       g = fry   )rj   rd   re   rf   rz   s    r   invalidate_token%WorkloadIdentityAuth.invalidate_token   s+    ZZ!%D6:D36:D3 ZZs   ,
:c                    U R                  5       n[        R                  " 5       nUS   nU R                     US   U l        X#-   U l        X R                  U5      -   U l        S S S 5        g ! , (       d  f       g = f)N
expires_inrJ   )_fetch_token_from_exchangetime	monotonicrj   rd   re   _refresh_delay_secondsrf   )rm   
token_datanowr   s       r   ru   %WorkloadIdentityAuth._perform_refresh   sf    446
nn-
ZZ!+N!;D696FD369<W<WXb<c6cD3 ZZs   ,A,,
A:c                   U R                  5       nU R                  S   S   n[        R                  U5      nUc3  [	        SU< SSR                  [        R                  5       5       35      e[        R                  " 5        nUR                  U R                  [        U R                  S   UUU R                  S   U R                  S   S	.S
S9nU R                  U5      sS S S 5        $ ! , (       d  f       g = f)Nr'   r   zUnsupported token type: z. Supported types: z, r$   r%   r&   )
grant_typer$   subject_tokensubject_token_typer%   r&   r:   )rP   r>   )_get_subject_tokenrc   SUBJECT_TOKEN_TYPESrK   r   joinkeysrL   rM   postr`   TOKEN_EXCHANGE_GRANT_TYPE_handle_token_response)rm   r   r   r   rR   rI   s         r   r   /WorkloadIdentityAuth._fetch_token_from_exchange   s    //1++J7E
044Z@%*:.8KDIIViVnVnVpLqKrs  \\^v{{''";!%!7!7!D%2*<,0,B,BCY,Z*.*@*@AU*V  # H ..x8 ^^s   AC++
C9c                    UR                   (       a  UR                  5       OS nUR                  S;   a	  [	        XS9eUR
                  (       a  Uc  [        S5      eUR                  S5      nUR                  S5      n[        U[        5      (       a  U(       d  [        S5      e[        U[        [        45      (       d  [        S5      eU[        U5      S.$ [        S	UR                   35      e! [         a    S n Nf = f)
N)i  i  i  )rI   bodyz4Token exchange succeeded but response body was emptyrJ   r   z<Token exchange response did not include a valid access_tokenz:Token exchange response did not include a valid expires_in)rJ   r   z"Token exchange failed with status )contentrP   
ValueErrorrO   r   
is_successr   rK   
isinstancer#   intfloat)rm   rI   r   rJ   r   s        r   r   +WorkloadIdentityAuth._handle_token_response
  s    	&.&6&68==?DD ?2h::|!"XYY88N3L,/JlC00!"`aaj3,77!"^__$0j@QRR01E1E0FG
 	
#  	D	s   #C2 2D Dc                \    U R                   S   nUS   " 5       nU(       d  [        S5      eU$ )Nr'   r   z>The workload identity provider returned an empty subject token)rc   r   )rm   r'   r   s      r   r   'WorkloadIdentityAuth._get_subject_token"  s4    ))*5 -/^__r   c                L    U R                   S L =(       d    U R                  5       $ ry   )rd   _token_expiredrz   s    r   rq   $WorkloadIdentityAuth._token_unusable)  s!    !!T)BT-@-@-BBr   c                b    U R                   c  g[        R                  " 5       U R                   :  $ )NT)re   r   r   rz   s    r   r   #WorkloadIdentityAuth._token_expired,  s)    22:~~4#J#JJJr   c                b    U R                   c  g[        R                  " 5       U R                   :  $ rb   )rf   r   r   rz   s    r   rs   #WorkloadIdentityAuth._needs_refresh1  s)    22:~~4#J#JJJr   c                z    U R                   R                  S[        5      n[        X!S-  5      n[	        X-
  S5      $ )Nr(   r   g        )rc   rK   DEFAULT_REFRESH_BUFFER_SECONDSminmax)rm   r   configured_buffereffective_buffers       r   r   +WorkloadIdentityAuth._refresh_delay_seconds6  s=     22667OQop0q.A:0#66r   )rd   re   rf   rl   rj   rg   r`   rc   N)rc   r!   r`   r#   r7   )r8   None)r8   dict[str, Any])rI   zhttpx.Responser8   r   )r8   bool)r   r   r8   r   )r   r   r   r   DEFAULT_TOKEN_EXCHANGE_URLrn   r   r{   r~   ru   r   r   r   rq   r   rs   r   r   r   r   r   r^   r^      s[    
 #=	: ,:  	: -:/;d92
0CK
K
7r   r^   )z3/var/run/secrets/kubernetes.io/serviceaccount/token)r4   z
str | Pathr8   r   )zhttps://management.azure.com/)rB   r#   r;   
str | Noner$   r   r<   r   r=   r#   r>   r   r?   httpx.Client | Noner8   r   )zhttps://api.openai.com/v1)rX   r#   r>   r   r?   r   r8   r   )
__future__r   r   rh   typingr   r   r   r   pathlibr   typing_extensionsr	   r
   rL   _exceptionsr   r   r   _utils._syncr   r   r   r   r   r   r!   r9   rU   r\   r^   r   r   r   <module>r      s7   "   1 1  2  L L $M B !%  2
5 !9 !
/y /& #X9994 499 ! !#'+9999 99 	99
 99 99 99 %99 99z 0)8 '+	)8)8 )8 %	)8
 )8XI7 I7r   