
    PL
jC                        d 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Zddlm	Z	 ddl
mZmZ ddlmZ ddlmZmZmZmZ ddlmZmZmZmZmZ  ej        e          Z e            dz  Zd	Zd
efdZ ded
dfdZ!de"d
e"fdZ#de"d
e$e"dz  e%f         fdZ&de"de"d
dfdZ'dde"de"dz  d
dfdZ(ddZ)ded
efdZ* G d d          Z+ G d de          Z,dS )zModal cloud execution environment using the native Modal SDK directly.

Uses ``Sandbox.create()`` + ``Sandbox.exec()`` instead of the older runtime
wrapper, while preserving Hermes' persistent snapshot behavior across sessions.
    N)Path)AnyOptional)get_hermes_home)BaseEnvironment_ThreadedProcessHandle_load_json_store_save_json_store)FileSyncManageriter_sync_filesquoted_mkdir_commandquoted_rm_commandunique_parent_dirszmodal_snapshots.jsondirectreturnc                  *    t          t                    S N)r	   _SNAPSHOT_STORE     </home/kuhnn/.hermes/hermes-agent/tools/environments/modal.py_load_snapshotsr   &   s    O,,,r   datac                 0    t          t          |            d S r   )r
   r   )r   s    r   _save_snapshotsr   *   s    _d+++++r   task_idc                     t            d|  S )N:)_DIRECT_SNAPSHOT_NAMESPACE)r   s    r   _direct_snapshot_keyr    .   s    (447444r   c                     t                      }t          |           }|                    |          }t          |t                    r|r|dfS |                    |           }t          |t                    r|r|dfS dS )NFT)NF)r   r    get
isinstancestr)r   	snapshotsnamespaced_keysnapshot_idlegacy_snapshot_ids        r   _get_snapshot_restore_candidater)   2   s    !!I)'22N--//K+s## " "E!!"w//$c** (/A (!4'';r   r'   c                     t                      }||t          |           <   |                    | d            t          |           d S r   )r   r    popr   )r   r'   r%   s      r   _store_direct_snapshotr,   >   sH    !!I/:I"7++,MM'4   Ir   c                     t                      }d}t          |           | fD ]:}|                    |          }||||k    r|                    |d            d};|rt	          |           d S d S )NFT)r   r    r"   r+   r   )r   r'   r%   updatedkeyvalues         r   _delete_direct_snapshotr1   E   s    !!IG$W--w7  c""=%;"6"6MM#t$$$G #	"""""# #r   c                      	 ddl m}   | dd           dS # t          $ r Y dS t          $ r!}t          t	          |                    d}~ww xY w)uG   Lazy-install modal on demand. Idempotent — fast no-op once installed.r   )ensurezterminal.modalF)promptN)tools.lazy_depsr3   ImportError	Exceptionr$   )_lazy_ensurees     r   _ensure_modal_sdkr:   S   s    "::::::%e444444    " " "#a&&!!!"s    
A	AA		A
image_specc                    t                       ddl}t          | t                    s| S |                     d          r|j                            |           S |                                 t          fddD                       }dg}|r|	                    dd           |j        
                    | |          S )	zConvert registry references or snapshot ids into Modal image objects.

    Includes add_python support for ubuntu/debian images (absorbed from PR 4511).
    r   Nzim-c              3       K   | ]}|v V  	d S r   r   ).0baselowers     r   	<genexpr>z'_resolve_modal_image.<locals>.<genexpr>n   s'      DDtTU]DDDDDDr   )ubuntudebianzRUN rm -rf /usr/local/lib/python*/site-packages/pip* 2>/dev/null; python -m ensurepip --upgrade --default-pip 2>/dev/null || truez^RUN apt-get update -qq && apt-get install -y -qq python3 python3-venv > /dev/null 2>&1 || true)setup_dockerfile_commands)r:   modalr#   r$   
startswithImagefrom_idr@   anyinsertfrom_registry)r;   _modal
add_pythonsetup_commandsr@   s       @r   _resolve_modal_imagerO   ^   s    
 j#&& U## 0|##J/// EDDDD/CDDDDDJ	JN  
al	
 	
 	
 <%%"0 &   r   c                   2    e Zd ZdZd Zd Zd Zd	dZd ZdS )
_AsyncWorkerzEBackground thread with its own event loop for async-safe Modal calls.c                 R    d | _         d | _        t          j                    | _        d S r   )_loop_thread	threadingEvent_startedselfs    r   __init__z_AsyncWorker.__init__   s#    :>
37!))r   c                     t          j        | j        d          | _        | j                                         | j                            d           d S )NT)targetdaemon   timeout)rU   Thread	_run_looprT   startrW   waitrX   s    r   rc   z_AsyncWorker.start   sN     't~dKKK2&&&&&r   c                     t          j                    | _        t          j        | j                   | j                                         | j                                         d S r   )asyncionew_event_looprS   set_event_looprW   setrun_foreverrX   s    r   rb   z_AsyncWorker._run_loop   sS    +--
tz***
     r   X  c                 ,   ddl m} | j        | j                                        r7t	          j        |          r|                                 t          d           ||| j                  }|t          d          |                    |          S )Nr   )safe_schedule_threadsafezAsyncWorker loop is not runningr_   )	agent.async_utilsrm   rS   	is_closedrf   iscoroutinecloseRuntimeErrorresult)rY   coror`   rm   futures        r   run_coroutinez_AsyncWorker.run_coroutine   s    >>>>>>:!5!5!7!7"4(( 

@AAA))$
;;>@AAA}}W}---r   c                     | j         r=| j                                         r$| j                             | j         j                   | j        r| j                            d           d S d S )N
   r_   )rS   
is_runningcall_soon_threadsafestoprT   joinrX   s    r   r{   z_AsyncWorker.stop   sn    : 	=$*//11 	=J++DJO<<<< 	*Lb)))))	* 	*r   N)rk   )	__name__
__module____qualname____doc__rZ   rc   rb   rv   r{   r   r   r   rQ   rQ      sj        OO* * *
' ' '
! ! !	. 	. 	. 	.* * * * *r   rQ   c                       e Zd ZdZdZdZ	 	 	 	 	 d#ded	ed
edee	ee
f                  dedef fdZdededdfdZdZdeeeef                  ddfdZdeddfdZdee         ddfdZd$dZdddddeded
ed edz  fd!Zd" Z xZS )%ModalEnvironmentzModal cloud execution via native Modal sandboxes.

    Spawn-per-call via _ThreadedProcessHandle wrapping async SDK calls.
    cancel_fn wired to sandbox.terminate for interrupt support.
    heredoc<   /rootNTdefaultimagecwdr`   modal_sandbox_kwargspersistent_filesystemr   c                 N   t                                          ||           || _        || _        d | _        d | _        t                      | _        d | _        t          |pi           d }d}| j        r<t          | j                  \  }}|r#t                              d|d d                    t                       dd lg 	 ddlm}	m}
m}  |	            D ]=}                    j                            |d         |d         	                     > |
            D ]=}                    j                            |d         |d         	                     > |            }|D ]=}                    j                            |d         |d         	                     >n2# t,          $ r%}t                              d
|           Y d }~nd }~ww xY w| j                                         dt2          ffd}	 |p|}	 t5          |          }| j                             ||          d          \  | _        | _        |r|rt9          | j        |           n# t,          $ r}|s t                              d|d d         |           t=          | j        |           t5          |          }| j                             ||          d          \  | _        | _        Y d }~nd }~ww xY wn(# t,          $ r | j                                          w xY wt                              d| j                   tA          d | j!        | j"        | j#        | j$                  | _        | j        %                    d           | &                                 d S )N)r   r`   Fz!Modal: restoring from snapshot %s   r   )get_credential_file_mountsiter_skills_filesiter_cache_files	host_pathcontainer_path)remote_pathz0Modal: could not load credential file mounts: %sr;   c                   K   j         j                            dd           d {V }t                    }r=t	          |                    dg                     }|                               ||d<    j        j        j        	 d
| |t          |                    dd                    d	| d {V }||fS )Nzhermes-agentT)create_if_missingmountssleepinfinityr`   i  )r   appr`   )r   r   )
Applookupaiodictlistr+   extendSandboxcreateint)r;   r   create_kwargsexisting_mountssandboxrL   cred_mountssandbox_kwargss        r   _create_sandboxz2ModalEnvironment.__init__.<locals>._create_sandbox   s     
)--nPT-UUUUUUUUC 00M :"&}'8'82'F'F"G"G&&{333*9h'5FN15# M--i>>??	 
         G <r   i,  r_   zBModal: failed to restore snapshot %s, retrying with base image: %sz Modal: sandbox created (task=%s)c                       t          d          S )Nz/root/.hermes)r   r   r   r   <lambda>z+ModalEnvironment.__init__.<locals>.<lambda>  s    !A!A r   )get_files_fn	upload_fn	delete_fnbulk_upload_fnbulk_download_fnT)force)'superrZ   _persistent_task_id_sandbox_apprQ   _worker_sync_managerr   r)   loggerinfor:   rE   tools.credential_filesr   r   r   appendMountfrom_local_filer7   debugrc   r   rO   rv   r,   warningr1   r{   r   _modal_upload_modal_delete_modal_bulk_upload_modal_bulk_downloadsyncinit_session)rY   r   r   r`   r   r   r   restored_snapshot_idrestored_from_legacy_keyr   r   r   mount_entryentrycache_filesr9   r   target_image_speceffective_imageexc
base_imagerL   r   r   	__class__s                        @@@r   rZ   zModalEnvironment.__init__   s    	S'2220	#~~5928b99##(  	\=\> >: ": $ \?AUVYWYVYAZ[[[	P           :9;;  ""L00#K0$/0@$A 1      +*,,  ""L00k*$)*:$; 1      +*,,K$  ""L00k*$)*:$; 1       	P 	P 	PLLKQOOOOOOOO	P 		 c 	  	  	  	  	  	  	  	  	 4 =P"67H"I"I+/<+E+E#OO44c ,F , ,(	4=  ( P,D P*4=:NOOO    + X("-s   (7KLLL1%88
+/<+E+E#OJ// ,F , ,(	4======  	 	 	L	 	6FFF,AA((2!6
 
 
 	d+++sQ   C$F0 0
G:GGL  AI+ L  +
K<5A=K72L  7K<<L   %L%r   r   r   c                     t          |                                          }t          j        |                              d          t          t          |          j                  }dt          j        |           dt          j        |            fd} j	        
                     |            d           dS )z4Upload a single file via base64 piped through stdin.asciiz	mkdir -p z && base64 -d > c                    K   j         j                            dd           d {V } d}j        }|t	                    k     ra| j                            |||z                       | j        j                                         d {V  ||z  }|t	                    k     a| j                                         | j        j                                         d {V  | j	                                         d {V  d S )Nbash-cr   )
r   execr   _STDIN_CHUNK_SIZElenstdinwritedrain	write_eofrd   )procoffset
chunk_sizeb64cmdrY   s      r   _writez.ModalEnvironment._modal_upload.<locals>._write1  s%     +//cBBBBBBBBDF/J3s88##
  VFZ,?%?!@AAAj&**,,,,,,,,,*$ 3s88## J  """*"&&((((((((()--//!!!!!!!!!r   r^   r_   N)r   
read_bytesbase64	b64encodedecoder$   parentshlexquoter   rv   )rY   r   r   contentcontainer_dirr   r   r   s   `     @@r   r   zModalEnvironment._modal_upload'  s    y//,,..w''..w77D--4556M22 6 6 ;{336 6 	

	" 
	" 
	" 
	" 
	" 
	" 
	" 	""6688R"88888r   i   filesc                 
   	
 |sdS t          j                    }t          j        |d          5 }|D ]/\  }}|                    ||                    d                     0	 ddd           n# 1 swxY w Y   t          j        |                                          	                    d          
t          |          }t          |          }| d		
 fd} j                             |            d	
           dS )a  Upload many files via tar archive piped through stdin.

        Builds a gzipped tar archive in memory and streams it into a
        ``base64 -d | tar xzf -`` pipeline via the process's stdin,
        avoiding the Modal SDK's 64 KB ``ARG_MAX_BYTES`` exec-arg limit.
        Nzw:gz)fileobjmode/)arcnamer   z && base64 -d | tar xzf - -C /c                    K   j         j                            dd           d {V } d}j        }|t	                    k     ra| j                            |||z                       | j        j                                         d {V  ||z  }|t	                    k     a| j                                         | j        j                                         d {V  | j	                                         d {V }|dk    r9| j
        j                                         d {V }t          d| d|           d S )Nr   r   r   zModal bulk upload failed (exit z): )r   r   r   r   r   r   r   r   r   rd   stderrreadrr   )r   r   r   	exit_codestderr_textr   payloadrY   s        r   _bulkz2ModalEnvironment._modal_bulk_upload.<locals>._bulkY  s}     +//cBBBBBBBBD F/J3w<<''
  0C)C!DEEEj&**,,,,,,,,,*$ 3w<<''
 J  """*"&&((((((((("immoo------IA~~$(K$4$8$8$:$:::::::"QiQQKQQ   ~r   x   r_   )ioBytesIOtarfileopenaddlstripr   r   getvaluer   r   r   r   rv   )rY   r   buftarr   r   parents
mkdir_partr   r   r   s   `        @@r   r   z#ModalEnvironment._modal_bulk_uploadE  su     	Fjll\#F333 	Ds*/ D D&	;	;+=+=c+B+BCCCCD	D 	D 	D 	D 	D 	D 	D 	D 	D 	D 	D 	D 	D 	D 	D "3<<>>2299'BB$U++)'22
;;;	 	 	 	 	 	 	, 	""5577C"88888s   3A11A58A5destc                       fd} j                              |            d          }t          |t                    r|                                }|                    |           dS )zDownload remote .hermes/ as a tar archive.

        Modal sandboxes always run as root, so /root/.hermes is hardcoded
        (consistent with iter_sync_files call on line 269).
        c                    K   j         j                            ddd           d {V } | j        j                                         d {V }| j                                         d {V }|dk    rt          d| d          |S )Nr   r   ztar cf - -C / root/.hermesr   z!Modal bulk download failed (exit ))r   r   r   stdoutr   rd   rr   )r   r   r   rY   s      r   	_downloadz8ModalEnvironment._modal_bulk_download.<locals>._downloadw  s      +//:       D )--////////D"immoo------IA~~"#Sy#S#S#STTTKr   r   r_   N)r   rv   r#   r$   encodewrite_bytes)rY   r  r  	tar_bytess   `   r   r   z%ModalEnvironment._modal_bulk_downloadq  s{    	 	 	 	 	 L..yy{{C.HH	i%% 	+!((**I#####r   remote_pathsc                 |     t          |           fd} j                             |            d           dS )z#Batch-delete remote files via exec.c                     K   j         j                            dd           d {V } | j                                         d {V  d S )Nr   r   )r   r   r   rd   )r   rm_cmdrY   s    r   _rmz+ModalEnvironment._modal_delete.<locals>._rm  s]      +//fEEEEEEEED)--//!!!!!!!!!r      r_   N)r   r   rv   )rY   r  r  r  s   `  @r   r   zModalEnvironment._modal_delete  sX    "<00	" 	" 	" 	" 	" 	" 	""3355""55555r   c                 8    | j                                          dS )zDSync files to sandbox via FileSyncManager (rate-limited internally).N)r   r   rX   s    r   _before_executez ModalEnvironment._before_execute  s    !!!!!r   Fr   )loginr`   
stdin_data
cmd_stringr  r  c                    | j         | j        fd}dt          t          t          f         ffd}t          ||          S )zEReturn a _ThreadedProcessHandle wrapping an async Modal sandbox exec.c                  d                          j                                        d           d S )Nr  r_   )rv   	terminater   )r   workers   r   cancelz*ModalEnvironment._run_bash.<locals>.cancel  s1      !2!6!6!8!8" EEEEEr   r   c                  X    fd}                       |             dz             S )Nc                  F  K   dg} r|                      ddg           n|                      dg            j        j        | d	i d {V }|j        j                                         d {V }|j        j                                         d {V }|j                                         d {V }t          |t                    r|	                    dd          }t          |t                    r|	                    dd          }|}|r|r| d| n|}||fS )	Nr   z-lr   r`   zutf-8replace)errors
)
r   r   r   r  r   r   rd   r#   bytesr   )
argsprocessr  r   r   outputr  r  r   r`   s
         r   _doz8ModalEnvironment._run_bash.<locals>.exec_fn.<locals>._do  so     x 4KKtZ 89999KKz 2333 0 0$ H H HHHHHHH&~26688888888&~26688888888"),"2"2"4"4444444	fe,, F#]]79]EEFfe,, F#]]79]EEF I6<H22&222&Fy((r   r^   r_   )rv   )r%  r  r  r   r`   r  s    r   exec_fnz+ModalEnvironment._run_bash.<locals>.exec_fn  sP    ) ) ) ) ) ) ) )& ''w|'DDDr   )	cancel_fn)r   r   tupler$   r   r   )	rY   r  r  r`   r  r  r&  r   r  s	    ```   @@r   	_run_bashzModalEnvironment._run_bash  s     -	F 	F 	F 	F 	F 	F	EsCx 	E 	E 	E 	E 	E 	E 	E 	E 	E 	E, &g@@@@r   c                 8     j         dS  j        r3t                              d            j                                          j        r	  fd}	  j                             |            d          }n# t          $ r d}Y nw xY w|r>t           j
        |           t                              d|dd          j
                   n2# t          $ r%}t                              d|           Y d}~nd}~ww xY w	  j                             j         j                                        d	           n# t          $ r Y nw xY w j                                         d _         d _        dS #  j                                         d _         d _        w xY w)
z>Snapshot the filesystem (if persistent) then stop the sandbox.Nz$Modal: syncing files from sandbox...c                  ^   K   j         j                                         d {V } | j        S r   )r   snapshot_filesystemr   	object_id)imgrY   s    r   	_snapshotz+ModalEnvironment.cleanup.<locals>._snapshot  s7       $ A E E G GGGGGGGC=(r   r   r_   z/Modal: saved filesystem snapshot %s for task %sr   z%Modal: filesystem snapshot failed: %sr  )r   r   r   r   	sync_backr   r   rv   r7   r,   r   r   r  r   r{   r   )rY   r/  r'   r9   s   `   r   cleanupzModalEnvironment.cleanup  s   = F 	+KK>???((*** 	KK) ) ) ) )'"&,"<"<YY[[RT"<"U"UKK  ' ' '"&KKK'  *4=+FFFKKI#CRC($-    K K KFJJJJJJJJK	L&&t}'>'B'B'D'Db&QQQQ 	 	 	D	 L DMDIII L DMDIsg   C $A8 7C 8BC BAC 
C:C55C:>8D7 6E0 7
EE0 EE0 0)F)r   r   NTr   r   N)r}   r~   r   r   _stdin_mode_snapshot_timeoutr$   r   r   r   r   boolrZ   r   r   r   r(  r   r   r   r   r  r)  r1  __classcell__)r   s   @r   r   r      s         K
 9=&* w ww w 	w
 'tCH~6w  $w w w w w w wr9s 9 9 9 9 9 98 (*9U38_(= *9$ *9 *9 *9 *9X$ $$ $ $ $ $*6$s) 6 6 6 6 6" " " " ;@!$+/ A  A  AC  A4  A A!Dj A  A  A  AD$ $ $ $ $ $ $r   r   r   r2  )-r   rf   r   r   loggingr   r   rU   pathlibr   typingr   r   hermes_constantsr   tools.environments.baser   r   r	   r
   tools.environments.file_syncr   r   r   r   r   	getLoggerr}   r   r   r   r   r   r   r$   r    r(  r5  r)   r,   r1   r:   rO   rQ   r   r   r   r   <module>r>     s      				                              , , , , , ,                         
	8	$	$!/##&<<% - - - - -,$ ,4 , , , ,5# 5# 5 5 5 5	S 	U3:t;K5L 	 	 	 	C c d    # #S #sTz #T # # # #" " " "S S    B"* "* "* "* "* "* "* "*Jz z z z z z z z z zr   