
    i&                        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m	Z	 ddl
mZmZmZ ddlmZmZmZ ddlmZ  ej        e          ZddZ e	d           G d d                      Z G d de          ZdS )z1Managed Modal environment backed by tool-gateway.    )annotationsN)	dataclass)AnyDictOptional)BaseModalExecutionEnvironmentModalExecStartPreparedModalExec)resolve_managed_tool_gatewaynamestrdefaultfloatreturnc                    	 t          t          j        | t          |                              }|dk    r|n|S # t          t
          f$ r |cY S w xY w)Nr   )r   osgetenvr   	TypeError
ValueError)r   r   values      E/home/ubuntu/.hermes/hermes-agent/tools/environments/managed_modal.py_request_timeout_envr      sa    bic'll3344		uuw.z"   s   8; AAT)frozenc                      e Zd ZU ded<   dS )_ManagedModalExecHandler   exec_idN)__name__
__module____qualname____annotations__     r   r   r      s         LLLLLr"   r   c                      e Zd ZdZ edd          Z edd          Z edd          ZdZdZ	d	Z
	 	 	 	 	 d=d> fdZd?dZd@d"ZdAd$ZdBd&Zd' ZdCd(ZdDd)Zdd*dd+dEd3ZdFd5ZedGd9            ZedHd<            Z xZS )IManagedModalEnvironmentzCGateway-owned Modal sandbox with Hermes-compatible execute/cleanup..TERMINAL_MANAGED_MODAL_CONNECT_TIMEOUT_SECONDSg      ?0TERMINAL_MANAGED_MODAL_POLL_READ_TIMEOUT_SECONDSg      @2TERMINAL_MANAGED_MODAL_CANCEL_READ_TIMEOUT_SECONDSg      $@z4[Command interrupted - Modal sandbox exec cancelled]Managed Modal exec failed/root<   NTr   imager   cwdtimeoutintmodal_sandbox_kwargsOptional[Dict[str, Any]]persistent_filesystembooltask_idc                   t                                          ||           |                                  t          d          }|t	          d          |j                            d          | _        |j        | _	        || _
        || _        || _        t          |pi           | _        t          t!          j                              | _        |                                 | _        d S )N)r,   r-   modalzDManaged Modal requires a configured tool gateway and Nous user token/)super__init__)_guard_unsupported_credential_passthroughr   r   gateway_originrstrip_gateway_originnous_user_token_nous_user_token_task_id_persistent_imagedict_sandbox_kwargsr   uuiduuid4_create_idempotency_key_create_sandbox_sandbox_id)	selfr+   r,   r-   r/   r1   r3   gateway	__class__s	           r   r8   z ManagedModalEnvironment.__init__.   s     	S'22266888.w77?cddd&5<<SAA ' 70#$8$>B??'*4:<<'8'8$//11r"   preparedr
   r   r	   c           	        t          t          j                              }||j        |j        t          |j        dz            d}|j        
|j        |d<   	 |                     dd| j	         d|d          }n=# t          $ r0}t          |                     d	|           
          cY d }~S d }~ww xY w|j        dk    r7t          |                     |                     d|                    
          S |                                }|                    d          }|dv rLt          |                     |                    dd          |                    dd                    
          S |                    d          |k    r#t          |                     d          
          S t          t%          |                    S )N  )execIdcommandr,   	timeoutMs	stdinDataPOST/v1/sandboxes/z/execs
   jsonr-   zManaged Modal exec failed: )immediate_result  r(   status>   failedr-   	cancelled	completedoutput 
returncode   rO   z<Managed Modal exec start did not return the expected exec id)r   )handle)r   rD   rE   rP   r,   r.   r-   
stdin_data_requestrH   	Exceptionr	   _error_resultstatus_code_format_errorrW   get_resultr   )rI   rL   r   payloadresponseexcbodyrZ   s           r   _start_modal_execz)ManagedModalEnvironment._start_modal_execH   s(   djll##'<X-455	#
 #
 *#+#6GK 
	}}9!1999	 %  HH  	 	 	!!%!3!34WRU4W4W!X!X        	
 3&&!!%!3!3&&'BHMM" "    }}(##DDD!!%HHXr**HH\1--" "    88H((!!%!3!3R" "    %<W%M%M%MNNNNs   "A< <
B6%B1+B61B6rb   r   dict | Nonec                N   	 |                      dd| j         d|j         | j        | j        f          }n/# t
          $ r"}|                     d|           cY d }~S d }~ww xY w|j        dk    r|                     d          S |j        dk    r)|                     |                     d	|                    S |	                                }|
                    d
          }|dv r>|                     |
                    dd          |
                    dd                    S d S )NGETrT   /execs/r-   z Managed Modal exec poll failed: i  zManaged Modal exec not foundrY   zManaged Modal exec poll failedrZ   >   r[   r-   r\   r]   r^   r_   r`   ra   )rd   rH   r   _CONNECT_TIMEOUT_SECONDS_POLL_READ_TIMEOUT_SECONDSre   rf   rg   rh   rW   ri   rj   )rI   rb   status_responserm   status_bodyrZ   s         r   _poll_modal_execz(ManagedModalEnvironment._poll_modal_execy   sj   	P"mmJ!1JJ&.JJ68WX ,  OO
  	P 	P 	P%%&N&N&NOOOOOOOO	P &#--%%&DEEE&#--%%""#C_UU   &**,,**DDD<<"--a00   ts   47 
A#AA#A#Nonec                :    |                      |j                   d S N)_cancel_execr   )rI   rb   s     r   _cancel_modal_execz*ManagedModalEnvironment._cancel_modal_exec   s    &.)))))r"   rB   c                6    |                      d| dd          S )Nz#Managed Modal exec timed out after s|   )rj   )rI   r-   s     r   _timeout_result_for_modalz1ManagedModalEnvironment._timeout_result_for_modal   s"    ||L'LLLcRRRr"   c                   t          | dd           sd S 	 |                     dd| j         dd| j        id           n2# t          $ r%}t
                              d|           Y d }~nd }~ww xY wd | _        d S # d | _        w xY w)	NrH   rS   rT   z
/terminatesnapshotBeforeTerminater*   rV   z Managed Modal cleanup failed: %s)getattrrd   rH   r@   re   loggerwarning)rI   rm   s     r   cleanupzManagedModalEnvironment.cleanup   s    t]D11 	F	$MM=!1===-t/?        	D 	D 	DNN=sCCCCCCCC	D  $DtD####s,   )? A: 
A.	A)$A: )A..A: :	Bc                L   |                      | j                            d          d          }|                      | j                            d| j                            d                    d          }|                      | j                            d| j                            d                    d           }| j        | j        ||dt          d	t          | j        d
z                      | j        | j	        d}|||d<   | 
                    dd|dd| j        i          }|j        dk    r#t          |                     d|                    |                                }|                    d          }t!          |t"                    r|st          d          |S )Ncpura   	memoryMiBmemoryi   ephemeral_diskdiskMiBi6 i rN   )r+   r,   r   r   rQ   idleTimeoutMspersistentFilesystem
logicalKeyrS   z/v1/sandboxesr*   zx-idempotency-keyrW   r-   extra_headersrY   zManaged Modal create failedidz0Managed Modal create did not return a sandbox id)_coerce_numberrC   ri   rA   r,   maxr.   r-   r@   r?   rd   rF   rg   RuntimeErrorrh   rW   
isinstancer   )rI   r   r   diskcreate_payloadrl   rn   
sandbox_ids           r   rG   z'ManagedModalEnvironment._create_sandbox   s   !!$"6":":5"A"A1EE$$ $$[$2F2J2J82T2TUU
 
 "" $$%5t7K7O7OPY7Z7Z[[
 
 [8" #dlT.A*B*BCC$($4-	
 	
 (,N9%==#T%A ! 
 
 3&&t112OQYZZ[[[}}XXd^^
*c** 	S* 	SQRRRr"   c                n    	 ddl m} n# t          $ r Y dS w xY w |            }|rt          d          dS )z;Managed Modal does not sync or mount host credential files.r   )get_credential_file_mountsNzManaged Modal does not support host credential-file passthrough. Use TERMINAL_MODAL_MODE=direct when skills or config require credential files inside the sandbox.)tools.credential_filesr   re   r   )rI   r   mountss      r   r9   zAManagedModalEnvironment._guard_unsupported_credential_passthrough   sw    	IIIIIII 	 	 	FF	 ,+-- 	7  	 	s   	 
   r   methodpathrW   Dict[str, Any] | Noner   Dict[str, str] | Nonerequests.Responsec                   d| j          dd}|r|                    |           t          j        || j         | |||          S )NzBearer zapplication/json)AuthorizationzContent-Type)headersrW   r-   )r>   updaterequestsrequestr<   )rI   r   r   rW   r-   r   r   s          r   rd   z ManagedModalEnvironment._request   st    
 ?t'<>>.
 
  	*NN=)))#+T++
 
 
 	
r"   r   c                    	 |                      dd| j         d| d| j        | j        f           d S # t          $ r&}t
                              d|           Y d }~d S d }~ww xY w)NrS   rT   rs   z/cancelrt   z$Managed Modal exec cancel failed: %s)rd   rH   ru   _CANCEL_READ_TIMEOUT_SECONDSre   r   r   )rI   r   rm   s      r   r}   z$ManagedModalEnvironment._cancel_exec   s    	HMMJ!1JJ'JJJ68YZ      
  	H 	H 	HNNA3GGGGGGGGG	Hs   04 
A$AA$r   r   r   c                \    	 | |S t          |           S # t          t          f$ r |cY S w xY wr|   )r   r   r   )r   r   s     r   r   z&ManagedModalEnvironment._coerce_number  sF    	}<<:& 	 	 	NNN	s     ++prefixrl   c                   	 |                                 }t          |t                    rx|                    d          p)|                    d          p|                    d          }t          |t                    r	|r|  d| S |  dt          j        |d           S n# t          $ r Y nw xY w|j                                        }|r|  d| S |  d|j	         S )Nerrormessagecodez: F)ensure_asciiz: HTTP )
rW   r   rB   ri   r   dumpsre   textstriprg   )r   rl   rk   r   r   s        r   rh   z%ManagedModalEnvironment._format_error  s   	mmooG'4(( N!++g.._'++i2H2H_GKKX^L_L_gs++ 2 2$11111 MMDJwU$K$K$KMMM	N
  	 	 	D	 }""$$ 	'&&&&&77!5777s   BB$ B$ $
B10B1)r)   r*   NTr   )r+   r   r,   r   r-   r.   r/   r0   r1   r2   r3   r   )rL   r
   r   r	   )rb   r   r   rp   )rb   r   r   rz   )r-   r.   r   rB   )r   r   )r   rz   )r   r   r   r   rW   r   r-   r.   r   r   r   r   )r   r   r   rz   )r   r   r   r   r   r   )r   r   rl   r   r   r   )r   r   r   __doc__r   ru   rv   r   _client_timeout_grace_seconds_interrupt_output_unexpected_error_prefixr8   ro   ry   r~   r   r   rG   r9   rd   r}   staticmethodr   rh   __classcell__)rK   s   @r   r$   r$   $   s       MM334dfijj!5!56hjm!n!n#7#78lnq#r#r $(!N:
 9=&* 2 2 2 2 2 2 24/O /O /O /Ob   6* * * *S S S S$ $ $$( ( ( (T     04 "8<
 
 
 
 
 
&H H H H    \ 8 8 8 \8 8 8 8 8r"   r$   )r   r   r   r   r   r   )r   
__future__r   rW   loggingr   r   rD   dataclassesr   typingr   r   r   tools.environments.modal_utilsr   r	   r
   tools.managed_tool_gatewayr   	getLoggerr   r   r   r   r$   r!   r"   r   <module>r      sh   7 7 " " " " " "   				   ! ! ! ! ! ! & & & & & & & & & &         
 D C C C C C		8	$	$    $       v8 v8 v8 v8 v8; v8 v8 v8 v8 v8r"   