@@ -31,7 +31,9 @@ use codex_protocol::ThreadId;
3131use codex_protocol:: error:: CodexErr ;
3232use codex_protocol:: error:: CodexErrorDetails ;
3333use codex_protocol:: error:: Result as CodexResult ;
34+ use codex_protocol:: protocol:: InterAgentCommunication ;
3435use codex_protocol:: protocol:: MultiAgentVersion ;
36+ use codex_protocol:: protocol:: Op ;
3537use codex_protocol:: protocol:: SessionSource ;
3638use codex_protocol:: protocol:: SubAgentSource ;
3739use codex_protocol:: protocol:: TokenUsage ;
@@ -44,6 +46,8 @@ use core_test_support::test_codex::test_codex;
4446use core_test_support:: wait_for_event;
4547use futures:: future:: BoxFuture ;
4648use pretty_assertions:: assert_eq;
49+ use serde_json:: json;
50+ use test_case:: test_case;
4751use wiremock:: MockServer ;
4852
4953struct TestAgentControl {
@@ -214,16 +218,21 @@ impl AgentControl for TestAgentControl {
214218
215219async fn test_with_host_control (
216220 server : & MockServer ,
221+ extra_features : & [ Feature ] ,
217222) -> anyhow:: Result < ( TestCodex , Arc < TestAgentControl > ) > {
218223 let thread_id = ThreadId :: new ( ) ;
219224 let controller = Arc :: new ( TestAgentControl :: new ( thread_id) ) ;
225+ let extra_features = extra_features. to_vec ( ) ;
220226 let test = test_codex ( )
221- . with_config ( |config| {
227+ . with_config ( move |config| {
222228 for feature in [
223229 Feature :: Collab ,
224230 Feature :: MultiAgentV2 ,
225231 Feature :: AgentMessageBoard ,
226- ] {
232+ ]
233+ . into_iter ( )
234+ . chain ( extra_features. iter ( ) . copied ( ) )
235+ {
227236 config
228237 . features
229238 . enable ( feature)
@@ -251,7 +260,7 @@ async fn test_with_host_control(
251260async fn host_threads_preserve_lineage_settings_and_resume_routing ( ) -> anyhow:: Result < ( ) > {
252261 skip_if_no_network ! ( Ok ( ( ) ) ) ;
253262 let server = responses:: start_mock_server ( ) . await ;
254- let ( test, controller) = test_with_host_control ( & server) . await ?;
263+ let ( test, controller) = test_with_host_control ( & server, & [ ] ) . await ?;
255264 let root_id = test. session_configured . thread_id ;
256265 let mut child_config = test. config . clone ( ) ;
257266 child_config. service_tier = None ;
@@ -490,7 +499,7 @@ async fn host_factory_follows_thread_lifecycle() -> anyhow::Result<()> {
490499async fn collaboration_tools_dispatch_to_the_host_controller ( ) -> anyhow:: Result < ( ) > {
491500 skip_if_no_network ! ( Ok ( ( ) ) ) ;
492501 let server = responses:: start_mock_server ( ) . await ;
493- let ( test, _controller) = test_with_host_control ( & server) . await ?;
502+ let ( test, _controller) = test_with_host_control ( & server, & [ ] ) . await ?;
494503 let mut events = vec ! [ responses:: ev_response_created( "tools" ) ] ;
495504 for ( call, name, arguments) in [
496505 (
@@ -541,7 +550,7 @@ use codex_protocol::protocol::InternalSessionSource;
541550#[ tokio:: test( flavor = "multi_thread" , worker_threads = 2 ) ]
542551async fn host_mailbox_notifies_waiting_agent_and_delivers_once ( ) -> anyhow:: Result < ( ) > {
543552 let server = responses:: start_mock_server ( ) . await ;
544- let ( test, controller) = test_with_host_control ( & server) . await ?;
553+ let ( test, controller) = test_with_host_control ( & server, & [ ] ) . await ?;
545554 responses:: mount_sse_once (
546555 & server,
547556 responses:: sse ( vec ! [
@@ -609,3 +618,80 @@ async fn host_mailbox_notifies_waiting_agent_and_delivers_once() -> anyhow::Resu
609618 test. codex . shutdown_and_wait ( ) . await ?;
610619 Ok ( ( ) )
611620}
621+
622+ /// Both delivery paths wake an opted-in completed root exactly once.
623+ #[ test_case( false ; "direct_op" ) ]
624+ #[ test_case( true ; "host_mailbox" ) ]
625+ #[ tokio:: test( flavor = "multi_thread" , worker_threads = 2 ) ]
626+ async fn agent_mail_wakes_completed_root_when_enabled (
627+ via_host_mailbox : bool ,
628+ ) -> anyhow:: Result < ( ) > {
629+ let server = responses:: start_mock_server ( ) . await ;
630+ let ( test, controller) =
631+ test_with_host_control ( & server, & [ Feature :: WakeOnAgentMessages ] ) . await ?;
632+ let mut final_answer = responses:: ev_assistant_message ( "first-answer" , "initial work finished" ) ;
633+ final_answer[ "item" ] [ "phase" ] = json ! ( "final_answer" ) ;
634+ let response = responses:: mount_sse_sequence (
635+ & server,
636+ vec ! [
637+ responses:: sse( vec![
638+ responses:: ev_response_created( "initial" ) ,
639+ final_answer,
640+ responses:: ev_completed( "initial" ) ,
641+ ] ) ,
642+ responses:: sse( vec![
643+ responses:: ev_response_created( "agent-mail" ) ,
644+ responses:: ev_assistant_message( "mail-answer" , "received the agent update" ) ,
645+ responses:: ev_completed( "agent-mail" ) ,
646+ ] ) ,
647+ ] ,
648+ )
649+ . await ;
650+ test. submit_text_turn ( "Finish the initial work." ) . await ?;
651+ assert_eq ! ( response. requests( ) . len( ) , 1 ) ;
652+
653+ let message = InterAgentCommunication :: new (
654+ AgentPath :: root ( ) . join ( "worker" ) . expect ( "valid worker path" ) ,
655+ AgentPath :: root ( ) ,
656+ Vec :: new ( ) ,
657+ "Late agent update" . into ( ) ,
658+ /*trigger_turn*/ false ,
659+ ) ;
660+ if via_host_mailbox {
661+ controller. mail . lock ( ) . expect ( "mail lock" ) . push ( message) ;
662+ controller. mailbox_state . send_replace ( true ) ;
663+ } else {
664+ test. codex
665+ . submit ( Op :: InterAgentCommunication {
666+ communication : message,
667+ start_options : Default :: default ( ) ,
668+ } )
669+ . await ?;
670+ }
671+ wait_for_event ( & test. codex , |event| {
672+ matches ! ( event, EventMsg :: TurnComplete ( _) )
673+ } )
674+ . await ;
675+ test. codex . shutdown_and_wait ( ) . await ?;
676+
677+ let requests = response. requests ( ) ;
678+ assert_eq ! ( requests. len( ) , 2 ) ;
679+ assert_eq ! (
680+ responses:: received_responses_requests( & server) . await . len( ) ,
681+ 2
682+ ) ;
683+ assert_ne ! (
684+ requests[ 0 ] . body_json( ) [ "client_metadata" ] [ "turn_id" ] ,
685+ requests[ 1 ] . body_json( ) [ "client_metadata" ] [ "turn_id" ]
686+ ) ;
687+ let messages = requests[ 1 ] . inputs_of_type ( "agent_message" ) ;
688+ assert_eq ! ( messages. len( ) , 1 ) ;
689+ assert_eq ! ( messages[ 0 ] [ "author" ] , "/root/worker" ) ;
690+ assert_eq ! ( messages[ 0 ] [ "recipient" ] , "/root" ) ;
691+ assert_eq ! (
692+ messages[ 0 ] [ "content" ] ,
693+ json!( [ { "type" : "input_text" , "text" : "Late agent update" } ] )
694+ ) ;
695+ assert ! ( controller. mail. lock( ) . expect( "mail lock" ) . is_empty( ) ) ;
696+ Ok ( ( ) )
697+ }
0 commit comments